Explorar o código

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 hai 6 meses
pai
achega
03fc89da00
Modificáronse 4 ficheiros con 29 adicións e 4 borrados
  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) {