From dba74e5b05f05fc16fd4c909e3c5c27df45d3175 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 7d4cdaf0a9..a769b8e4fe 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 95a4313fa1..ea37c5bfde 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 31f9ec6d81..1ebb058150 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 ef466d8d44..964ab665c8 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) {