mirror of
https://github.com/alibaba/higress.git
synced 2026-07-24 05:10:34 +08:00
feat(ai-token,cluster-key-ratelimit): multi-rule OR overlay (breaking change) (#4011)
Signed-off-by: geekspeng <geekspeng@icloud.com> Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
3
.gitignore
vendored
3
.gitignore
vendored
@@ -17,3 +17,6 @@ target/
|
||||
tools/hack/cluster.conf
|
||||
envoy/1.20
|
||||
istio/1.12
|
||||
|
||||
# Local working notes (design specs, implementation plans) — not for upstream PRs
|
||||
docs/superpowers/
|
||||
|
||||
@@ -4,12 +4,23 @@ keywords: [ AI网关, AI token限流 ]
|
||||
description: AI Token限流插件配置参考
|
||||
---
|
||||
|
||||
> ⚠️ **行为变更提示(无版本号变化)**
|
||||
>
|
||||
> 自本次更新起,`rule_items` 的匹配语义从 **first-match-wins**(命中第一条即返回)改为 **all-match OR 叠加**(所有命中规则都评估,任一触发即拒绝)。同时解除 `global_threshold` 与 `rule_items` 的互斥约束,支持混合配置。
|
||||
>
|
||||
> - 老配置(单条 `rule_items` 或仅 `global_threshold`):行为不变
|
||||
> - 老配置(多条 `rule_items` 期望短路匹配):**行为会变** —— 所有命中规则都会评估
|
||||
> - Redis key 格式新增 `{rule_name}` hash tag 以兼容 Redis Cluster,老计数器数据不兼容
|
||||
>
|
||||
> 详见下方"功能说明"与"配置示例"。
|
||||
|
||||
## 功能说明
|
||||
|
||||
`ai-token-ratelimit`插件基于 Redis 实现了 AI Token 限流功能,支持以下两种限流模式:
|
||||
`ai-token-ratelimit`插件基于 Redis 实现了 AI Token 限流功能,支持以下三种限流模式:
|
||||
|
||||
- **规则级全局限流**:依据相同的`rule_name`与`global_threshold`配置,为自定义规则组设置全局 token 限流阈值
|
||||
- **Key 级动态限流**:根据请求中的动态 Key(包括 URL 参数、请求头、客户端 IP、Consumer 名称或 Cookie 字段等)进行分组 token 限流
|
||||
- **混合限流**:同时配置 `global_threshold`(全局兜底)和 `rule_items`(按维度细分),所有命中的规则**叠加生效**,任一触发即拒绝。
|
||||
|
||||
## 运行属性
|
||||
|
||||
@@ -21,8 +32,8 @@ description: AI Token限流插件配置参考
|
||||
| 配置项 | 类型 | 必填 | 默认值 | 说明 |
|
||||
| ----------------------- | ------ | ---- | ------ |---------------------------------------------------------------------------|
|
||||
| rule_name | string | 是 | - | 限流规则名称,根据限流规则名称+限流类型+限流key名称+限流key对应的实际值来拼装redis key |
|
||||
| global_threshold | Object | 否,`global_threshold` 或 `rule_items` 选填一项 | - | 对整个自定义规则组进行限流 |
|
||||
| rule_items | array of object | 否,`global_threshold` 或 `rule_items` 选填一项 | - | 限流规则项,按照rule_items下的排列顺序,匹配第一个rule_item后命中限流规则,后续规则将被忽略 |
|
||||
| global_threshold | Object | 否,至少一项;可同时配置 | - | 对整个自定义规则组进行限流 |
|
||||
| rule_items | array of object | 否,至少一项;可同时配置 | - | 限流规则项,最多支持 **10 条**。所有满足匹配条件的 `rule_item` 都会参与限流,规则之间是"或"关系,任一触发即拒绝;规则的执行顺序不影响最终结果。详见下方"配置说明"展开。 |
|
||||
| rejected_code | int | 否 | 429 | 请求被限流时,返回的HTTP状态码 |
|
||||
| rejected_msg | string | 否 | Too many requests | 请求被限流时,返回的响应体 |
|
||||
| redis | object | 是 | - | redis相关配置 |
|
||||
@@ -51,6 +62,14 @@ description: AI Token限流插件配置参考
|
||||
| limit_by_per_ip | string | 否,`limit_by_*`中选填一项 | - | 按规则匹配特定 IP,并对每个 IP 分别计算限流,配置获取限流键值的来源 IP 参数名称,从请求头获取,以`from-header-对应的header名`,示例:`from-header-x-forwarded-for`,直接获取对端socket ip,配置为`from-remote-addr` |
|
||||
| limit_keys | array of object | 是 | - | 配置匹配键值后的限流次数 |
|
||||
|
||||
#### `rule_items` 多规则匹配语义
|
||||
|
||||
`rule_items` 是一个数组,**所有满足匹配条件**的 `rule_item` 都会被评估(均发起 Redis 计数器检查),规则之间是"或"关系,**任一触发即拒绝**。
|
||||
|
||||
> **注意:** 拒绝时按 `global_threshold`(如有)→ `rule_items`(数组顺序)的优先级**取最先触发的规则**报告响应头(`X-RateLimit-Reset` 等),而非"所有触发规则叠加报告"。因此,若同时配置了 `global_threshold` 和 `rule_items`,`global_threshold` 触发时总会优先报告,不会被 `rule_items` 覆盖。
|
||||
|
||||
> `rule_items` 数组最多支持 **10 条**规则。每条 `rule_item` 会按命中的 `limit_keys` 产生独立的 Redis 计数器。
|
||||
|
||||
`limit_keys`中每一项的配置字段说明。
|
||||
|
||||
| 配置项 | 类型 | 必填 | 默认值 | 说明 |
|
||||
|
||||
@@ -4,12 +4,23 @@ keywords: [ AI Gateway, AI Token Rate Limiting ]
|
||||
description: AI Token Rate Limiting Plugin Configuration Reference
|
||||
---
|
||||
|
||||
> ⚠️ **Behavior Change Notice (no version bump)**
|
||||
>
|
||||
> As of this update, `rule_items` matching semantics changed from **first-match-wins** (returns on first hit) to **all-match OR overlay** (all matched rules are evaluated, any trigger rejects). The mutual exclusion between `global_threshold` and `rule_items` is also removed to support hybrid configuration.
|
||||
>
|
||||
> - Existing config (single `rule_items` or only `global_threshold`): behavior unchanged
|
||||
> - Existing config (multiple `rule_items` expecting short-circuit match): **behavior will change** — all matched rules are evaluated
|
||||
> - Redis key format adds `{rule_name}` hash tag for Redis Cluster compatibility; old counter data is incompatible
|
||||
>
|
||||
> See "Features" and "Configuration Examples" below for details.
|
||||
|
||||
## Function Description
|
||||
|
||||
The `ai-token-ratelimit` plugin implements AI Token rate limiting based on Redis, supporting the following two rate limiting modes:
|
||||
The `ai-token-ratelimit` plugin implements AI Token rate limiting based on Redis, supporting the following three rate limiting modes:
|
||||
|
||||
- **Rule-level Global Rate Limiting**: Sets a global token rate limit threshold for custom rule groups based on the same `rule_name` and `global_threshold` configurations.
|
||||
- **Key-level Dynamic Rate Limiting**: Performs grouped token rate limiting based on dynamic keys in requests (including URL parameters, request headers, client IP, Consumer name, or Cookie fields, etc.).
|
||||
- **Hybrid rate limiting**: Configure `global_threshold` (global fallback) and `rule_items` (per-dimension) simultaneously. All matched rules take effect together; any trigger rejects the request.
|
||||
|
||||
|
||||
## Runtime Properties
|
||||
@@ -23,8 +34,8 @@ Plugin execution priority: `600`
|
||||
| Configuration Item | Type | Required | Default Value | Description |
|
||||
|--------------------------|----------------|----------|---------------|-------------------------------------------------------------------------------------------------|
|
||||
| rule_name | string | Yes | - | Name of the rate limiting rule. The Redis key is assembled based on the rate limiting rule name + rate limiting type + rate limiting key name + actual value corresponding to the rate limiting key. |
|
||||
| global_threshold | Object | No, either `global_threshold` or `rule_items` is required | - | Rate limits the entire custom rule group |
|
||||
| rule_items | array of object| No, either `global_threshold` or `rule_items` is required | - | Rate limiting rule items. The first matching `rule_item` in the order of `rule_items` triggers the rate limiting rule, and subsequent rules are ignored. |
|
||||
| global_threshold | Object | No, at least one of them is required; can be configured simultaneously | - | Rate limits the entire custom rule group |
|
||||
| rule_items | array of object| No, at least one of them is required; can be configured simultaneously | - | Rate limiting rule items. Supports up to **10** rules. All matched `rule_item`s are evaluated and combined with an OR relationship; any trigger rejects the request. The execution order of rules does not affect the final result. See the expanded `rule_items` notes below for details. |
|
||||
| rejected_code | int | No | 429 | HTTP status code returned when a request is rate-limited |
|
||||
| rejected_msg | string | No | Too many requests | Response body returned when a request is rate-limited |
|
||||
| redis | object | Yes | - | Redis-related configurations |
|
||||
@@ -55,6 +66,14 @@ Plugin execution priority: `600`
|
||||
| limit_by_per_ip | string | No, one of `limit_by_*` is required | - | Matches specific IPs by rule and calculates rate limits for each IP separately. Configures the source of the rate limiting key value as the IP parameter name, obtained from the request header in the format `from-header-corresponding_header_name` (e.g., `from-header-x-forwarded-for`), or directly obtains the peer socket IP by configuring `from-remote-addr`. |
|
||||
| limit_keys | array of object | Yes | - | Configures the rate limiting count after matching the key value |
|
||||
|
||||
#### `rule_items` Multi-Rule Matching Semantics
|
||||
|
||||
`rule_items` is an array. **All** `rule_item`s whose match conditions are satisfied are evaluated (each triggers an independent Redis counter check), and the rules are combined with an OR relationship: **any trigger rejects the request**.
|
||||
|
||||
> **Note:** On rejection, the response headers (`X-RateLimit-Reset`, etc.) are reported from the **first triggered rule** in priority order: `global_threshold` (if present) → `rule_items` (in array order). This is **not** an "all triggered rules combined" report. Therefore, when both `global_threshold` and `rule_items` are configured, a triggered `global_threshold` always takes reporting priority over `rule_items`.
|
||||
|
||||
> The `rule_items` array supports up to **10** rules. Each `rule_item` produces independent Redis counters for each matched `limit_keys`.
|
||||
|
||||
|
||||
### Description of Configuration Fields in `limit_keys`
|
||||
|
||||
|
||||
@@ -46,6 +46,9 @@ const (
|
||||
SecondsPerMinute = 60 * Second
|
||||
SecondsPerHour = 60 * SecondsPerMinute
|
||||
SecondsPerDay = 24 * SecondsPerHour
|
||||
|
||||
// MaxRuleItems 限制 rule_items 数组最大长度,对齐商业版"最多同时命中 10 条规则"。
|
||||
MaxRuleItems = 10
|
||||
)
|
||||
|
||||
var timeWindows = map[string]int64{
|
||||
@@ -181,8 +184,6 @@ func initLimitRule(json gjson.Result, config *AiTokenRateLimitConfig) error {
|
||||
hasRule := ruleItemsResult.Exists()
|
||||
if !hasGlobal && !hasRule {
|
||||
return errors.New("at least one of 'global_threshold' or 'rule_items' must be set")
|
||||
} else if hasGlobal && hasRule {
|
||||
return errors.New("'global_threshold' and 'rule_items' cannot be set at the same time")
|
||||
}
|
||||
|
||||
// 处理全局限流配置
|
||||
@@ -192,14 +193,19 @@ func initLimitRule(json gjson.Result, config *AiTokenRateLimitConfig) error {
|
||||
return fmt.Errorf("failed to parse global_threshold: %w", err)
|
||||
}
|
||||
config.GlobalThreshold = threshold
|
||||
return nil
|
||||
}
|
||||
|
||||
// 处理条件限流规则
|
||||
if !hasRule {
|
||||
return nil
|
||||
}
|
||||
items := ruleItemsResult.Array()
|
||||
if len(items) == 0 {
|
||||
return errors.New("config rule_items cannot be empty")
|
||||
}
|
||||
if len(items) > MaxRuleItems {
|
||||
return fmt.Errorf("rule_items length %d exceeds maximum %d", len(items), MaxRuleItems)
|
||||
}
|
||||
|
||||
var ruleItems []LimitRuleItem
|
||||
// 用于记录已出现的LimitType和Key的组合
|
||||
@@ -298,7 +304,7 @@ func parseLimitRuleItem(item gjson.Result) (*LimitRuleItem, error) {
|
||||
}
|
||||
|
||||
if limitType == "" {
|
||||
return nil, errors.New("only one of 'limit_by_header' and 'limit_by_param' and 'limit_by_consumer' and 'limit_by_cookie' and 'limit_by_per_header' and 'limit_by_per_param' and 'limit_by_per_consumer' and 'limit_by_per_cookie' and 'limit_by_per_ip' can be set")
|
||||
return nil, errors.New("at least one of 'limit_by_header', 'limit_by_param', 'limit_by_consumer', 'limit_by_cookie', 'limit_by_per_header', 'limit_by_per_param', 'limit_by_per_consumer', 'limit_by_per_cookie', 'limit_by_per_ip' must be set")
|
||||
}
|
||||
ruleItem.LimitType = limitType
|
||||
|
||||
|
||||
@@ -2,6 +2,8 @@ package config
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
@@ -166,13 +168,40 @@ func TestParseAiTokenRateLimitConfig(t *testing.T) {
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "Conflict_GlobalThresholdAndRuleItems",
|
||||
name: "Both_ShouldSucceed",
|
||||
json: `{
|
||||
"rule_name": "test-conflict",
|
||||
"rule_name": "test-both",
|
||||
"global_threshold": {"token_per_second": 100},
|
||||
"rule_items": [{"limit_by_header": "x-test"}]
|
||||
"rule_items": [
|
||||
{
|
||||
"limit_by_header": "x-test",
|
||||
"limit_keys": [{"key": "k1", "token_per_minute": 10}]
|
||||
}
|
||||
]
|
||||
}`,
|
||||
expectedErr: errors.New("'global_threshold' and 'rule_items' cannot be set at the same time"),
|
||||
expected: AiTokenRateLimitConfig{
|
||||
RuleName: "test-both",
|
||||
GlobalThreshold: &GlobalThreshold{
|
||||
Count: 100,
|
||||
TimeWindow: Second,
|
||||
},
|
||||
RuleItems: []LimitRuleItem{
|
||||
{
|
||||
LimitType: LimitByHeaderType,
|
||||
Key: "x-test",
|
||||
ConfigItems: []LimitConfigItem{
|
||||
{
|
||||
ConfigType: ExactType,
|
||||
Key: "k1",
|
||||
Count: 10,
|
||||
TimeWindow: SecondsPerMinute,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
RejectedCode: DefaultRejectedCode,
|
||||
RejectedMsg: DefaultRejectedMsg,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "Missing_GlobalThresholdAndRuleItems",
|
||||
@@ -199,6 +228,16 @@ func TestParseAiTokenRateLimitConfig(t *testing.T) {
|
||||
RejectedMsg: "Forbidden",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "LimitByType_AtLeastOne_Error",
|
||||
json: `{
|
||||
"rule_name": "no-limit-by",
|
||||
"rule_items": [
|
||||
{"limit_keys": [{"key": "k", "token_per_second": 1}]}
|
||||
]
|
||||
}`,
|
||||
expectedErr: errors.New("failed to parse rule_item in rule_items: at least one of 'limit_by_header', 'limit_by_param', 'limit_by_consumer', 'limit_by_cookie', 'limit_by_per_header', 'limit_by_per_param', 'limit_by_per_consumer', 'limit_by_per_cookie', 'limit_by_per_ip' must be set"),
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
@@ -216,3 +255,32 @@ func TestParseAiTokenRateLimitConfig(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseAiTokenRateLimitConfig_RuleItemsLimit(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
itemCnt int
|
||||
expectErr string
|
||||
}{
|
||||
{"Exactly10", 10, ""},
|
||||
{"Over10", 11, "rule_items length 11 exceeds maximum 10"},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
items := make([]string, 0, tt.itemCnt)
|
||||
for i := 0; i < tt.itemCnt; i++ {
|
||||
items = append(items,
|
||||
fmt.Sprintf(`{"limit_by_header":"h%d","limit_keys":[{"key":"k","token_per_second":1}]}`, i+1))
|
||||
}
|
||||
jsonStr := fmt.Sprintf(`{"rule_name":"t","rule_items":[%s]}`, strings.Join(items, ","))
|
||||
var config AiTokenRateLimitConfig
|
||||
err := ParseAiTokenRateLimitConfig(gjson.Parse(jsonStr), &config)
|
||||
if tt.expectErr == "" {
|
||||
assert.NoError(t, err)
|
||||
assert.Len(t, config.RuleItems, tt.itemCnt)
|
||||
} else {
|
||||
assert.EqualError(t, err, tt.expectErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,8 +6,6 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20250822030947-8345453fddd0 h1:YGdj8KBzVjabU3STUfwMZghB+VlX6YLfJtLbrsWaOD0=
|
||||
github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20250822030947-8345453fddd0/go.mod h1:tRI2LfMudSkKHhyv1uex3BWzcice2s/l8Ah8axporfA=
|
||||
github.com/higress-group/wasm-go v1.0.2-0.20250821081215-b573359becf8 h1:rs+AH1wfZy4swzuAyiRXT7xPUm8gycXt9Gwy0tqOq0o=
|
||||
github.com/higress-group/wasm-go v1.0.2-0.20250821081215-b573359becf8/go.mod h1:9k7L730huS/q4V5iH9WLDgf5ZUHEtfhM/uXcegKDG/M=
|
||||
github.com/higress-group/wasm-go v1.0.6-0.20251103065747-41d65dbb2f9e h1:wYW/DXjyQniQLaB26c+J9NQk3+AhqByzS1r18NShvB4=
|
||||
github.com/higress-group/wasm-go v1.0.6-0.20251103065747-41d65dbb2f9e/go.mod h1:B8C6+OlpnyYyZUBEdUXA7tYZYD+uwZTNjfkE5FywA+A=
|
||||
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
|
||||
|
||||
@@ -23,6 +23,7 @@ import (
|
||||
|
||||
"ai-token-ratelimit/config"
|
||||
"ai-token-ratelimit/util"
|
||||
|
||||
"github.com/higress-group/proxy-wasm-go-sdk/proxywasm"
|
||||
"github.com/higress-group/proxy-wasm-go-sdk/proxywasm/types"
|
||||
"github.com/higress-group/wasm-go/pkg/log"
|
||||
@@ -45,55 +46,59 @@ func init() {
|
||||
|
||||
const (
|
||||
RedisKeyPrefix string = "higress-token-ratelimit"
|
||||
// AiTokenGlobalRateLimitFormat 全局限流模式 redis key 为 RedisKeyPrefix:限流规则名称:global_threshold:时间窗口
|
||||
AiTokenGlobalRateLimitFormat = RedisKeyPrefix + ":%s:global_threshold:%d"
|
||||
// AiTokenRateLimitFormat 规则限流模式 redis key 为 RedisKeyPrefix:限流规则名称:限流类型:时间窗口:限流key名称:限流key对应的实际值
|
||||
AiTokenRateLimitFormat = RedisKeyPrefix + ":%s:%s:%d:%s:%s"
|
||||
RequestPhaseFixedWindowScript = `
|
||||
local current = redis.call('get', KEYS[1])
|
||||
local ttl = redis.call('ttl', KEYS[1])
|
||||
local threshold = tonumber(ARGV[1])
|
||||
local window = tonumber(ARGV[2])
|
||||
|
||||
-- 键不存在时,返回初始状态(计数0,窗口时间为过期时间)
|
||||
if not current then
|
||||
return {threshold, 0, window}
|
||||
end
|
||||
// 使用 {rule_name} hash tag 让多规则多键操作在 Redis Cluster 下落到同一 slot
|
||||
// AiTokenGlobalRateLimitFormat 全局限流模式 redis key 为 RedisKeyPrefix:{限流规则名称}:global_threshold:时间窗口
|
||||
AiTokenGlobalRateLimitFormat = RedisKeyPrefix + ":{%s}:global_threshold:%d"
|
||||
// AiTokenRateLimitFormat 规则限流模式 redis key 为 RedisKeyPrefix:{限流规则名称}:限流类型:时间窗口:限流key名称:限流key对应的实际值
|
||||
AiTokenRateLimitFormat = RedisKeyPrefix + ":{%s}:%s:%d:%s:%s"
|
||||
// MultiKeyRequestPhaseScript 多规则请求阶段只读检查脚本
|
||||
// KEYS = [key1, ..., keyN]
|
||||
// ARGV = [threshold1, window1, threshold2, window2, ..., thresholdN, windowN]
|
||||
// 返回嵌套数组 {{threshold_i, current_i, ttl_i}, ...}
|
||||
MultiKeyRequestPhaseScript = `
|
||||
local results = {}
|
||||
for i = 1, #KEYS do
|
||||
local threshold = tonumber(ARGV[2*i - 1])
|
||||
local window = tonumber(ARGV[2*i])
|
||||
local current = redis.call('get', KEYS[i])
|
||||
local ttl = redis.call('ttl', KEYS[i])
|
||||
|
||||
-- 修复异常过期时间(确保窗口有效)
|
||||
if ttl < 0 then
|
||||
ttl = window
|
||||
-- 键不存在时,返回初始状态(计数0,窗口时间为过期时间)
|
||||
if not current then
|
||||
table.insert(results, {threshold, 0, window})
|
||||
else
|
||||
-- 修复异常过期时间(确保窗口有效)
|
||||
if ttl < 0 then
|
||||
ttl = window
|
||||
end
|
||||
-- 返回窗口状态:阈值、当前计数、剩余时间
|
||||
table.insert(results, {threshold, tonumber(current), ttl})
|
||||
end
|
||||
end
|
||||
|
||||
-- 返回窗口状态:阈值、当前计数、剩余时间
|
||||
return {threshold, tonumber(current), ttl}
|
||||
return results
|
||||
`
|
||||
// MultiKeyResponsePhaseScript 多规则响应阶段累加脚本(仅 ai-token-ratelimit 使用)
|
||||
// KEYS = [key1, ..., keyN]
|
||||
// ARGV = [threshold1, window1, count1, ..., thresholdN, windowN, countN]
|
||||
// 每条规则独立判断 current <= threshold 才累加;返回 KEYS 数量
|
||||
MultiKeyResponsePhaseScript = `
|
||||
for i = 1, #KEYS do
|
||||
local threshold = tonumber(ARGV[3*i - 2])
|
||||
local window = tonumber(ARGV[3*i - 1])
|
||||
local added = tonumber(ARGV[3*i])
|
||||
local current = tonumber(redis.call('get', KEYS[i]) or "0")
|
||||
if current <= threshold then
|
||||
current = redis.call('incrby', KEYS[i], added)
|
||||
if current == added then
|
||||
redis.call('expire', KEYS[i], window)
|
||||
else
|
||||
local ttl = redis.call('ttl', KEYS[i])
|
||||
if ttl < 0 then redis.call('expire', KEYS[i], window) end
|
||||
end
|
||||
end
|
||||
end
|
||||
return #KEYS
|
||||
`
|
||||
ResponsePhaseFixedWindowScript = `
|
||||
local key = KEYS[1]
|
||||
local threshold = tonumber(ARGV[1])
|
||||
local window = tonumber(ARGV[2])
|
||||
local added = tonumber(ARGV[3]) -- 需要累加的token数量
|
||||
|
||||
local current = tonumber(redis.call('get', key) or "0")
|
||||
|
||||
-- 只有当前计数未超过阈值时才执行累加
|
||||
if current <= threshold then
|
||||
current = redis.call('incrby', key, added)
|
||||
-- 第一次设置值时初始化过期时间
|
||||
if current == added then
|
||||
redis.call('expire', key, window)
|
||||
else
|
||||
-- 非首次设置时检查过期时间,确保窗口有效性
|
||||
local ttl = redis.call('ttl', key)
|
||||
if ttl < 0 then
|
||||
redis.call('expire', key, window)
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
-- 返回当前窗口状态:阈值、当前计数、剩余时间
|
||||
return {threshold, current, redis.call('ttl', key)}
|
||||
`
|
||||
|
||||
LimitRedisContextKey = "LimitRedisContext"
|
||||
|
||||
@@ -110,10 +115,16 @@ type LimitContext struct {
|
||||
reset int
|
||||
}
|
||||
|
||||
// MatchedRule 表示请求阶段命中的单条限流规则(global 或 rule_item)
|
||||
type MatchedRule struct {
|
||||
key string // 完整 Redis key
|
||||
count int64 // 时间窗口内的限额(与 LimitConfigItem.Count / GlobalThreshold.Count 同义)
|
||||
window int64 // 时间窗口大小(秒)
|
||||
}
|
||||
|
||||
// LimitRedisContext 暂存请求阶段命中的全部规则,供响应阶段多键 INCRBY 使用
|
||||
type LimitRedisContext struct {
|
||||
key string
|
||||
count int64
|
||||
window int64
|
||||
rules []MatchedRule
|
||||
}
|
||||
|
||||
func parseConfig(json gjson.Result, cfg *config.AiTokenRateLimitConfig) error {
|
||||
@@ -132,59 +143,73 @@ func parseConfig(json gjson.Result, cfg *config.AiTokenRateLimitConfig) error {
|
||||
|
||||
func onHttpRequestHeaders(ctx wrapper.HttpContext, cfg config.AiTokenRateLimitConfig) types.Action {
|
||||
ctx.DisableReroute()
|
||||
limitKey, count, timeWindow := "", int64(0), int64(0)
|
||||
|
||||
if cfg.GlobalThreshold != nil {
|
||||
// 全局限流模式
|
||||
limitKey = fmt.Sprintf(AiTokenGlobalRateLimitFormat, cfg.RuleName, cfg.GlobalThreshold.TimeWindow)
|
||||
count = cfg.GlobalThreshold.Count
|
||||
timeWindow = cfg.GlobalThreshold.TimeWindow
|
||||
} else {
|
||||
// 规则限流模式
|
||||
val, ruleItem, configItem := checkRequestAgainstLimitRule(ctx, cfg.RuleItems)
|
||||
if ruleItem == nil || configItem == nil {
|
||||
// 没有匹配到限流规则直接返回
|
||||
return types.ActionContinue
|
||||
}
|
||||
|
||||
limitKey = fmt.Sprintf(AiTokenRateLimitFormat, cfg.RuleName, ruleItem.LimitType, configItem.TimeWindow, ruleItem.Key, val)
|
||||
count = configItem.Count
|
||||
timeWindow = configItem.TimeWindow
|
||||
matched := collectMatchedRules(ctx, cfg)
|
||||
if len(matched) == 0 {
|
||||
// 无任何规则命中:直接放行,不发起 Redis 调用
|
||||
log.Debugf("ai-token-ratelimit: no rule matched, path=%s host=%s", ctx.Path(), ctx.Host())
|
||||
return types.ActionContinue
|
||||
}
|
||||
|
||||
ctx.SetContext(LimitRedisContextKey, LimitRedisContext{
|
||||
key: limitKey,
|
||||
count: count,
|
||||
window: timeWindow,
|
||||
})
|
||||
log.Debugf("ai-token-ratelimit: request phase matched %d rule(s), path=%s host=%s",
|
||||
len(matched), ctx.Path(), ctx.Host())
|
||||
for i, m := range matched {
|
||||
log.Debugf("ai-token-ratelimit: rule[%d] key=%s threshold=%d window=%ds",
|
||||
i, m.key, m.count, m.window)
|
||||
}
|
||||
|
||||
// 执行限流逻辑
|
||||
keys := []interface{}{limitKey}
|
||||
args := []interface{}{count, timeWindow}
|
||||
err := cfg.RedisClient.Eval(RequestPhaseFixedWindowScript, 1, keys, args, func(response resp.Value) {
|
||||
resultArray := response.Array()
|
||||
if len(resultArray) != 3 {
|
||||
log.Errorf("redis response parse error, response: %v", response)
|
||||
proxywasm.ResumeHttpRequest()
|
||||
n := len(matched)
|
||||
keys := make([]interface{}, n)
|
||||
args := make([]interface{}, 0, n*2)
|
||||
for i, m := range matched {
|
||||
keys[i] = m.key
|
||||
args = append(args, m.count, m.window)
|
||||
}
|
||||
|
||||
// 暂存命中规则,供响应阶段多键 INCRBY 使用
|
||||
ctx.SetContext(LimitRedisContextKey, LimitRedisContext{rules: matched})
|
||||
|
||||
err := cfg.RedisClient.Eval(MultiKeyRequestPhaseScript, n, keys, args, func(response resp.Value) {
|
||||
arr := response.Array()
|
||||
if len(arr) != n {
|
||||
log.Errorf("redis response length mismatch: got %d, want %d", len(arr), n)
|
||||
_ = proxywasm.ResumeHttpRequest()
|
||||
return
|
||||
}
|
||||
|
||||
// 获取限流结果
|
||||
threshold, current, ttl := resultArray[0].Integer(), resultArray[1].Integer(), resultArray[2].Integer()
|
||||
context := LimitContext{
|
||||
count: threshold,
|
||||
remaining: threshold - current,
|
||||
reset: ttl,
|
||||
}
|
||||
if current > threshold {
|
||||
// 触发限流
|
||||
ctx.SetUserAttribute("token_ratelimit_status", "limited")
|
||||
ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
|
||||
rejected(cfg, context)
|
||||
} else {
|
||||
proxywasm.ResumeHttpRequest()
|
||||
// 单次遍历:触发即 return;未触发则放行。
|
||||
// ai-token-ratelimit 不对外暴露 LimitContext(与 cluster-key-ratelimit 不同,
|
||||
// 后者通过 X-RateLimit-* 头可观测 tightest 选择),因此此处不再写入 Context。
|
||||
for i, ruleResult := range arr {
|
||||
ruleState := ruleResult.Array()
|
||||
if len(ruleState) != 3 {
|
||||
log.Errorf("redis sub-array length mismatch: got %d, want 3", len(ruleState))
|
||||
_ = proxywasm.ResumeHttpRequest()
|
||||
return
|
||||
}
|
||||
threshold, current, ttl := ruleState[0].Integer(), ruleState[1].Integer(), ruleState[2].Integer()
|
||||
log.Debugf("ai-token-ratelimit: eval rule[%d] key=%s threshold=%d current=%d ttl=%ds",
|
||||
i, matched[i].key, threshold, current, ttl)
|
||||
|
||||
if current > threshold {
|
||||
// 命中触发的第一条规则(按 collectMatchedRules 顺序,global 优先)
|
||||
log.Debugf("ai-token-ratelimit: rule[%d] key=%s triggered (current=%d > threshold=%d), rejecting with code %d",
|
||||
i, matched[i].key, current, threshold, cfg.RejectedCode)
|
||||
ctx.SetUserAttribute("token_ratelimit_status", "limited")
|
||||
_ = ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
|
||||
rejected(cfg, LimitContext{
|
||||
count: threshold,
|
||||
remaining: threshold - current,
|
||||
reset: ttl,
|
||||
})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
log.Debugf("ai-token-ratelimit: all %d rule(s) within threshold, allowing request to proceed", n)
|
||||
_ = proxywasm.ResumeHttpRequest()
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
log.Errorf("redis call failed: %v", err)
|
||||
return types.ActionContinue
|
||||
@@ -194,39 +219,90 @@ func onHttpRequestHeaders(ctx wrapper.HttpContext, cfg config.AiTokenRateLimitCo
|
||||
|
||||
func onHttpStreamingBody(ctx wrapper.HttpContext, cfg config.AiTokenRateLimitConfig, data []byte, endOfStream bool) []byte {
|
||||
if usage := tokenusage.GetTokenUsage(ctx, data); usage.TotalToken > 0 {
|
||||
log.Debugf("ai-token-ratelimit: token usage detected input=%d output=%d total=%d",
|
||||
usage.InputToken, usage.OutputToken, usage.TotalToken)
|
||||
ctx.SetContext(tokenusage.CtxKeyInputToken, usage.InputToken)
|
||||
ctx.SetContext(tokenusage.CtxKeyOutputToken, usage.OutputToken)
|
||||
}
|
||||
if endOfStream {
|
||||
if ctx.GetContext(tokenusage.CtxKeyInputToken) == nil || ctx.GetContext(tokenusage.CtxKeyOutputToken) == nil {
|
||||
return data
|
||||
}
|
||||
inputToken := ctx.GetContext(tokenusage.CtxKeyInputToken).(int64)
|
||||
outputToken := ctx.GetContext(tokenusage.CtxKeyOutputToken).(int64)
|
||||
limitRedisContext, ok := ctx.GetContext(LimitRedisContextKey).(LimitRedisContext)
|
||||
if !ok {
|
||||
return data
|
||||
}
|
||||
keys := []interface{}{limitRedisContext.key}
|
||||
args := []interface{}{limitRedisContext.count, limitRedisContext.window, inputToken + outputToken}
|
||||
err := cfg.RedisClient.Eval(ResponsePhaseFixedWindowScript, 1, keys, args, nil)
|
||||
if err != nil {
|
||||
log.Errorf("redis call failed: %v", err)
|
||||
}
|
||||
if !endOfStream {
|
||||
return data
|
||||
}
|
||||
|
||||
inputTokenRaw := ctx.GetContext(tokenusage.CtxKeyInputToken)
|
||||
outputTokenRaw := ctx.GetContext(tokenusage.CtxKeyOutputToken)
|
||||
if inputTokenRaw == nil || outputTokenRaw == nil {
|
||||
log.Debugf("ai-token-ratelimit: response phase end-of-stream reached but no token usage recorded, skipping")
|
||||
return data
|
||||
}
|
||||
inputToken, ok1 := inputTokenRaw.(int64)
|
||||
outputToken, ok2 := outputTokenRaw.(int64)
|
||||
if !ok1 || !ok2 {
|
||||
log.Errorf("ai-token-ratelimit: response phase token usage context has unexpected types: input=%T output=%T",
|
||||
inputTokenRaw, outputTokenRaw)
|
||||
return data
|
||||
}
|
||||
|
||||
limitRedisContextRaw := ctx.GetContext(LimitRedisContextKey)
|
||||
if limitRedisContextRaw == nil {
|
||||
log.Debugf("ai-token-ratelimit: response phase reached with no LimitRedisContext, skipping accumulation")
|
||||
return data
|
||||
}
|
||||
limitRedisContext, ok := limitRedisContextRaw.(LimitRedisContext)
|
||||
if !ok || len(limitRedisContext.rules) == 0 {
|
||||
log.Debugf("ai-token-ratelimit: response phase LimitRedisContext is empty, skipping accumulation")
|
||||
return data
|
||||
}
|
||||
|
||||
// 多键 INCRBY:每条规则一组 (threshold, window, added)
|
||||
n := len(limitRedisContext.rules)
|
||||
keys := make([]interface{}, n)
|
||||
args := make([]interface{}, 0, n*3)
|
||||
added := inputToken + outputToken
|
||||
log.Debugf("ai-token-ratelimit: response phase accumulating tokens input=%d output=%d total=%d across %d rule(s)",
|
||||
inputToken, outputToken, added, n)
|
||||
for i, r := range limitRedisContext.rules {
|
||||
keys[i] = r.key
|
||||
args = append(args, r.count, r.window, added)
|
||||
log.Debugf("ai-token-ratelimit: rule[%d] key=%s threshold=%d window=%ds added=%d",
|
||||
i, r.key, r.count, r.window, added)
|
||||
}
|
||||
|
||||
err := cfg.RedisClient.Eval(MultiKeyResponsePhaseScript, n, keys, args, func(response resp.Value) {
|
||||
incremented := response.Integer()
|
||||
log.Debugf("ai-token-ratelimit: response phase INCRBY done, %d/%d rule(s) updated", incremented, n)
|
||||
})
|
||||
if err != nil {
|
||||
log.Errorf("redis call failed: %v", err)
|
||||
}
|
||||
return data
|
||||
}
|
||||
|
||||
func checkRequestAgainstLimitRule(ctx wrapper.HttpContext, ruleItems []config.LimitRuleItem) (string, *config.LimitRuleItem, *config.LimitConfigItem) {
|
||||
if len(ruleItems) > 0 {
|
||||
for _, rule := range ruleItems {
|
||||
val, ruleItem, configItem := hitRateRuleItem(ctx, rule)
|
||||
if ruleItem != nil && configItem != nil {
|
||||
return val, ruleItem, configItem
|
||||
}
|
||||
// collectMatchedRules 遍历 global_threshold 和 rule_items,返回所有命中规则。
|
||||
// 顺序:global_threshold(如有)→ rule_items 中所有命中项(按数组顺序追加)。
|
||||
// 该顺序决定了"触发时优先报告哪条规则"以及"未触发时 tightest 的选择范围"。
|
||||
func collectMatchedRules(ctx wrapper.HttpContext, cfg config.AiTokenRateLimitConfig) []MatchedRule {
|
||||
var matched []MatchedRule
|
||||
|
||||
if cfg.GlobalThreshold != nil {
|
||||
matched = append(matched, MatchedRule{
|
||||
key: fmt.Sprintf(AiTokenGlobalRateLimitFormat, cfg.RuleName, cfg.GlobalThreshold.TimeWindow),
|
||||
count: cfg.GlobalThreshold.Count,
|
||||
window: cfg.GlobalThreshold.TimeWindow,
|
||||
})
|
||||
}
|
||||
|
||||
for _, ruleItem := range cfg.RuleItems {
|
||||
val, hitRule, hitItem := hitRateRuleItem(ctx, ruleItem)
|
||||
if hitRule != nil && hitItem != nil {
|
||||
matched = append(matched, MatchedRule{
|
||||
key: fmt.Sprintf(AiTokenRateLimitFormat, cfg.RuleName, hitRule.LimitType, hitItem.TimeWindow, hitRule.Key, val),
|
||||
count: hitItem.Count,
|
||||
window: hitItem.TimeWindow,
|
||||
})
|
||||
}
|
||||
}
|
||||
return "", nil, nil
|
||||
|
||||
return matched
|
||||
}
|
||||
|
||||
func hitRateRuleItem(ctx wrapper.HttpContext, rule config.LimitRuleItem) (string, *config.LimitRuleItem, *config.LimitConfigItem) {
|
||||
|
||||
@@ -16,11 +16,14 @@ package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/higress-group/proxy-wasm-go-sdk/proxywasm/types"
|
||||
"github.com/higress-group/wasm-go/pkg/test"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/resp"
|
||||
)
|
||||
|
||||
// 测试配置:全局限流配置
|
||||
@@ -191,6 +194,58 @@ var regexpLimitConfig = func() json.RawMessage {
|
||||
return data
|
||||
}()
|
||||
|
||||
// 测试配置:混合限流(global_threshold + rule_items 同时配置)
|
||||
var hybridLimitConfig = func() json.RawMessage {
|
||||
data, _ := json.Marshal(map[string]interface{}{
|
||||
"rule_name": "ai-token-hybrid",
|
||||
"global_threshold": map[string]interface{}{
|
||||
"token_per_minute": 10000,
|
||||
},
|
||||
"rule_items": []map[string]interface{}{
|
||||
{
|
||||
"limit_by_header": "x-api-key",
|
||||
"limit_keys": []map[string]interface{}{
|
||||
{
|
||||
"key": "vip-key",
|
||||
"token_per_minute": 100,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"redis": map[string]interface{}{
|
||||
"service_name": "redis.static",
|
||||
"service_port": 6379,
|
||||
},
|
||||
})
|
||||
return data
|
||||
}()
|
||||
|
||||
// 测试配置:多条 rule_items 同时命中
|
||||
var multiRuleItemsConfig = func() json.RawMessage {
|
||||
data, _ := json.Marshal(map[string]interface{}{
|
||||
"rule_name": "ai-token-multi-items",
|
||||
"rule_items": []map[string]interface{}{
|
||||
{
|
||||
"limit_by_header": "x-api-key",
|
||||
"limit_keys": []map[string]interface{}{
|
||||
{"key": "k1", "token_per_minute": 100},
|
||||
},
|
||||
},
|
||||
{
|
||||
"limit_by_param": "apikey",
|
||||
"limit_keys": []map[string]interface{}{
|
||||
{"key": "k2", "token_per_minute": 50},
|
||||
},
|
||||
},
|
||||
},
|
||||
"redis": map[string]interface{}{
|
||||
"service_name": "redis.static",
|
||||
"service_port": 6379,
|
||||
},
|
||||
})
|
||||
return data
|
||||
}()
|
||||
|
||||
func TestParseConfig(t *testing.T) {
|
||||
test.RunGoTest(t, func(t *testing.T) {
|
||||
// 测试全局限流配置解析
|
||||
@@ -291,8 +346,8 @@ func TestOnHttpRequestHeaders(t *testing.T) {
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
// 模拟 Redis 调用响应(允许请求)
|
||||
// 返回 [threshold, current, ttl] 格式
|
||||
resp := test.CreateRedisRespArray([]interface{}{1000, 1, 60})
|
||||
// 返回 [[threshold, current, ttl]] 嵌套格式
|
||||
resp := multiRuleResp([3]int{1000, 1, 60})
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
host.CompleteHttp()
|
||||
@@ -316,7 +371,7 @@ func TestOnHttpRequestHeaders(t *testing.T) {
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
// 模拟 Redis 调用响应(允许请求)
|
||||
resp := test.CreateRedisRespArray([]interface{}{100, 1, 60})
|
||||
resp := multiRuleResp([3]int{100, 1, 60})
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
host.CompleteHttp()
|
||||
@@ -339,7 +394,7 @@ func TestOnHttpRequestHeaders(t *testing.T) {
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
// 模拟 Redis 调用响应(允许请求)
|
||||
resp := test.CreateRedisRespArray([]interface{}{50, 1, 60})
|
||||
resp := multiRuleResp([3]int{50, 1, 60})
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
host.CompleteHttp()
|
||||
@@ -363,7 +418,7 @@ func TestOnHttpRequestHeaders(t *testing.T) {
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
// 模拟 Redis 调用响应(允许请求)
|
||||
resp := test.CreateRedisRespArray([]interface{}{200, 1, 60})
|
||||
resp := multiRuleResp([3]int{200, 1, 60})
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
host.CompleteHttp()
|
||||
@@ -387,7 +442,7 @@ func TestOnHttpRequestHeaders(t *testing.T) {
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
// 模拟 Redis 调用响应(允许请求)
|
||||
resp := test.CreateRedisRespArray([]interface{}{75, 1, 60})
|
||||
resp := multiRuleResp([3]int{75, 1, 60})
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
host.CompleteHttp()
|
||||
@@ -410,8 +465,8 @@ func TestOnHttpRequestHeaders(t *testing.T) {
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
// 模拟 Redis 调用响应(触发限流)
|
||||
// 返回 [threshold, current, ttl] 格式,current > threshold 表示触发限流
|
||||
resp := test.CreateRedisRespArray([]interface{}{1000, 1001, 60})
|
||||
// 返回 [[threshold, current, ttl]] 嵌套格式,current > threshold 表示触发限流
|
||||
resp := multiRuleResp([3]int{1000, 1001, 60})
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
// 检查是否发送了限流响应
|
||||
@@ -440,6 +495,213 @@ func TestOnHttpRequestHeaders(t *testing.T) {
|
||||
// 应该返回 ActionContinue,因为没有匹配到限流规则
|
||||
require.Equal(t, types.ActionContinue, action)
|
||||
})
|
||||
|
||||
// 混合限流:global + rule_item 同时命中,返回 2-key Eval 响应
|
||||
t.Run("hybrid limit both match", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(hybridLimitConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/test"},
|
||||
{":method", "POST"},
|
||||
{"x-api-key", "vip-key"},
|
||||
})
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
resp := multiRuleResp(
|
||||
[3]int{10000, 1, 60},
|
||||
[3]int{100, 1, 60},
|
||||
)
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// 多条 rule_items 同时命中
|
||||
t.Run("multi rule_items all match", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(multiRuleItemsConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/test?apikey=k2"},
|
||||
{":method", "POST"},
|
||||
{"x-api-key", "k1"},
|
||||
})
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
resp := multiRuleResp(
|
||||
[3]int{100, 1, 60},
|
||||
[3]int{50, 1, 60},
|
||||
)
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// 混合限流:global 触发(current > threshold),rule_item 未触发
|
||||
t.Run("hybrid limit global triggered", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(hybridLimitConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/test"},
|
||||
{":method", "POST"},
|
||||
{"x-api-key", "vip-key"},
|
||||
})
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
resp := multiRuleResp(
|
||||
[3]int{10000, 10001, 60},
|
||||
[3]int{100, 1, 60},
|
||||
)
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
localResponse := host.GetLocalResponse()
|
||||
require.NotNil(t, localResponse)
|
||||
require.Equal(t, uint32(429), localResponse.StatusCode)
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// Redis 响应数组长度与规则数不匹配
|
||||
t.Run("redis response length mismatch", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(hybridLimitConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/test"},
|
||||
{":method", "POST"},
|
||||
{"x-api-key", "vip-key"},
|
||||
})
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
resp := multiRuleResp([3]int{10000, 1, 60})
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
require.Nil(t, host.GetLocalResponse())
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// Redis 子数组长度异常(少于 3)
|
||||
t.Run("redis sub-array length mismatch", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(hybridLimitConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/test"},
|
||||
{":method", "POST"},
|
||||
{"x-api-key", "vip-key"},
|
||||
})
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
badResp, err := resp.ArrayValue([]resp.Value{
|
||||
resp.ArrayValue([]resp.Value{
|
||||
resp.IntegerValue(10000),
|
||||
resp.IntegerValue(1),
|
||||
}),
|
||||
}).MarshalRESP()
|
||||
require.NoError(t, err)
|
||||
host.CallOnRedisCall(0, badResp)
|
||||
|
||||
require.Nil(t, host.GetLocalResponse())
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// 混合配置:global 命中,rule_items 不命中(请求头不含 x-api-key)
|
||||
t.Run("hybrid limit rule_items no match", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(hybridLimitConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/test"},
|
||||
{":method", "POST"},
|
||||
})
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
resp := multiRuleResp([3]int{10000, 1, 60})
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
require.Nil(t, host.GetLocalResponse())
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// 多规则同时触发,验证 rejected 报告 global(reset=60 而非 rule_item 的 30)
|
||||
t.Run("hybrid limit both triggered reports global first", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(hybridLimitConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/test"},
|
||||
{":method", "POST"},
|
||||
{"x-api-key", "vip-key"},
|
||||
})
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
resp := multiRuleResp(
|
||||
[3]int{10000, 10001, 60},
|
||||
[3]int{100, 101, 30},
|
||||
)
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
localResponse := host.GetLocalResponse()
|
||||
require.NotNil(t, localResponse)
|
||||
require.Equal(t, uint32(429), localResponse.StatusCode)
|
||||
|
||||
// LocalHttpResponse.Headers 类型为 [][2]string,按切片迭代查找目标头
|
||||
var resetHeader string
|
||||
for _, h := range localResponse.Headers {
|
||||
if strings.EqualFold(h[0], RateLimitResetHeader) && h[1] != "" {
|
||||
resetHeader = h[1]
|
||||
break
|
||||
}
|
||||
}
|
||||
require.Equal(t, "60", resetHeader, "应报告 global 规则的 reset 时间")
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// 多规则未触发:ai-token-ratelimit 不对外暴露 tightest 选择(无 X-RateLimit-* 头),
|
||||
// 此处仅验证不触发拒绝;tightest 行为由 cluster-key-ratelimit 的测试覆盖。
|
||||
t.Run("multi-rule no trigger does not reject", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(hybridLimitConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/test"},
|
||||
{":method", "POST"},
|
||||
{"x-api-key", "vip-key"},
|
||||
})
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
resp := multiRuleResp(
|
||||
[3]int{10000, 9000, 60},
|
||||
[3]int{100, 10, 60},
|
||||
)
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
require.Nil(t, host.GetLocalResponse())
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
@@ -459,7 +721,7 @@ func TestOnHttpStreamingBody(t *testing.T) {
|
||||
})
|
||||
|
||||
// 模拟 Redis 调用响应
|
||||
resp := test.CreateRedisRespArray([]interface{}{1000, 1, 60})
|
||||
resp := multiRuleResp([3]int{1000, 1, 60})
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
// 处理流式响应体
|
||||
@@ -499,7 +761,7 @@ func TestOnHttpStreamingBody(t *testing.T) {
|
||||
})
|
||||
|
||||
// 模拟 Redis 调用响应
|
||||
resp := test.CreateRedisRespArray([]interface{}{1000, 1, 60})
|
||||
resp := multiRuleResp([3]int{1000, 1, 60})
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
// 处理流式响应体
|
||||
@@ -514,6 +776,72 @@ func TestOnHttpStreamingBody(t *testing.T) {
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// 多规则下响应阶段应触发 2-key INCRBY。
|
||||
// Strengthened: 用 GetRedisCalloutAttributes() 断言响应阶段确实发起了一次 Eval。
|
||||
// Callout 计数模式:0 → 1 (请求阶段) → 0 (CallOnRedisCall 消费) → 1 (响应阶段)
|
||||
// 如果 response 阶段被改回 1-key 或不调用 Redis,最终计数会停在 0 而非 1。
|
||||
t.Run("streaming body multi-rule incrby", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(hybridLimitConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
require.Equal(t, 0, len(host.GetRedisCalloutAttributes()),
|
||||
"请求前应无 Redis 调用")
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/test"},
|
||||
{":method", "POST"},
|
||||
{"x-api-key", "vip-key"},
|
||||
})
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
require.Equal(t, 1, len(host.GetRedisCalloutAttributes()),
|
||||
"请求阶段应发起 1 次多键 Eval")
|
||||
|
||||
resp := multiRuleResp(
|
||||
[3]int{10000, 1, 60},
|
||||
[3]int{100, 1, 60},
|
||||
)
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
body := []byte(`{"choices":[{"message":{"content":"hi"}}],"usage":{"prompt_tokens":10,"completion_tokens":15,"total_tokens":25}}`)
|
||||
// 注意:ai-token 的 onHttpStreamingBody 注册为 ProcessStreamingResponseBody,
|
||||
// 对应测试方法 CallOnHttpStreamingResponseBody(不是 CallOnHttpStreamingRequestBody)。
|
||||
// 这是因为 ai-token 处理的是 LLM 的流式响应体(usage 字段在响应里)。
|
||||
streamAction := host.CallOnHttpStreamingResponseBody(body, true)
|
||||
require.Equal(t, types.ActionContinue, streamAction)
|
||||
// 响应阶段应再发起 1 次多键 INCRBY(callout 数 0 → 1)。
|
||||
// 如果 response 阶段被改回 1-key 或不调用 Redis,此断言会失败。
|
||||
require.Equal(t, 1, len(host.GetRedisCalloutAttributes()),
|
||||
"响应阶段应再发起 1 次多键 INCRBY(callout 数 0 → 1)")
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// 流式响应体处理时无前置 context(覆盖 line 216-225, 227-234):
|
||||
// 使用 headerLimitConfig(仅 x-api-key 规则),ensureContextInitialized 自动触发的
|
||||
// onHttpRequestHeaders 因为缺少 x-api-key 头而 matched 为空,返回 ActionContinue,
|
||||
// 此时 token context / LimitRedisContext 均未被设置。后续 streaming body 应通过
|
||||
// GetContext 判空分支安全返回 data,不 panic、不调用 Redis。
|
||||
t.Run("streaming body without prior context", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(headerLimitConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
// 直接处理流式响应体(缺少 x-api-key,matched 为空,SetContext 未调用)
|
||||
body := []byte(`{"choices":[{"message":{"content":"hi"}}],"usage":{"prompt_tokens":10,"completion_tokens":15,"total_tokens":25}}`)
|
||||
streamAction := host.CallOnHttpStreamingResponseBody(body, true)
|
||||
|
||||
// 应安全返回 ActionContinue
|
||||
require.Equal(t, types.ActionContinue, streamAction)
|
||||
|
||||
// 不应触发任何 Redis callout(前置 context 缺失时跳过响应阶段)
|
||||
require.Equal(t, 0, len(host.GetRedisCalloutAttributes()),
|
||||
"missing token context should skip response-phase INCRBY")
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
@@ -537,7 +865,7 @@ func TestCompleteFlow(t *testing.T) {
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
// 2. 模拟 Redis 调用响应
|
||||
resp := test.CreateRedisRespArray([]interface{}{100, 1, 60})
|
||||
resp := multiRuleResp([3]int{100, 1, 60})
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
// 3. 处理流式响应体
|
||||
@@ -555,3 +883,24 @@ func TestCompleteFlow(t *testing.T) {
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
// multiRuleResp 构建多规则嵌套 Redis 响应(RESP wire format)。
|
||||
// test.CreateRedisRespArray 不支持嵌套数组,因此直接用 resp.Writer 构造。
|
||||
// 每个 [3]int 元组为 {threshold, current, ttl}。
|
||||
// 注意:不能使用 resp.Value.Bytes()——它返回的是显示用字符串,不是 RESP wire 格式。
|
||||
// 必须用 resp.Writer.WriteArray 或 Value.MarshalRESP() 输出真正的 RESP 字节流。
|
||||
func multiRuleResp(items ...[3]int) []byte {
|
||||
values := make([]resp.Value, len(items))
|
||||
for i, it := range items {
|
||||
values[i] = resp.ArrayValue([]resp.Value{
|
||||
resp.IntegerValue(it[0]),
|
||||
resp.IntegerValue(it[1]),
|
||||
resp.IntegerValue(it[2]),
|
||||
})
|
||||
}
|
||||
b, err := resp.ArrayValue(values).MarshalRESP()
|
||||
if err != nil {
|
||||
panic(fmt.Sprintf("failed to marshal multiRuleResp: %v", err))
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
@@ -4,14 +4,25 @@ keywords: [higress, rate-limit]
|
||||
description: 基于 Key 集群限流插件配置参考
|
||||
---
|
||||
|
||||
> ⚠️ **行为变更提示(无版本号变化)**
|
||||
>
|
||||
> 自本次更新起,`rule_items` 的匹配语义从 **first-match-wins**(命中第一条即返回)改为 **all-match OR 叠加**(所有命中规则都评估,任一触发即拒绝)。同时解除 `global_threshold` 与 `rule_items` 的互斥约束,支持混合配置。
|
||||
>
|
||||
> - 老配置(单条 `rule_items` 或仅 `global_threshold`):行为不变
|
||||
> - 老配置(多条 `rule_items` 期望短路匹配):**行为会变** —— 所有命中规则都会评估
|
||||
> - Redis key 格式新增 `{rule_name}` hash tag 以兼容 Redis Cluster,老计数器数据不兼容
|
||||
>
|
||||
> 详见下方"功能说明"与"配置示例"。
|
||||
|
||||
## 功能说明
|
||||
|
||||
`cluster-key-rate-limit` 插件基于 Redis 实现**集群级限流**,适用于需要跨多个 Higress Gateway 实例进行**全局一致速率限制**的场景。
|
||||
|
||||
支持两种限流模式:
|
||||
支持三种限流模式:
|
||||
|
||||
- **规则级全局限流**:基于相同的 `rule_name` 和 `global_threshold` 配置,对自定义规则组设置全局限流阈值
|
||||
- **Key 级动态限流**:根据请求中的动态 Key(如 URL 参数、请求头、客户端 IP、Consumer 名称或 Cookie 字段)进行分组限流
|
||||
- **混合限流**:同时配置 `global_threshold`(全局兜底)和 `rule_items`(按维度细分),所有命中的规则**叠加生效**,任一触发即拒绝。
|
||||
|
||||
## 运行属性
|
||||
|
||||
@@ -23,8 +34,8 @@ description: 基于 Key 集群限流插件配置参考
|
||||
| 配置项 | 类型 | 必填 | 默认值 | 说明 |
|
||||
| ----------------------- | ------ | ---- | ------ |-----------------------------------------------------------------------------|
|
||||
| rule_name | string | 是 | - | 限流规则名称,根据限流规则名称 + 限流类型 + 限流 key 名称 + 限流 key 对应的实际值来拼装 redis key |
|
||||
| global_threshold | Object | 否,`global_threshold` 或 `rule_items` 选填一项 | - | 对整个自定义规则组进行限流 |
|
||||
| rule_items | array of object | 否,`global_threshold` 或 `rule_items` 选填一项 | - | 限流规则项,按照 rule_items 下的排列顺序,匹配第一个 rule_item 后命中限流规则,后续规则将被忽略 |
|
||||
| global_threshold | Object | 否,至少一项;可同时配置 | - | 对整个自定义规则组进行限流 |
|
||||
| rule_items | array of object | 否,至少一项;可同时配置 | - | 限流规则项,最多支持 **10 条**。所有满足匹配条件的 `rule_item` 都会参与限流,规则之间是"或"关系,任一触发即拒绝;规则的执行顺序不影响最终结果。详见下方"配置说明"展开。 |
|
||||
| show_limit_quota_header | bool | 否 | false | 响应头中是否显示 `X-RateLimit-Limit`(限制的总请求数)和 `X-RateLimit-Remaining`(剩余还可以发送的请求数) |
|
||||
| rejected_code | int | 否 | 429 | 请求被限流时,返回的 HTTP 状态码 |
|
||||
| rejected_msg | string | 否 | Too many requests | 请求被限流时,返回的响应体 |
|
||||
@@ -54,6 +65,18 @@ description: 基于 Key 集群限流插件配置参考
|
||||
| limit_by_per_ip | string | 否,`limit_by_*` 中选填一项 | - | 按规则匹配特定 IP,并对每个 IP 分别计算限流,配置获取限流键值的来源 IP 参数名称,从请求头获取,以 `from-header-对应的header名`,示例:`from-header-x-forwarded-for`,直接获取对端 socket ip,配置为 `from-remote-addr` |
|
||||
| limit_keys | array of object | 是 | - | 配置匹配键值后的限流次数 |
|
||||
|
||||
#### `rule_items` 多规则匹配语义
|
||||
|
||||
`rule_items` 是一个数组,**所有满足匹配条件**的 `rule_item` 都会被评估,规则之间是"或"关系,**任一触发即拒绝**。规则的执行顺序不影响最终结果。
|
||||
|
||||
> `rule_items` 数组最多支持 **10 条**规则。每条 `rule_item` 会按命中的 `limit_keys` 产生独立的 Redis 计数器。
|
||||
|
||||
##### 多规则场景下的 X-RateLimit-* 头
|
||||
|
||||
当多条规则同时未触发、配置 `show_limit_quota_header: true` 时:
|
||||
- `X-RateLimit-Limit` / `X-RateLimit-Remaining`:取剩余比例最小(最紧约束)的命中规则
|
||||
- `X-RateLimit-Reset`(触发限流时返回):取第一条触发的规则(按 `rule_items` 数组顺序,全局优先)
|
||||
|
||||
`limit_keys` 中每一项的配置字段说明。
|
||||
|
||||
| 配置项 | 类型 | 必填 | 默认值 | 说明 |
|
||||
|
||||
@@ -5,15 +5,26 @@ description: Configuration reference for the Key-based cluster rate limiting plu
|
||||
|
||||
---
|
||||
|
||||
> ⚠️ **Behavior Change Notice (no version bump)**
|
||||
>
|
||||
> As of this update, `rule_items` matching semantics changed from **first-match-wins** (returns on first hit) to **all-match OR overlay** (all matched rules are evaluated, any trigger rejects). The mutual exclusion between `global_threshold` and `rule_items` is also removed to support hybrid configuration.
|
||||
>
|
||||
> - Existing config (single `rule_items` or only `global_threshold`): behavior unchanged
|
||||
> - Existing config (multiple `rule_items` expecting short-circuit match): **behavior will change** — all matched rules are evaluated
|
||||
> - Redis key format adds `{rule_name}` hash tag for Redis Cluster compatibility; old counter data is incompatible
|
||||
>
|
||||
> See "Function Description" and "Configuration Examples" below for details.
|
||||
|
||||
## Function Description
|
||||
|
||||
The `cluster-key-rate-limit` plugin implements **cluster-level rate limiting** based on Redis, suitable for scenarios
|
||||
requiring **globally consistent rate limiting across multiple Higress Gateway instances**.
|
||||
|
||||
It supports two rate limiting modes:
|
||||
It supports three rate limiting modes:
|
||||
|
||||
- **Rule-Level Global Rate Limiting**: Applies a unified rate limit threshold to custom rule groups based on identical `rule_name` and `global_threshold` configurations.
|
||||
- **Key-Level Dynamic Rate Limiting**: Groups and limits requests by dynamic keys extracted from requests, such as URL parameters, request headers, client IPs, consumer names, or cookie fields.
|
||||
- **Hybrid rate limiting**: Configure `global_threshold` (global fallback) and `rule_items` (per-dimension) simultaneously. All matched rules take effect together; any trigger rejects the request.
|
||||
|
||||
## Operational Attributes
|
||||
|
||||
@@ -25,8 +36,8 @@ It supports two rate limiting modes:
|
||||
| Configuration Item | Type | Required | Default Value | Description |
|
||||
|--------------------------|---------------|-------------------------------------------|---------------------|----------------------------------------------------------------------------|
|
||||
| rule_name | string | Yes | - | Name of the rate limiting rule. Used to construct the Redis key in the format: `rule_name:rate_limit_type:key_name:key_value`. |
|
||||
| global_threshold | Object | No (choose either `global_threshold` or `rule_items`) | - | Apply rate limiting to the entire custom rule group.|
|
||||
| rule_items | array of object | No (choose either `global_threshold` or `rule_items`) | - | Rate limiting rule items. Rules are matched in the order of the array; once the first matching rule is hit, subsequent rules are ignored. |
|
||||
| global_threshold | Object | No, at least one of them is required; can be configured simultaneously | - | Apply rate limiting to the entire custom rule group.|
|
||||
| rule_items | array of object | No, at least one of them is required; can be configured simultaneously | - | Rate limiting rule items. Supports up to **10** rules. All matched `rule_item`s are evaluated and combined with an OR relationship; any trigger rejects the request. The execution order of rules does not affect the final result. See the expanded `rule_items` notes below for details. |
|
||||
| show_limit_quota_header | bool | No | false | Whether to display `X-RateLimit-Limit` (total allowed requests) and `X-RateLimit-Remaining` (remaining allowed requests) in the response header. |
|
||||
| rejected_code | int | No | 429 | HTTP status code returned when a request is rate-limited. |
|
||||
| rejected_msg | string | No | Too many requests | Response body returned when a request is rate-limited. |
|
||||
@@ -56,6 +67,18 @@ It supports two rate limiting modes:
|
||||
| limit_by_per_ip | string | No (choose one of `limit_by_*` fields) | - | Matches specific IPs by rule and calculates rate limits for each IP. The IP can be extracted from a request header (formatted as `from-header-<header_name>`, e.g., `from-header-x-forwarded-for`) or directly from the peer socket IP (configured as `from-remote-addr`). |
|
||||
| limit_keys | array of object | Yes | - | Configures the rate limits for matched key values. |
|
||||
|
||||
#### `rule_items` Multi-Rule Matching Semantics
|
||||
|
||||
`rule_items` is an array. **All** `rule_item`s whose match conditions are satisfied are evaluated, and the rules are combined with an OR relationship: **any trigger rejects the request**. The execution order of rules does not affect the final result.
|
||||
|
||||
> The `rule_items` array supports up to **10** rules. Each `rule_item` produces independent Redis counters for each matched `limit_keys`.
|
||||
|
||||
##### X-RateLimit-* headers in multi-rule scenarios
|
||||
|
||||
When multiple rules are matched and none trigger (with `show_limit_quota_header: true`):
|
||||
- `X-RateLimit-Limit` / `X-RateLimit-Remaining`: from the matched rule with the smallest remaining ratio (tightest constraint)
|
||||
- `X-RateLimit-Reset` (returned when triggered): from the first triggered rule (in `rule_items` array order, global first)
|
||||
|
||||
### Configuration Fields for `limit_keys`
|
||||
|
||||
| Configuration Item | Type | Required | Default Value | Description |
|
||||
|
||||
@@ -47,6 +47,9 @@ const (
|
||||
SecondsPerMinute = 60 * Second
|
||||
SecondsPerHour = 60 * SecondsPerMinute
|
||||
SecondsPerDay = 24 * SecondsPerHour
|
||||
|
||||
// MaxRuleItems 限制 rule_items 数组最大长度
|
||||
MaxRuleItems = 10
|
||||
)
|
||||
|
||||
var timeWindows = map[string]int64{
|
||||
@@ -169,8 +172,6 @@ func initLimitRule(json gjson.Result, config *ClusterKeyRateLimitConfig) error {
|
||||
hasRule := ruleItemsResult.Exists()
|
||||
if !hasGlobal && !hasRule {
|
||||
return errors.New("at least one of 'global_threshold' or 'rule_items' must be set")
|
||||
} else if hasGlobal && hasRule {
|
||||
return errors.New("'global_threshold' and 'rule_items' cannot be set at the same time")
|
||||
}
|
||||
|
||||
// 处理全局限流配置
|
||||
@@ -180,14 +181,19 @@ func initLimitRule(json gjson.Result, config *ClusterKeyRateLimitConfig) error {
|
||||
return fmt.Errorf("failed to parse global_threshold: %w", err)
|
||||
}
|
||||
config.GlobalThreshold = threshold
|
||||
return nil
|
||||
}
|
||||
|
||||
// 处理条件限流规则
|
||||
if !hasRule {
|
||||
return nil
|
||||
}
|
||||
items := ruleItemsResult.Array()
|
||||
if len(items) == 0 {
|
||||
return errors.New("config rule_items cannot be empty")
|
||||
}
|
||||
if len(items) > MaxRuleItems {
|
||||
return fmt.Errorf("rule_items length %d exceeds maximum %d", len(items), MaxRuleItems)
|
||||
}
|
||||
|
||||
var ruleItems []LimitRuleItem
|
||||
// 用于记录已出现的LimitType和Key的组合
|
||||
@@ -286,7 +292,7 @@ func parseLimitRuleItem(item gjson.Result) (*LimitRuleItem, error) {
|
||||
}
|
||||
|
||||
if limitType == "" {
|
||||
return nil, errors.New("only one of 'limit_by_header' and 'limit_by_param' and 'limit_by_consumer' and 'limit_by_cookie' and 'limit_by_per_header' and 'limit_by_per_param' and 'limit_by_per_consumer' and 'limit_by_per_cookie' and 'limit_by_per_ip' can be set")
|
||||
return nil, errors.New("at least one of 'limit_by_header', 'limit_by_param', 'limit_by_consumer', 'limit_by_cookie', 'limit_by_per_header', 'limit_by_per_param', 'limit_by_per_consumer', 'limit_by_per_cookie', 'limit_by_per_ip' must be set")
|
||||
}
|
||||
ruleItem.LimitType = limitType
|
||||
|
||||
|
||||
@@ -2,6 +2,8 @@ package config
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
@@ -166,13 +168,41 @@ func TestParseClusterKeyRateLimitConfig(t *testing.T) {
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "Conflict_GlobalThresholdAndRuleItems",
|
||||
name: "Both_ShouldSucceed",
|
||||
json: `{
|
||||
"rule_name": "test-conflict",
|
||||
"rule_name": "test-both",
|
||||
"global_threshold": {"query_per_second": 100},
|
||||
"rule_items": [{"limit_by_header": "x-test"}]
|
||||
"rule_items": [
|
||||
{
|
||||
"limit_by_header": "x-test",
|
||||
"limit_keys": [{"key": "k1", "query_per_minute": 10}]
|
||||
}
|
||||
]
|
||||
}`,
|
||||
expectedErr: errors.New("'global_threshold' and 'rule_items' cannot be set at the same time"),
|
||||
expected: ClusterKeyRateLimitConfig{
|
||||
RuleName: "test-both",
|
||||
GlobalThreshold: &GlobalThreshold{
|
||||
Count: 100,
|
||||
TimeWindow: Second,
|
||||
},
|
||||
RuleItems: []LimitRuleItem{
|
||||
{
|
||||
LimitType: LimitByHeaderType,
|
||||
Key: "x-test",
|
||||
ConfigItems: []LimitConfigItem{
|
||||
{
|
||||
ConfigType: ExactType,
|
||||
Key: "k1",
|
||||
Count: 10,
|
||||
TimeWindow: SecondsPerMinute,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
ShowLimitQuotaHeader: false,
|
||||
RejectedCode: DefaultRejectedCode,
|
||||
RejectedMsg: DefaultRejectedMsg,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "Missing_GlobalThresholdAndRuleItems",
|
||||
@@ -217,6 +247,16 @@ func TestParseClusterKeyRateLimitConfig(t *testing.T) {
|
||||
RejectedMsg: DefaultRejectedMsg,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "LimitByType_AtLeastOne_Error",
|
||||
json: `{
|
||||
"rule_name": "no-limit-by",
|
||||
"rule_items": [
|
||||
{"limit_keys": [{"key": "k", "query_per_second": 1}]}
|
||||
]
|
||||
}`,
|
||||
expectedErr: errors.New("failed to parse rule_item in rule_items: at least one of 'limit_by_header', 'limit_by_param', 'limit_by_consumer', 'limit_by_cookie', 'limit_by_per_header', 'limit_by_per_param', 'limit_by_per_consumer', 'limit_by_per_cookie', 'limit_by_per_ip' must be set"),
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
@@ -234,3 +274,32 @@ func TestParseClusterKeyRateLimitConfig(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseClusterKeyRateLimitConfig_RuleItemsLimit(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
itemCnt int
|
||||
expectErr string
|
||||
}{
|
||||
{"Exactly10", 10, ""},
|
||||
{"Over10", 11, "rule_items length 11 exceeds maximum 10"},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
items := make([]string, 0, tt.itemCnt)
|
||||
for i := 0; i < tt.itemCnt; i++ {
|
||||
items = append(items,
|
||||
fmt.Sprintf(`{"limit_by_header":"h%d","limit_keys":[{"key":"k","query_per_second":1}]}`, i+1))
|
||||
}
|
||||
jsonStr := fmt.Sprintf(`{"rule_name":"t","rule_items":[%s]}`, strings.Join(items, ","))
|
||||
var config ClusterKeyRateLimitConfig
|
||||
err := ParseClusterKeyRateLimitConfig(gjson.Parse(jsonStr), &config)
|
||||
if tt.expectErr == "" {
|
||||
assert.NoError(t, err)
|
||||
assert.Len(t, config.RuleItems, tt.itemCnt)
|
||||
} else {
|
||||
assert.EqualError(t, err, tt.expectErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -16,6 +16,7 @@ package main
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
"net"
|
||||
"net/url"
|
||||
"strconv"
|
||||
@@ -46,30 +47,37 @@ func init() {
|
||||
const (
|
||||
// RedisKeyPrefix 集群限流插件在 Redis 中 key 的统一前缀
|
||||
RedisKeyPrefix = "higress-cluster-key-rate-limit"
|
||||
// 使用 {rule_name} hash tag 让多规则多键操作在 Redis Cluster 下落到同一 slot
|
||||
// ClusterGlobalRateLimitFormat 全局限流模式 redis key 为 RedisKeyPrefix:限流规则名称:global_threshold:时间窗口
|
||||
ClusterGlobalRateLimitFormat = RedisKeyPrefix + ":%s:global_threshold:%d"
|
||||
ClusterGlobalRateLimitFormat = RedisKeyPrefix + ":{%s}:global_threshold:%d"
|
||||
// ClusterRateLimitFormat 规则限流模式 redis key 为 RedisKeyPrefix:限流规则名称:限流类型:时间窗口:限流key名称:限流key对应的实际值
|
||||
ClusterRateLimitFormat = RedisKeyPrefix + ":%s:%s:%d:%s:%s"
|
||||
FixedWindowScript = `
|
||||
local key = KEYS[1]
|
||||
local threshold = tonumber(ARGV[1])
|
||||
local window = tonumber(ARGV[2])
|
||||
|
||||
local current = tonumber(redis.call('get', key) or "0")
|
||||
|
||||
-- 只有超过阈值时才停止累加,达到阈值时仍允许(此时是最后一次允许)
|
||||
if current > threshold then
|
||||
return {threshold, current, redis.call('ttl', key)}
|
||||
ClusterRateLimitFormat = RedisKeyPrefix + ":{%s}:%s:%d:%s:%s"
|
||||
// MultiKeyFixedWindowScript 多规则请求阶段 check + incr 合一脚本(cluster-key-rate-limit 使用)
|
||||
// KEYS = [key1, ..., keyN]
|
||||
// ARGV = [threshold1, window1, ..., thresholdN, windowN]
|
||||
// 返回嵌套数组 {{threshold_i, current_i, ttl_i}, ...}(每个 key 独立判断是否 incr)
|
||||
MultiKeyFixedWindowScript = `
|
||||
local results = {}
|
||||
for i = 1, #KEYS do
|
||||
local threshold = tonumber(ARGV[2*i - 1])
|
||||
local window = tonumber(ARGV[2*i])
|
||||
local current = tonumber(redis.call('get', KEYS[i]) or "0")
|
||||
local ttl = redis.call('ttl', KEYS[i])
|
||||
if ttl < 0 then ttl = window end
|
||||
|
||||
if current > threshold then
|
||||
-- 已超阈值,不再 incr
|
||||
table.insert(results, {threshold, current, ttl})
|
||||
else
|
||||
-- 未超阈值,原子 incr
|
||||
current = redis.call('incr', KEYS[i])
|
||||
if current == 1 then
|
||||
redis.call('expire', KEYS[i], window)
|
||||
end
|
||||
table.insert(results, {threshold, current, redis.call('ttl', KEYS[i])})
|
||||
end
|
||||
end
|
||||
|
||||
-- 计数未超过阈值,执行累加
|
||||
current = redis.call('incr', key)
|
||||
-- 第一次累加时设置过期时间
|
||||
if current == 1 then
|
||||
redis.call('expire', key, window)
|
||||
end
|
||||
|
||||
return {threshold, current, redis.call('ttl', key)}
|
||||
return results
|
||||
`
|
||||
|
||||
LimitContextKey = "LimitContext" // 限流上下文信息
|
||||
@@ -87,6 +95,13 @@ type LimitContext struct {
|
||||
reset int
|
||||
}
|
||||
|
||||
// MatchedRule 表示请求阶段命中的单条限流规则(global 或 rule_item)
|
||||
type MatchedRule struct {
|
||||
key string // 完整 Redis key
|
||||
count int64 // 时间窗口内的限额
|
||||
window int64 // 时间窗口大小(秒)
|
||||
}
|
||||
|
||||
func parseConfig(json gjson.Result, cfg *config.ClusterKeyRateLimitConfig) error {
|
||||
err := config.InitRedisClusterClient(json, cfg)
|
||||
if err != nil {
|
||||
@@ -101,51 +116,82 @@ func parseConfig(json gjson.Result, cfg *config.ClusterKeyRateLimitConfig) error
|
||||
|
||||
func onHttpRequestHeaders(ctx wrapper.HttpContext, cfg config.ClusterKeyRateLimitConfig) types.Action {
|
||||
ctx.DisableReroute()
|
||||
limitKey, count, timeWindow := "", int64(0), int64(0)
|
||||
|
||||
if cfg.GlobalThreshold != nil {
|
||||
// 全局限流模式
|
||||
limitKey = fmt.Sprintf(ClusterGlobalRateLimitFormat, cfg.RuleName, cfg.GlobalThreshold.TimeWindow)
|
||||
count = cfg.GlobalThreshold.Count
|
||||
timeWindow = cfg.GlobalThreshold.TimeWindow
|
||||
} else {
|
||||
// 规则限流模式
|
||||
val, ruleItem, configItem := checkRequestAgainstLimitRule(ctx, cfg.RuleItems)
|
||||
if ruleItem == nil || configItem == nil {
|
||||
// 没有匹配到限流规则直接返回
|
||||
return types.ActionContinue
|
||||
}
|
||||
|
||||
limitKey = fmt.Sprintf(ClusterRateLimitFormat, cfg.RuleName, ruleItem.LimitType, configItem.TimeWindow, ruleItem.Key, val)
|
||||
count = configItem.Count
|
||||
timeWindow = configItem.TimeWindow
|
||||
matched := collectMatchedRules(ctx, cfg)
|
||||
if len(matched) == 0 {
|
||||
// 无任何规则命中:直接放行,不发起 Redis 调用
|
||||
log.Debugf("cluster-key-rate-limit: no rule matched, path=%s host=%s", ctx.Path(), ctx.Host())
|
||||
return types.ActionContinue
|
||||
}
|
||||
|
||||
// 执行限流逻辑
|
||||
keys := []interface{}{limitKey}
|
||||
args := []interface{}{count, timeWindow}
|
||||
err := cfg.RedisClient.Eval(FixedWindowScript, 1, keys, args, func(response resp.Value) {
|
||||
resultArray := response.Array()
|
||||
if len(resultArray) != 3 {
|
||||
log.Errorf("redis response parse error, response: %v", response)
|
||||
proxywasm.ResumeHttpRequest()
|
||||
log.Debugf("cluster-key-rate-limit: request phase matched %d rule(s), path=%s host=%s",
|
||||
len(matched), ctx.Path(), ctx.Host())
|
||||
for i, m := range matched {
|
||||
log.Debugf("cluster-key-rate-limit: rule[%d] key=%s threshold=%d window=%ds",
|
||||
i, m.key, m.count, m.window)
|
||||
}
|
||||
|
||||
n := len(matched)
|
||||
keys := make([]interface{}, n)
|
||||
args := make([]interface{}, 0, n*2)
|
||||
for i, m := range matched {
|
||||
keys[i] = m.key
|
||||
args = append(args, m.count, m.window)
|
||||
}
|
||||
|
||||
err := cfg.RedisClient.Eval(MultiKeyFixedWindowScript, n, keys, args, func(response resp.Value) {
|
||||
arr := response.Array()
|
||||
if len(arr) != n {
|
||||
log.Errorf("redis response length mismatch: got %d, want %d", len(arr), n)
|
||||
_ = proxywasm.ResumeHttpRequest()
|
||||
return
|
||||
}
|
||||
|
||||
// 获取限流结果
|
||||
threshold, current, ttl := resultArray[0].Integer(), resultArray[1].Integer(), resultArray[2].Integer()
|
||||
context := LimitContext{
|
||||
count: threshold,
|
||||
remaining: threshold - current,
|
||||
reset: ttl,
|
||||
}
|
||||
if current > threshold {
|
||||
// 触发限流
|
||||
rejected(cfg, context)
|
||||
} else {
|
||||
ctx.SetContext(LimitContextKey, context)
|
||||
proxywasm.ResumeHttpRequest()
|
||||
// 使用 math.MaxFloat64 初始化,避免在 arr[0] 校验前预先读取
|
||||
tightestIdx := 0
|
||||
tightestRatio := math.MaxFloat64
|
||||
|
||||
for i, sub := range arr {
|
||||
a := sub.Array()
|
||||
if len(a) != 3 {
|
||||
log.Errorf("redis sub-array length mismatch: got %d, want 3", len(a))
|
||||
_ = proxywasm.ResumeHttpRequest()
|
||||
return
|
||||
}
|
||||
threshold, current, ttl := a[0].Integer(), a[1].Integer(), a[2].Integer()
|
||||
log.Debugf("cluster-key-rate-limit: eval rule[%d] key=%s threshold=%d current=%d ttl=%ds",
|
||||
i, matched[i].key, threshold, current, ttl)
|
||||
|
||||
if current > threshold {
|
||||
// 命中触发的第一条规则(按 collectMatchedRules 顺序,global 优先)
|
||||
log.Debugf("cluster-key-rate-limit: rule[%d] key=%s triggered (current=%d > threshold=%d), rejecting with code %d",
|
||||
i, matched[i].key, current, threshold, cfg.RejectedCode)
|
||||
rejected(cfg, LimitContext{
|
||||
count: threshold,
|
||||
remaining: threshold - current,
|
||||
reset: ttl,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
if ratio := float64(threshold-current) / float64(threshold); ratio < tightestRatio {
|
||||
tightestIdx = i
|
||||
tightestRatio = ratio
|
||||
}
|
||||
}
|
||||
|
||||
// 未触发:写入 tightest 规则到 LimitContext,供 onHttpResponseHeaders 读取 X-RateLimit-* 头
|
||||
tightSub := arr[tightestIdx].Array()
|
||||
tightThreshold, tightCurrent, tightTtl := tightSub[0].Integer(), tightSub[1].Integer(), tightSub[2].Integer()
|
||||
log.Debugf("cluster-key-rate-limit: all %d rule(s) within threshold, tightest rule[%d] key=%s (count=%d remaining=%d reset=%ds)",
|
||||
n, tightestIdx, matched[tightestIdx].key, tightThreshold, tightThreshold-tightCurrent, tightTtl)
|
||||
ctx.SetContext(LimitContextKey, LimitContext{
|
||||
count: tightThreshold,
|
||||
remaining: tightThreshold - tightCurrent,
|
||||
reset: tightTtl,
|
||||
})
|
||||
|
||||
_ = proxywasm.ResumeHttpRequest()
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
@@ -158,25 +204,43 @@ func onHttpRequestHeaders(ctx wrapper.HttpContext, cfg config.ClusterKeyRateLimi
|
||||
func onHttpResponseHeaders(ctx wrapper.HttpContext, config config.ClusterKeyRateLimitConfig) types.Action {
|
||||
limitContext, ok := ctx.GetContext(LimitContextKey).(LimitContext)
|
||||
if !ok {
|
||||
log.Debugf("cluster-key-rate-limit: response phase reached with no LimitContext, skipping X-RateLimit-* headers")
|
||||
return types.ActionContinue
|
||||
}
|
||||
if config.ShowLimitQuotaHeader {
|
||||
_ = proxywasm.ReplaceHttpResponseHeader(RateLimitLimitHeader, strconv.Itoa(limitContext.count))
|
||||
_ = proxywasm.ReplaceHttpResponseHeader(RateLimitRemainingHeader, strconv.Itoa(limitContext.remaining))
|
||||
log.Debugf("cluster-key-rate-limit: response phase wrote X-RateLimit-Limit=%d X-RateLimit-Remaining=%d",
|
||||
limitContext.count, limitContext.remaining)
|
||||
}
|
||||
return types.ActionContinue
|
||||
}
|
||||
|
||||
func checkRequestAgainstLimitRule(ctx wrapper.HttpContext, ruleItems []config.LimitRuleItem) (string, *config.LimitRuleItem, *config.LimitConfigItem) {
|
||||
if len(ruleItems) > 0 {
|
||||
for _, rule := range ruleItems {
|
||||
val, ruleItem, configItem := hitRateRuleItem(ctx, rule)
|
||||
if ruleItem != nil && configItem != nil {
|
||||
return val, ruleItem, configItem
|
||||
}
|
||||
// collectMatchedRules 遍历 global_threshold 和 rule_items,返回所有命中规则。
|
||||
// 顺序:global_threshold(如有)→ rule_items 中所有命中项(按数组顺序追加)。
|
||||
func collectMatchedRules(ctx wrapper.HttpContext, cfg config.ClusterKeyRateLimitConfig) []MatchedRule {
|
||||
var matched []MatchedRule
|
||||
|
||||
if cfg.GlobalThreshold != nil {
|
||||
matched = append(matched, MatchedRule{
|
||||
key: fmt.Sprintf(ClusterGlobalRateLimitFormat, cfg.RuleName, cfg.GlobalThreshold.TimeWindow),
|
||||
count: cfg.GlobalThreshold.Count,
|
||||
window: cfg.GlobalThreshold.TimeWindow,
|
||||
})
|
||||
}
|
||||
|
||||
for _, ruleItem := range cfg.RuleItems {
|
||||
val, hitRule, hitItem := hitRateRuleItem(ctx, ruleItem)
|
||||
if hitRule != nil && hitItem != nil {
|
||||
matched = append(matched, MatchedRule{
|
||||
key: fmt.Sprintf(ClusterRateLimitFormat, cfg.RuleName, hitRule.LimitType, hitItem.TimeWindow, hitRule.Key, val),
|
||||
count: hitItem.Count,
|
||||
window: hitItem.TimeWindow,
|
||||
})
|
||||
}
|
||||
}
|
||||
return "", nil, nil
|
||||
|
||||
return matched
|
||||
}
|
||||
|
||||
func hitRateRuleItem(ctx wrapper.HttpContext, rule config.LimitRuleItem) (string, *config.LimitRuleItem, *config.LimitConfigItem) {
|
||||
|
||||
@@ -16,6 +16,8 @@ package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"cluster-key-rate-limit/config"
|
||||
@@ -23,6 +25,7 @@ import (
|
||||
"github.com/higress-group/proxy-wasm-go-sdk/proxywasm/types"
|
||||
"github.com/higress-group/wasm-go/pkg/test"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/resp"
|
||||
)
|
||||
|
||||
// 测试配置:全局限流配置
|
||||
@@ -222,6 +225,56 @@ var regexpLimitConfig = func() json.RawMessage {
|
||||
return data
|
||||
}()
|
||||
|
||||
// 测试配置:混合限流(global + rule_items 同时生效)
|
||||
var hybridLimitConfig = func() json.RawMessage {
|
||||
data, _ := json.Marshal(map[string]interface{}{
|
||||
"rule_name": "cluster-hybrid",
|
||||
"global_threshold": map[string]interface{}{
|
||||
"query_per_minute": 10000,
|
||||
},
|
||||
"rule_items": []map[string]interface{}{
|
||||
{
|
||||
"limit_by_header": "x-api-key",
|
||||
"limit_keys": []map[string]interface{}{
|
||||
{"key": "vip-key", "query_per_minute": 100},
|
||||
},
|
||||
},
|
||||
},
|
||||
"redis": map[string]interface{}{
|
||||
"service_name": "redis.static",
|
||||
"service_port": 6379,
|
||||
},
|
||||
"show_limit_quota_header": true,
|
||||
})
|
||||
return data
|
||||
}()
|
||||
|
||||
// 测试配置:多条 rule_items 同时命中
|
||||
var multiRuleItemsConfig = func() json.RawMessage {
|
||||
data, _ := json.Marshal(map[string]interface{}{
|
||||
"rule_name": "cluster-multi-items",
|
||||
"rule_items": []map[string]interface{}{
|
||||
{
|
||||
"limit_by_header": "x-api-key",
|
||||
"limit_keys": []map[string]interface{}{
|
||||
{"key": "k1", "query_per_minute": 100},
|
||||
},
|
||||
},
|
||||
{
|
||||
"limit_by_param": "apikey",
|
||||
"limit_keys": []map[string]interface{}{
|
||||
{"key": "k2", "query_per_minute": 50},
|
||||
},
|
||||
},
|
||||
},
|
||||
"redis": map[string]interface{}{
|
||||
"service_name": "redis.static",
|
||||
"service_port": 6379,
|
||||
},
|
||||
})
|
||||
return data
|
||||
}()
|
||||
|
||||
func TestParseConfig(t *testing.T) {
|
||||
test.RunGoTest(t, func(t *testing.T) {
|
||||
// 测试全局限流配置解析
|
||||
@@ -385,7 +438,7 @@ func TestOnHttpRequestHeaders(t *testing.T) {
|
||||
// 由于需要调用 Redis,应该返回 HeaderStopAllIterationAndWatermark
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
resp := test.CreateRedisRespArray([]interface{}{1000, 999, 60})
|
||||
resp := multiRuleResp([3]int{1000, 999, 60})
|
||||
// 模拟 Redis 调用响应(允许请求)
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
@@ -409,7 +462,7 @@ func TestOnHttpRequestHeaders(t *testing.T) {
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
// 模拟 Redis 调用响应(允许请求)
|
||||
resp := test.CreateRedisRespArray([]interface{}{10, 9, 60})
|
||||
resp := multiRuleResp([3]int{10, 9, 60})
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
host.CompleteHttp()
|
||||
@@ -433,7 +486,7 @@ func TestOnHttpRequestHeaders(t *testing.T) {
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
// 模拟 Redis 调用响应(允许请求)
|
||||
resp := test.CreateRedisRespArray([]interface{}{10, 9, 60})
|
||||
resp := multiRuleResp([3]int{10, 9, 60})
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
host.CompleteHttp()
|
||||
@@ -457,7 +510,7 @@ func TestOnHttpRequestHeaders(t *testing.T) {
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
// 模拟 Redis 调用响应(允许请求)
|
||||
resp := test.CreateRedisRespArray([]interface{}{10, 9, 1})
|
||||
resp := multiRuleResp([3]int{10, 9, 1})
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
host.CompleteHttp()
|
||||
@@ -481,7 +534,7 @@ func TestOnHttpRequestHeaders(t *testing.T) {
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
// 模拟 Redis 调用响应(允许请求)
|
||||
resp := test.CreateRedisRespArray([]interface{}{10, 9, 60})
|
||||
resp := multiRuleResp([3]int{10, 9, 60})
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
host.CompleteHttp()
|
||||
@@ -505,7 +558,7 @@ func TestOnHttpRequestHeaders(t *testing.T) {
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
// 模拟 Redis 调用响应(允许请求)
|
||||
resp := test.CreateRedisRespArray([]interface{}{10, 9, 86400})
|
||||
resp := multiRuleResp([3]int{10, 9, 86400})
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
host.CompleteHttp()
|
||||
@@ -529,7 +582,7 @@ func TestOnHttpRequestHeaders(t *testing.T) {
|
||||
|
||||
// 模拟 Redis 调用响应(触发限流)
|
||||
// 当前请求数(1001)超过阈值(1000),触发限流
|
||||
resp := test.CreateRedisRespArray([]interface{}{1000, 1001, 60})
|
||||
resp := multiRuleResp([3]int{1000, 1001, 60})
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
// 检查是否发送了限流响应
|
||||
@@ -540,6 +593,224 @@ func TestOnHttpRequestHeaders(t *testing.T) {
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// 混合限流:global + rule_item 同时命中
|
||||
t.Run("hybrid limit both match", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(hybridLimitConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/test"},
|
||||
{":method", "GET"},
|
||||
{"x-api-key", "vip-key"},
|
||||
})
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
resp := multiRuleResp(
|
||||
[3]int{10000, 1, 60},
|
||||
[3]int{100, 1, 60},
|
||||
)
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// 多条 rule_items 同时命中
|
||||
t.Run("multi rule_items all match", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(multiRuleItemsConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/test?apikey=k2"},
|
||||
{":method", "GET"},
|
||||
{"x-api-key", "k1"},
|
||||
})
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
resp := multiRuleResp(
|
||||
[3]int{100, 1, 60},
|
||||
[3]int{50, 1, 60},
|
||||
)
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// 混合限流:global 触发
|
||||
t.Run("hybrid limit global triggered", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(hybridLimitConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/test"},
|
||||
{":method", "GET"},
|
||||
{"x-api-key", "vip-key"},
|
||||
})
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
resp := multiRuleResp(
|
||||
[3]int{10000, 10001, 60},
|
||||
[3]int{100, 1, 60},
|
||||
)
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
localResponse := host.GetLocalResponse()
|
||||
require.NotNil(t, localResponse)
|
||||
require.Equal(t, uint32(429), localResponse.StatusCode)
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// Redis 响应数组长度与规则数不匹配
|
||||
t.Run("redis response length mismatch", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(hybridLimitConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/test"},
|
||||
{":method", "GET"},
|
||||
{"x-api-key", "vip-key"},
|
||||
})
|
||||
|
||||
resp := multiRuleResp([3]int{10000, 1, 60}) // 期望 2 条,返回 1 条
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
require.Nil(t, host.GetLocalResponse())
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// Redis 子数组长度异常(少于 3)
|
||||
t.Run("redis sub-array length mismatch", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(hybridLimitConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/test"},
|
||||
{":method", "GET"},
|
||||
{"x-api-key", "vip-key"},
|
||||
})
|
||||
|
||||
badResp, err := resp.ArrayValue([]resp.Value{
|
||||
resp.ArrayValue([]resp.Value{
|
||||
resp.IntegerValue(10000),
|
||||
resp.IntegerValue(1),
|
||||
}),
|
||||
}).MarshalRESP() // 关键:用 MarshalRESP() 而非 Bytes()
|
||||
require.NoError(t, err)
|
||||
host.CallOnRedisCall(0, badResp)
|
||||
|
||||
require.Nil(t, host.GetLocalResponse())
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// 混合配置:global 命中,rule_items 不命中
|
||||
t.Run("hybrid limit rule_items no match", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(hybridLimitConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/test"},
|
||||
{":method", "GET"},
|
||||
})
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
resp := multiRuleResp([3]int{10000, 1, 60})
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
require.Nil(t, host.GetLocalResponse())
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// 多规则同时触发,验证 rejected 报告 global(reset=60 而非 rule_item 的 30)
|
||||
t.Run("hybrid limit both triggered reports global first", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(hybridLimitConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/test"},
|
||||
{":method", "GET"},
|
||||
{"x-api-key", "vip-key"},
|
||||
})
|
||||
|
||||
resp := multiRuleResp(
|
||||
[3]int{10000, 10001, 60},
|
||||
[3]int{100, 101, 30},
|
||||
)
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
localResponse := host.GetLocalResponse()
|
||||
require.NotNil(t, localResponse)
|
||||
require.Equal(t, uint32(429), localResponse.StatusCode)
|
||||
|
||||
// LocalHttpResponse.Headers 类型为 [][2]string,按切片迭代查找目标头
|
||||
var resetHeader string
|
||||
for _, h := range localResponse.Headers {
|
||||
if strings.EqualFold(h[0], "X-RateLimit-Reset") {
|
||||
resetHeader = h[1]
|
||||
break
|
||||
}
|
||||
}
|
||||
require.Equal(t, "60", resetHeader, "应报告 global 规则的 reset 时间")
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// 多规则未触发,剩余比例不同 → LimitContext 取 tightest
|
||||
// cluster-key 可以通过 X-RateLimit-* 响应头验证(ai-token 不能)
|
||||
t.Run("multi-rule no trigger takes tightest", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(hybridLimitConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/test"},
|
||||
{":method", "GET"},
|
||||
{"x-api-key", "vip-key"},
|
||||
})
|
||||
|
||||
resp := multiRuleResp(
|
||||
[3]int{10000, 9000, 60}, // global: 剩余 10%
|
||||
[3]int{100, 10, 60}, // rule_item: 剩余 90%
|
||||
)
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
require.Nil(t, host.GetLocalResponse())
|
||||
|
||||
host.CallOnHttpResponseHeaders([][2]string{
|
||||
{":status", "200"},
|
||||
{"content-type", "application/json"},
|
||||
})
|
||||
|
||||
responseHeaders := host.GetResponseHeaders()
|
||||
require.True(t, test.HasHeader(responseHeaders, "x-ratelimit-limit"))
|
||||
require.True(t, test.HasHeader(responseHeaders, "x-ratelimit-remaining"))
|
||||
|
||||
for _, h := range responseHeaders {
|
||||
if strings.EqualFold(h[0], "x-ratelimit-limit") {
|
||||
require.Equal(t, "10000", h[1], "X-RateLimit-Limit 应为 tightest(global) 的 threshold")
|
||||
}
|
||||
if strings.EqualFold(h[0], "x-ratelimit-remaining") {
|
||||
require.Equal(t, "1000", h[1], "X-RateLimit-Remaining 应为 tightest(global) 的 threshold-current")
|
||||
}
|
||||
}
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
@@ -559,7 +830,7 @@ func TestOnHttpResponseHeaders(t *testing.T) {
|
||||
})
|
||||
|
||||
// 模拟 Redis 调用响应
|
||||
resp := test.CreateRedisRespArray([]interface{}{1000, 999, 60})
|
||||
resp := multiRuleResp([3]int{1000, 999, 60})
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
// 处理响应头
|
||||
@@ -609,7 +880,7 @@ func TestOnHttpResponseHeaders(t *testing.T) {
|
||||
})
|
||||
|
||||
// 模拟 Redis 调用响应
|
||||
resp := test.CreateRedisRespArray([]interface{}{1000, 999, 60})
|
||||
resp := multiRuleResp([3]int{1000, 999, 60})
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
// 处理响应头
|
||||
@@ -628,6 +899,88 @@ func TestOnHttpResponseHeaders(t *testing.T) {
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// 测试响应头处理时 context 中无 LimitContext(覆盖 GetContext 返回 nil 的分支)。
|
||||
// 场景:使用 headerLimitConfig(仅 x-ca-key 规则),ensureContextInitialized 调用的
|
||||
// onHttpRequestHeaders 因为缺少 x-ca-key 头而 matched 为空,返回 ActionContinue,
|
||||
// 此时 SetContext 未被调用。后续 onHttpResponseHeaders 应通过类型断言失败分支安全返回。
|
||||
t.Run("no prior context", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(headerLimitConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
// 直接处理响应头:测试框架会通过 ensureContextInitialized 触发默认请求头处理,
|
||||
// 但缺少 x-ca-key,因此 collectMatchedRules 返回空,SetContext 未调用。
|
||||
action := host.CallOnHttpResponseHeaders([][2]string{
|
||||
{":status", "200"},
|
||||
{"content-type", "application/json"},
|
||||
})
|
||||
|
||||
// 应该返回 ActionContinue,且不 panic
|
||||
require.Equal(t, types.ActionContinue, action)
|
||||
|
||||
// 不应写入限流配额响应头
|
||||
responseHeaders := host.GetResponseHeaders()
|
||||
require.False(t, test.HasHeader(responseHeaders, "x-ratelimit-limit"))
|
||||
require.False(t, test.HasHeader(responseHeaders, "x-ratelimit-remaining"))
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
// TestRejectedHidesQuotaHeadersWhenDisabled 验证 rejected() 在 show_limit_quota_header=false
|
||||
// 时只设置 X-RateLimit-Reset,不设置 X-RateLimit-Limit / X-RateLimit-Remaining(覆盖 line 339-341)。
|
||||
func TestRejectedHidesQuotaHeadersWhenDisabled(t *testing.T) {
|
||||
test.RunTest(t, func(t *testing.T) {
|
||||
hideQuotaConfig := func() json.RawMessage {
|
||||
data, _ := json.Marshal(map[string]interface{}{
|
||||
"rule_name": "routeA-rejected-hide-quota",
|
||||
"global_threshold": map[string]interface{}{
|
||||
"query_per_minute": 1,
|
||||
},
|
||||
"redis": map[string]interface{}{
|
||||
"service_name": "redis.static",
|
||||
"service_port": 6379,
|
||||
},
|
||||
"show_limit_quota_header": false,
|
||||
"rejected_code": 429,
|
||||
"rejected_msg": "Too many requests",
|
||||
})
|
||||
return data
|
||||
}()
|
||||
|
||||
host, status := test.NewTestHost(hideQuotaConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/api/test"},
|
||||
{":method", "GET"},
|
||||
})
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
// 触发限流:current (2) > threshold (1)
|
||||
resp := multiRuleResp([3]int{1, 2, 60})
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
// 应返回 429 本地响应
|
||||
localResponse := host.GetLocalResponse()
|
||||
require.NotNil(t, localResponse)
|
||||
require.Equal(t, uint32(429), localResponse.StatusCode)
|
||||
|
||||
// X-RateLimit-Reset 应保留(始终设置)
|
||||
require.True(t, test.HasHeader(localResponse.Headers, "x-ratelimit-reset"),
|
||||
"X-RateLimit-Reset should always be present on rejection")
|
||||
|
||||
// X-RateLimit-Limit / X-RateLimit-Remaining 应缺失(show_limit_quota_header=false)
|
||||
require.False(t, test.HasHeader(localResponse.Headers, "x-ratelimit-limit"),
|
||||
"X-RateLimit-Limit should be hidden when show_limit_quota_header=false")
|
||||
require.False(t, test.HasHeader(localResponse.Headers, "x-ratelimit-remaining"),
|
||||
"X-RateLimit-Remaining should be hidden when show_limit_quota_header=false")
|
||||
|
||||
host.CompleteHttp()
|
||||
})
|
||||
}
|
||||
|
||||
@@ -649,7 +1002,7 @@ func TestCompleteFlow(t *testing.T) {
|
||||
require.Equal(t, types.HeaderStopAllIterationAndWatermark, action)
|
||||
|
||||
// 2. 模拟 Redis 调用响应
|
||||
resp := test.CreateRedisRespArray([]interface{}{1000, 1, 60})
|
||||
resp := multiRuleResp([3]int{1000, 1, 60})
|
||||
host.CallOnRedisCall(0, resp)
|
||||
|
||||
// 3. 处理响应头
|
||||
@@ -672,3 +1025,24 @@ func TestCompleteFlow(t *testing.T) {
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
// multiRuleResp 构建多规则嵌套 Redis 响应(RESP wire format)。
|
||||
// test.CreateRedisRespArray 不支持嵌套数组,因此直接用 resp 包构造。
|
||||
// 每个 [3]int 元组为 {threshold, current, ttl}。
|
||||
// 注意:不能使用 resp.Value.Bytes()——它返回的是显示用字符串,不是 RESP wire 格式。
|
||||
// 必须用 resp.Writer.WriteArray 或 Value.MarshalRESP() 输出真正的 RESP 字节流。
|
||||
func multiRuleResp(items ...[3]int) []byte {
|
||||
values := make([]resp.Value, len(items))
|
||||
for i, it := range items {
|
||||
values[i] = resp.ArrayValue([]resp.Value{
|
||||
resp.IntegerValue(it[0]),
|
||||
resp.IntegerValue(it[1]),
|
||||
resp.IntegerValue(it[2]),
|
||||
})
|
||||
}
|
||||
b, err := resp.ArrayValue(values).MarshalRESP()
|
||||
if err != nil {
|
||||
panic(fmt.Sprintf("failed to marshal multiRuleResp: %v", err))
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user