Selaa lähdekoodia

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.
CaIon 4 kuukautta sitten
vanhempi
sitoutus
03fc89da00
4 muutettua tiedostoa jossa 29 lisäystä ja 4 poistoa
  1. 18 0
      common/str.go
  2. 7 1
      controller/channel-test.go
  3. 1 1
      relay/common/relay_info.go
  4. 3 2
      service/pre_consume_quota.go

+ 18 - 0
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

+ 7 - 1
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{

+ 1 - 1
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

+ 3 - 2
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) {