From 03fc89da00ef6f5dda778646e3c98eceb59078ce Mon Sep 17 00:00:00 2001
From: CaIon <i@caion.me>
Date: Fri, 15 Aug 2025 12:50:27 +0800
Subject: [PATCH] refactor: add email masking function and enhance RelayInfo
 logging

This commit introduces a new function, MaskEmail, to mask user email addresses in logs, preventing PII leakage. Additionally, the RelayInfo logging has been updated to utilize this new masking function, ensuring sensitive information is properly handled. The channel test logic has also been improved to dynamically determine the relay format based on the request path.
---
 common/str.go                | 18 ++++++++++++++++++
 controller/channel-test.go   |  8 +++++++-
 relay/common/relay_info.go   |  2 +-
 service/pre_consume_quota.go |  5 +++--
 4 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/common/str.go b/common/str.go
index 7d4cdaf0..a769b8e4 100644
--- a/common/str.go
+++ b/common/str.go
@@ -99,6 +99,24 @@ func GetJsonString(data any) string {
 	return string(b)
 }
 
+// MaskEmail masks a user email to prevent PII leakage in logs
+// Returns "***masked***" if email is empty, otherwise shows only the domain part
+func MaskEmail(email string) string {
+	if email == "" {
+		return "***masked***"
+	}
+
+	// Find the @ symbol
+	atIndex := strings.Index(email, "@")
+	if atIndex == -1 {
+		// No @ symbol found, return masked
+		return "***masked***"
+	}
+
+	// Return only the domain part with @ symbol
+	return "***@" + email[atIndex+1:]
+}
+
 // MaskSensitiveInfo masks sensitive information like URLs, IPs, and domain names in a string
 // Example:
 // http://example.com -> http://***.com
diff --git a/controller/channel-test.go b/controller/channel-test.go
index 95a4313f..ea37c5bf 100644
--- a/controller/channel-test.go
+++ b/controller/channel-test.go
@@ -134,7 +134,13 @@ func testChannel(channel *model.Channel, testModel string) testResult {
 	}
 	request := buildTestRequest(testModel)
 
-	info, err := relaycommon.GenRelayInfo(c, types.RelayFormatOpenAI, request, nil)
+	// Determine relay format based on request path
+	relayFormat := types.RelayFormatOpenAI
+	if c.Request.URL.Path == "/v1/embeddings" {
+		relayFormat = types.RelayFormatEmbedding
+	}
+
+	info, err := relaycommon.GenRelayInfo(c, relayFormat, request, nil)
 
 	if err != nil {
 		return testResult{
diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go
index 31f9ec6d..1ebb0581 100644
--- a/relay/common/relay_info.go
+++ b/relay/common/relay_info.go
@@ -178,7 +178,7 @@ func (info *RelayInfo) ToString() string {
 
 	// User & token info (mask secrets)
 	fmt.Fprintf(b, "User{ Id: %d, Email: %q, Group: %q, UsingGroup: %q, Quota: %d }, ",
-		info.UserId, info.UserEmail, info.UserGroup, info.UsingGroup, info.UserQuota)
+		info.UserId, common.MaskEmail(info.UserEmail), info.UserGroup, info.UsingGroup, info.UserQuota)
 	fmt.Fprintf(b, "Token{ Id: %d, Unlimited: %t, Key: ***masked*** }, ", info.TokenId, info.TokenUnlimited)
 
 	// Time info
diff --git a/service/pre_consume_quota.go b/service/pre_consume_quota.go
index ef466d8d..964ab665 100644
--- a/service/pre_consume_quota.go
+++ b/service/pre_consume_quota.go
@@ -3,14 +3,15 @@ package service
 import (
 	"errors"
 	"fmt"
-	"github.com/bytedance/gopkg/util/gopool"
-	"github.com/gin-gonic/gin"
 	"net/http"
 	"one-api/common"
 	"one-api/logger"
 	"one-api/model"
 	relaycommon "one-api/relay/common"
 	"one-api/types"
+
+	"github.com/bytedance/gopkg/util/gopool"
+	"github.com/gin-gonic/gin"
 )
 
 func ReturnPreConsumedQuota(c *gin.Context, relayInfo *relaycommon.RelayInfo, preConsumedQuota int) {