refactor: optimize the code structure and introduce log auditing

Author: dancing-ui

core/llm_token_ratelimit/config.go

@@ -86,7 +86,7 @@ type KeyItem struct {
 	Time  Time   `json:"time" yaml:"time"`
 }
 
-type RuleItem struct {
+type SpecificItem struct {
 	Identifier Identifier `json:"identifier" yaml:"identifier"`
 	KeyItems   []*KeyItem `json:"keyItems" yaml:"keyItems"`
 }
@@ -332,18 +332,18 @@ func (p TokenEncoderProvider) String() string {
 	}
 }
 
-func (ri *RuleItem) String() string {
-	if ri == nil {
-		return "RuleItem{nil}"
+func (si *SpecificItem) String() string {
+	if si == nil {
+		return "SpecificItem{nil}"
 	}
 
 	var sb strings.Builder
-	sb.WriteString("RuleItem{")
-	sb.WriteString(fmt.Sprintf("Identifier:%s", ri.Identifier.String()))
+	sb.WriteString("SpecificItem{")
+	sb.WriteString(fmt.Sprintf("Identifier:%s", si.Identifier.String()))
 
-	if len(ri.KeyItems) > 0 {
+	if len(si.KeyItems) > 0 {
 		sb.WriteString(", KeyItems:[")
-		for i, item := range ri.KeyItems {
+		for i, item := range si.KeyItems {
 			if i > 0 {
 				sb.WriteString(", ")
 			}

core/llm_token_ratelimit/config_test.go

@@ -359,24 +359,24 @@ func TestKeyItem_String(t *testing.T) {
 	}
 }
 
-func TestRuleItem_String(t *testing.T) {
+func TestSpecificItem_String(t *testing.T) {
 	tests := []struct {
 		name     string
-		ri       *RuleItem
+		ri       *SpecificItem
 		expected string
 	}{
-		{"nil ruleitem", nil, "RuleItem{nil}"},
+		{"nil specificItem", nil, "SpecificItem{nil}"},
 		{
 			"empty keyitems",
-			&RuleItem{
+			&SpecificItem{
 				Identifier: Identifier{Type: AllIdentifier, Value: ".*"},
 				KeyItems:   []*KeyItem{},
 			},
-			"RuleItem{Identifier:Identifier{Type:all, Value:.*}, KeyItems:[]}",
+			"SpecificItem{Identifier:Identifier{Type:all, Value:.*}, KeyItems:[]}",
 		},
 		{
 			"single keyitem",
-			&RuleItem{
+			&SpecificItem{
 				Identifier: Identifier{Type: Header, Value: "user-id"},
 				KeyItems: []*KeyItem{
 					{
@@ -386,11 +386,11 @@ func TestRuleItem_String(t *testing.T) {
 					},
 				},
 			},
-			"RuleItem{Identifier:Identifier{Type:header, Value:user-id}, KeyItems:[KeyItem{Key:limit1, Token:Token{Number:1000, CountStrategy:total-tokens}, Time:Time{Value:3600 second}}]}",
+			"SpecificItem{Identifier:Identifier{Type:header, Value:user-id}, KeyItems:[KeyItem{Key:limit1, Token:Token{Number:1000, CountStrategy:total-tokens}, Time:Time{Value:3600 second}}]}",
 		},
 		{
 			"multiple keyitems",
-			&RuleItem{
+			&SpecificItem{
 				Identifier: Identifier{Type: Header, Value: "api-key"},
 				KeyItems: []*KeyItem{
 					{
@@ -405,7 +405,7 @@ func TestRuleItem_String(t *testing.T) {
 					},
 				},
 			},
-			"RuleItem{Identifier:Identifier{Type:header, Value:api-key}, KeyItems:[KeyItem{Key:limit1, Token:Token{Number:500, CountStrategy:input-tokens}, Time:Time{Value:1800 second}}, KeyItem{Key:limit2, Token:Token{Number:300, CountStrategy:output-tokens}, Time:Time{Value:3600 second}}]}",
+			"SpecificItem{Identifier:Identifier{Type:header, Value:api-key}, KeyItems:[KeyItem{Key:limit1, Token:Token{Number:500, CountStrategy:input-tokens}, Time:Time{Value:1800 second}}, KeyItem{Key:limit2, Token:Token{Number:300, CountStrategy:output-tokens}, Time:Time{Value:3600 second}}]}",
 		},
 	}
 
@@ -425,8 +425,7 @@ rules:
   - id: "rule1"
     resource: "/api/chat"
     strategy: fixed-window
-    ruleName: "chat-limit"
-    ruleItems:
+    specificItems:
       - identifier:
           type: header
           value: "user-id"
@@ -478,17 +477,17 @@ errorMessage: "Rate limit exceeded"
 		t.Errorf("Expected FixedWindow strategy, got %v", rule.Strategy)
 	}
 
-	if len(rule.RuleItems) != 1 {
-		t.Errorf("Expected 1 rule item, got %d", len(rule.RuleItems))
+	if len(rule.SpecificItems) != 1 {
+		t.Errorf("Expected 1 rule item, got %d", len(rule.SpecificItems))
 	}
 
-	ruleItem := rule.RuleItems[0]
-	if ruleItem.Identifier.Type != Header {
-		t.Errorf("Expected Header identifier type, got %v", ruleItem.Identifier.Type)
+	specificItem := rule.SpecificItems[0]
+	if specificItem.Identifier.Type != Header {
+		t.Errorf("Expected Header identifier type, got %v", specificItem.Identifier.Type)
 	}
 
-	if len(ruleItem.KeyItems) != 2 {
-		t.Errorf("Expected 2 key items, got %d", len(ruleItem.KeyItems))
+	if len(specificItem.KeyItems) != 2 {
+		t.Errorf("Expected 2 key items, got %d", len(specificItem.KeyItems))
 	}
 
 	// Test Redis config
@@ -516,7 +515,7 @@ func TestYAMLUnmarshalingErrors(t *testing.T) {
 			"invalid identifier type",
 			`
 rules:
-  - ruleItems:
+  - specificItems:
       - identifier:
           type: invalid
 `,
@@ -526,7 +525,7 @@ rules:
 			"invalid count strategy",
 			`
 rules:
-  - ruleItems:
+  - specificItems:
       - keyItems:
           - token:
               countStrategy: invalid-strategy
@@ -537,7 +536,7 @@ rules:
 			"invalid time unit",
 			`
 rules:
-  - ruleItems:
+  - specificItems:
       - keyItems:
           - time:
               unit: invalid-unit
@@ -600,8 +599,8 @@ func BenchmarkToken_String(b *testing.B) {
 	}
 }
 
-func BenchmarkRuleItem_String(b *testing.B) {
-	ri := &RuleItem{
+func BenchmarkSpecificItem_String(b *testing.B) {
+	ri := &SpecificItem{
 		Identifier: Identifier{Type: Header, Value: "user-id"},
 		KeyItems: []*KeyItem{
 			{

core/llm_token_ratelimit/constant.go

@@ -16,9 +16,7 @@ package llmtokenratelimit
 
 // ================================= Config ====================================
 const (
-	DefaultResource               string = "default-resource"
 	DefaultResourcePattern        string = ".*"
-	DefaultRuleName               string = "overall-rule"
 	DefaultIdentifierValuePattern string = ".*"
 	DefaultKeyPattern             string = ".*"
 
@@ -48,13 +46,13 @@ const (
 	KeyRequestInfos    string = "SentinelLLMTokenRatelimitReqInfos"
 	KeyUsedTokenInfos  string = "SentinelLLMTokenRatelimitUsedTokenInfos"
 	KeyMatchedRules    string = "SentinelLLMTokenRatelimitMatchedRules"
-	KeyLLMPrompts      string = "SentinelLLMTokenRatelimitLLMPrompts"
 	KeyResponseHeaders string = "SentinelLLMTokenRatelimitResponseHeaders"
+	KeyRequestID       string = "SentinelLLMTokenRatelimitRequestID"
 )
 
 // ================================= RedisRatelimitKeyFormat ==================
 const (
-	RedisRatelimitKeyFormat string = "sentinel-go:llm-token-ratelimit:%s:%s:%s:%d:%s" // ruleName, strategy, identifierType, timeWindow, tokenCountStrategy
+	RedisRatelimitKeyFormat string = "sentinel-go:llm-token-ratelimit:resource-%s:%s:%s:%d:%s" // hashedResource, strategy, identifierType, timeWindow, tokenCountStrategy
 )
 
 // ================================= ResponseHeader ==================
@@ -67,11 +65,13 @@ const (
 
 // ================================= PETAStrategy =============================
 const (
-	PETANoWaiting              int64  = 0
-	PETACorrectOK              int64  = 1
-	PETASlidingWindowKeyFormat string = "{shard-%s}:sliding-window:%s" // hashTag, redisRatelimitKey
-	PETATokenBucketKeyFormat   string = "{shard-%s}:token-bucket:%s"   // hashTag, redisRatelimitKey
-	PETARandomStringLength     int    = 16
+	PETANoWaiting                 int64  = 0
+	PETACorrectOK                 int64  = 0
+	PETACorrectUnderestimateError int64  = 1
+	PETACorrectOverestimateError  int64  = 2
+	PETASlidingWindowKeyFormat    string = "{shard-%s}:sliding-window:%s" // hashTag, redisRatelimitKey
+	PETATokenBucketKeyFormat      string = "{shard-%s}:token-bucket:%s"   // hashTag, redisRatelimitKey
+	PETARandomStringLength        int    = 16
 )
 
 // ================================= Generate Random String ===================
@@ -86,108 +86,3 @@ const (
 const (
 	TokenEncoderKeyFormat string = "%s:token-encoder:%s:%s" // redisRatelimitKey, provider, model
 )
-
-// ================================= RedisKeyForbiddenChars ===================
-var RedisKeyForbiddenChars = map[string]string{
-	// Control characters
-	" ":    "space",
-	"\n":   "newline",
-	"\r":   "carriage return",
-	"\t":   "tab",
-	"\x00": "null byte",
-	"\x01": "start of heading",
-	"\x02": "start of text",
-	"\x03": "end of text",
-	"\x04": "end of transmission",
-	"\x05": "enquiry",
-	"\x06": "acknowledge",
-	"\x07": "bell",
-	"\x08": "backspace",
-	"\x0B": "vertical tab",
-	"\x0C": "form feed",
-	"\x0E": "shift out",
-	"\x0F": "shift in",
-	"\x10": "data link escape",
-	"\x11": "device control 1",
-	"\x12": "device control 2",
-	"\x13": "device control 3",
-	"\x14": "device control 4",
-	"\x15": "negative acknowledge",
-	"\x16": "synchronous idle",
-	"\x17": "end of transmission block",
-	"\x18": "cancel",
-	"\x19": "end of medium",
-	"\x1A": "substitute",
-	"\x1B": "escape",
-	"\x1C": "file separator",
-	"\x1D": "group separator",
-	"\x1E": "record separator",
-	"\x1F": "unit separator",
-	"\x7F": "delete",
-
-	// Special characters that may cause issues
-	"*":  "asterisk (reserved for wildcard)",
-	"?":  "question mark (reserved for wildcard)",
-	"[":  "left bracket (reserved for character class)",
-	"]":  "right bracket (reserved for character class)",
-	"{":  "left brace (reserved for quantifier)",
-	"}":  "right brace (reserved for quantifier)",
-	"(":  "left parenthesis (reserved for grouping)",
-	")":  "right parenthesis (reserved for grouping)",
-	"|":  "pipe (reserved for alternation)",
-	"^":  "caret (reserved for start anchor)",
-	"$":  "dollar (reserved for end anchor)",
-	"+":  "plus (reserved for quantifier)",
-	".":  "dot (reserved for any character)",
-	"\\": "backslash (reserved for escape)",
-	"/":  "forward slash (potential path separator)",
-
-	// Redis protocol related characters
-	"\r\n": "CRLF sequence",
-
-	// Unicode control characters (partial)
-	"\u0080": "padding character",
-	"\u0081": "high octet preset",
-	"\u0082": "break permitted here",
-	"\u0083": "no break here",
-	"\u0084": "index",
-	"\u0085": "next line",
-	"\u0086": "start of selected area",
-	"\u0087": "end of selected area",
-	"\u0088": "character tabulation set",
-	"\u0089": "character tabulation with justification",
-	"\u008A": "line tabulation set",
-	"\u008B": "partial line forward",
-	"\u008C": "partial line backward",
-	"\u008D": "reverse line feed",
-	"\u008E": "single shift two",
-	"\u008F": "single shift three",
-	"\u0090": "device control string",
-	"\u0091": "private use one",
-	"\u0092": "private use two",
-	"\u0093": "set transmit state",
-	"\u0094": "cancel character",
-	"\u0095": "message waiting",
-	"\u0096": "start of guarded area",
-	"\u0097": "end of guarded area",
-	"\u0098": "start of string",
-	"\u0099": "single graphic character introducer",
-	"\u009A": "single character introducer",
-	"\u009B": "control sequence introducer",
-	"\u009C": "string terminator",
-	"\u009D": "operating system command",
-	"\u009E": "privacy message",
-	"\u009F": "application program command",
-
-	// Other potentially problematic characters
-	"\"": "double quote (potential JSON escape)",
-	"'":  "single quote (potential SQL injection)",
-	"`":  "backtick (potential command injection)",
-	"<":  "less than (potential XSS)",
-	">":  "greater than (potential XSS)",
-	"&":  "ampersand (potential HTML entity)",
-	";":  "semicolon (potential command separator)",
-	":":  "colon (potential protocol separator)",
-	"=":  "equals (potential assignment)",
-	",":  "comma (potential CSV separator)",
-}

core/llm_token_ratelimit/context.go

@@ -57,29 +57,13 @@ func (ctx *Context) Get(key string) interface{} {
 	return nil
 }
 
-func extractContextFromArgs(ctx *base.EntryContext) *Context {
-	if ctx == nil || ctx.Input == nil || ctx.Input.Args == nil {
-		return nil
-	}
-	for _, arg := range ctx.Input.Args {
-		if llmCtx, ok := arg.(*Context); ok && llmCtx != nil {
-			return llmCtx
-		}
-	}
-	return nil
-}
-
-func extractContextFromData(ctx *base.EntryContext) *Context {
-	if ctx == nil || ctx.Data == nil {
-		return nil
+func (ctx *Context) extractArgs(entryCtx *base.EntryContext) {
+	if ctx == nil || entryCtx == nil || entryCtx.Input == nil || entryCtx.Input.Args == nil {
+		return
 	}
-	for key, value := range ctx.Data {
-		if key != KeyContext {
-			continue
-		}
-		if llmCtx, ok := value.(*Context); ok && llmCtx != nil {
-			return llmCtx
+	for _, arg := range entryCtx.Input.Args {
+		if reqInfos, ok := arg.(*RequestInfos); ok && reqInfos != nil {
+			ctx.Set(KeyRequestInfos, reqInfos)
 		}
 	}
-	return nil
 }