Copilot commented on code in PR #13191:
URL: https://github.com/apache/apisix/pull/13191#discussion_r3057744723
##########
apisix/plugins/ai-rate-limiting.lua:
##########
@@ -136,8 +149,42 @@ local limit_conf_cache = core.lrucache.new({
})
+-- safe math functions allowed in cost expressions
+local expr_safe_env = {
+ math = math,
+ abs = math.abs,
+ ceil = math.ceil,
+ floor = math.floor,
+ max = math.max,
+ min = math.min,
+}
Review Comment:
The runtime env allows upstream-provided usage keys to overwrite
reserved/safe names (e.g., a usage field named `math`/`abs` would shadow the
safe functions), which can break evaluation and potentially bypass rate
limiting. Consider preventing `raw` keys from overwriting names present in
`expr_safe_env` (skip those keys), and avoid exposing the full `math` table if
the intent is to allow only a small whitelist (e.g., provide a minimal `math`
table with only `abs/ceil/floor/max/min`).
##########
apisix/plugins/ai-rate-limiting.lua:
##########
@@ -264,7 +311,51 @@ function _M.check_instance_status(conf, ctx, instance_name)
end
+local function eval_cost_expr(conf_cost_expr, raw)
+ local fn_code = "return " .. conf_cost_expr
+ -- build environment: safe math + usage variables (missing vars default to
0)
+ local env = setmetatable({}, {
+ __index = function(_, k)
+ local v = expr_safe_env[k]
+ if v ~= nil then
+ return v
+ end
+ return 0
+ end
+ })
+ for k, v in pairs(raw) do
+ if type(v) == "number" then
Review Comment:
The runtime env allows upstream-provided usage keys to overwrite
reserved/safe names (e.g., a usage field named `math`/`abs` would shadow the
safe functions), which can break evaluation and potentially bypass rate
limiting. Consider preventing `raw` keys from overwriting names present in
`expr_safe_env` (skip those keys), and avoid exposing the full `math` table if
the intent is to allow only a small whitelist (e.g., provide a minimal `math`
table with only `abs/ceil/floor/max/min`).
```suggestion
local safe_math = {
abs = math.abs,
ceil = math.ceil,
floor = math.floor,
max = math.max,
min = math.min,
}
local reserved_names = {
math = true,
abs = true,
ceil = true,
floor = true,
max = true,
min = true,
}
-- build environment: safe math + usage variables (missing vars default
to 0)
local env = setmetatable({}, {
__index = function(_, k)
if k == "math" then
return safe_math
end
local v = safe_math[k]
if v ~= nil then
return v
end
v = expr_safe_env[k]
if v ~= nil then
return v
end
return 0
end
})
for k, v in pairs(raw) do
if type(v) == "number" and not reserved_names[k] and
expr_safe_env[k] == nil then
```
##########
apisix/plugins/ai-rate-limiting.lua:
##########
@@ -264,7 +311,51 @@ function _M.check_instance_status(conf, ctx, instance_name)
end
+local function eval_cost_expr(conf_cost_expr, raw)
+ local fn_code = "return " .. conf_cost_expr
+ -- build environment: safe math + usage variables (missing vars default to
0)
+ local env = setmetatable({}, {
+ __index = function(_, k)
+ local v = expr_safe_env[k]
+ if v ~= nil then
+ return v
+ end
+ return 0
+ end
+ })
+ for k, v in pairs(raw) do
+ if type(v) == "number" then
+ env[k] = v
+ end
+ end
+ local fn, err = load(fn_code, "cost_expr", "t", env)
+ if not fn then
+ return nil, "failed to compile cost_expr: " .. err
+ end
+ local ok, result = pcall(fn)
Review Comment:
The expression is recompiled (`load`) on every request. This adds avoidable
overhead in a hot path. Consider caching the compiled function keyed by
`conf.cost_expr` (e.g., via an `lrucache`) and only rebuilding the per-request
variable environment (or setting a per-call env via `setfenv`/equivalent) while
reusing the compiled chunk.
##########
apisix/plugins/ai-rate-limiting.lua:
##########
@@ -264,7 +311,51 @@ function _M.check_instance_status(conf, ctx, instance_name)
end
+local function eval_cost_expr(conf_cost_expr, raw)
+ local fn_code = "return " .. conf_cost_expr
Review Comment:
The expression is recompiled (`load`) on every request. This adds avoidable
overhead in a hot path. Consider caching the compiled function keyed by
`conf.cost_expr` (e.g., via an `lrucache`) and only rebuilding the per-request
variable environment (or setting a per-call env via `setfenv`/equivalent) while
reusing the compiled chunk.
##########
t/plugin/ai-rate-limiting-expression.t:
##########
@@ -0,0 +1,551 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+BEGIN {
+ $ENV{TEST_ENABLE_CONTROL_API_V1} = "0";
+}
+
+use t::APISIX 'no_plan';
+
+log_level("info");
+repeat_each(1);
+no_long_string();
+no_shuffle();
+no_root_location();
+
+add_block_preprocessor(sub {
+ my ($block) = @_;
+
+ if (!defined $block->request) {
+ $block->set_value("request", "GET /t");
+ }
+
+ my $http_config = $block->http_config // <<_EOC_;
+ server {
+ server_name anthropic;
+ listen 16725;
+
+ default_type 'application/json';
+
+ location /v1/messages {
+ content_by_lua_block {
+ local json = require("cjson.safe")
+ local ngx = ngx
+
+ ngx.req.read_body()
+ local body = ngx.req.get_body_data()
+ body = json.decode(body)
+
+ if not body or not body.messages then
+ ngx.status = 400
+
ngx.say('{"type":"error","error":{"type":"invalid_request_error","message":"missing
messages"}}')
+ return
+ end
+
+ local api_key = ngx.req.get_headers()["x-api-key"]
+ if api_key ~= "test-key" then
+ ngx.status = 401
+
ngx.say('{"type":"error","error":{"type":"authentication_error","message":"invalid
x-api-key"}}')
+ return
+ end
+
+ if body.stream then
+ ngx.header["Content-Type"] = "text/event-stream"
+
+ -- message_start with input_tokens and cache tokens
+ local message_start = json.encode({
+ type = "message_start",
+ message = {
+ id = "msg_test123",
+ type = "message",
+ role = "assistant",
+ model = body.model or
"claude-sonnet-4-20250514",
+ content = {},
+ usage = {
+ input_tokens = 50,
+ output_tokens = 0,
+ cache_creation_input_tokens = 100,
+ cache_read_input_tokens = 200,
+ },
+ },
+ })
+ ngx.say("event: message_start")
+ ngx.say("data: " .. message_start)
+ ngx.say("")
+
+ -- content_block_start
+ ngx.say("event: content_block_start")
+ ngx.say('data:
{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}')
+ ngx.say("")
+
+ -- content_block_delta
+ ngx.say("event: content_block_delta")
+ ngx.say('data:
{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello
from Claude!"}}')
+ ngx.say("")
+
+ -- content_block_stop
+ ngx.say("event: content_block_stop")
+ ngx.say('data:
{"type":"content_block_stop","index":0}')
+ ngx.say("")
+
+ -- message_delta with output_tokens
+ local message_delta = json.encode({
+ type = "message_delta",
+ delta = { stop_reason = "end_turn" },
+ usage = {
+ output_tokens = 30,
+ },
+ })
+ ngx.say("event: message_delta")
+ ngx.say("data: " .. message_delta)
+ ngx.say("")
+
+ -- message_stop
+ ngx.say("event: message_stop")
+ ngx.say("data: {}")
+ ngx.say("")
+ else
+ ngx.status = 200
+ ngx.say(json.encode({
+ id = "msg_test456",
+ type = "message",
+ role = "assistant",
+ model = body.model or "claude-sonnet-4-20250514",
+ content = {{
+ type = "text",
+ text = "Hello from Claude!",
+ }},
+ stop_reason = "end_turn",
+ usage = {
+ input_tokens = 50,
+ output_tokens = 30,
+ cache_creation_input_tokens = 100,
+ cache_read_input_tokens = 200,
+ },
+ }))
+ end
+ }
+ }
+ }
+_EOC_
+
+ $block->set_value("http_config", $http_config);
+});
+
+run_tests();
+
+__DATA__
+
+=== TEST 1: schema validation - expression strategy requires cost_expr
+--- config
+ location /t {
+ content_by_lua_block {
+ local plugin = require("apisix.plugins.ai-rate-limiting")
+ local configs = {
+ -- expression without cost_expr
+ {
+ limit = 100,
+ time_window = 60,
+ limit_strategy = "expression",
+ },
+ -- expression with empty cost_expr
+ {
+ limit = 100,
+ time_window = 60,
+ limit_strategy = "expression",
+ cost_expr = "",
+ },
+ -- expression with invalid cost_expr syntax
+ {
+ limit = 100,
+ time_window = 60,
+ limit_strategy = "expression",
+ cost_expr = "invalid $$$ syntax %%%",
+ },
+ -- valid expression
+ {
+ limit = 100,
+ time_window = 60,
+ limit_strategy = "expression",
+ cost_expr = "input_tokens + output_tokens",
+ },
+ -- valid complex expression
+ {
+ limit = 100,
+ time_window = 60,
+ limit_strategy = "expression",
+ cost_expr = "(input_tokens - cache_read_input_tokens) +
cache_creation_input_tokens * 1.25 + output_tokens",
+ },
Review Comment:
Given expressions can subtract fields, it would be useful to add a test case
where the expression evaluates to a negative value (e.g., `input_tokens -
cache_read_input_tokens` with cache_read > input) and assert the intended
behavior (clamp to 0 vs reject vs treat as error). This will lock in the safety
behavior and prevent future bypasses.
##########
apisix/plugins/ai-rate-limiting.lua:
##########
@@ -264,7 +311,51 @@ function _M.check_instance_status(conf, ctx, instance_name)
end
+local function eval_cost_expr(conf_cost_expr, raw)
Review Comment:
Using `load()` on an arbitrary expression string allows constructs that are
not 'arithmetic' (e.g., `(function() while true do end end)()`), which can hang
a worker or consume excessive CPU/memory if configured incorrectly. If the
feature is intended to support arithmetic only, add a strict validator at
config-time (e.g., deny Lua keywords like `function/while/for/repeat/do/end`,
quotes, braces, semicolons, etc., and allow only
identifiers/numbers/operators/parentheses/whitespace) before calling `load()`.
```suggestion
local function validate_cost_expr(conf_cost_expr)
if type(conf_cost_expr) ~= "string" or conf_cost_expr == "" then
return nil, "cost_expr must be a non-empty string"
end
-- Allow only arithmetic-style expressions: identifiers, numbers,
-- arithmetic operators, parentheses, commas, decimal points and
whitespace.
if conf_cost_expr:find("[\"'%[%]{};:=]") then
return nil, "cost_expr contains disallowed syntax"
end
if conf_cost_expr:find("%-%-") or conf_cost_expr:find("%.%.") then
return nil, "cost_expr contains disallowed token sequence"
end
if conf_cost_expr:find("[^%w_%s%+%-%*%/%^%%%(%),%.]") then
return nil, "cost_expr contains invalid characters"
end
local reserved_words = {
"and", "break", "do", "else", "elseif", "end", "false", "for",
"function", "goto", "if", "in", "local", "nil", "not", "or",
"repeat", "return", "then", "true", "until", "while"
}
for _, word in ipairs(reserved_words) do
if conf_cost_expr:find("%f[%a_]" .. word .. "%f[^%w_]") then
return nil, "cost_expr contains disallowed keyword: " .. word
end
end
local depth = 0
for i = 1, #conf_cost_expr do
local ch = conf_cost_expr:sub(i, i)
if ch == "(" then
depth = depth + 1
elseif ch == ")" then
depth = depth - 1
if depth < 0 then
return nil, "cost_expr has unbalanced parentheses"
end
end
end
if depth ~= 0 then
return nil, "cost_expr has unbalanced parentheses"
end
return true
end
local function eval_cost_expr(conf_cost_expr, raw)
local ok, err = validate_cost_expr(conf_cost_expr)
if not ok then
return nil, "invalid cost_expr: " .. err
end
```
##########
apisix/plugins/ai-rate-limiting.lua:
##########
@@ -264,7 +311,51 @@ function _M.check_instance_status(conf, ctx, instance_name)
end
+local function eval_cost_expr(conf_cost_expr, raw)
+ local fn_code = "return " .. conf_cost_expr
+ -- build environment: safe math + usage variables (missing vars default to
0)
+ local env = setmetatable({}, {
+ __index = function(_, k)
+ local v = expr_safe_env[k]
+ if v ~= nil then
+ return v
+ end
+ return 0
+ end
+ })
+ for k, v in pairs(raw) do
+ if type(v) == "number" then
+ env[k] = v
+ end
+ end
+ local fn, err = load(fn_code, "cost_expr", "t", env)
Review Comment:
Using `load()` on an arbitrary expression string allows constructs that are
not 'arithmetic' (e.g., `(function() while true do end end)()`), which can hang
a worker or consume excessive CPU/memory if configured incorrectly. If the
feature is intended to support arithmetic only, add a strict validator at
config-time (e.g., deny Lua keywords like `function/while/for/repeat/do/end`,
quotes, braces, semicolons, etc., and allow only
identifiers/numbers/operators/parentheses/whitespace) before calling `load()`.
##########
apisix/plugins/ai-rate-limiting.lua:
##########
@@ -136,8 +149,42 @@ local limit_conf_cache = core.lrucache.new({
})
+-- safe math functions allowed in cost expressions
+local expr_safe_env = {
+ math = math,
+ abs = math.abs,
+ ceil = math.ceil,
+ floor = math.floor,
+ max = math.max,
+ min = math.min,
+}
+
+local function compile_cost_expr(expr_str)
+ local fn_code = "return " .. expr_str
+ -- validate syntax by loading first
+ local fn, err = load(fn_code, "cost_expr", "t", expr_safe_env)
+ if not fn then
+ return nil, err
Review Comment:
Using `load()` on an arbitrary expression string allows constructs that are
not 'arithmetic' (e.g., `(function() while true do end end)()`), which can hang
a worker or consume excessive CPU/memory if configured incorrectly. If the
feature is intended to support arithmetic only, add a strict validator at
config-time (e.g., deny Lua keywords like `function/while/for/repeat/do/end`,
quotes, braces, semicolons, etc., and allow only
identifiers/numbers/operators/parentheses/whitespace) before calling `load()`.
```suggestion
local cost_expr_disallowed_keywords = {
"and", "break", "do", "else", "elseif", "end", "false", "for",
"function", "goto", "if", "in", "local", "nil", "not", "or",
"repeat", "return", "then", "true", "until", "while",
}
local function validate_cost_expr(expr_str)
if type(expr_str) ~= "string" or expr_str == "" then
return nil, "cost_expr must be a non-empty string"
end
if expr_str:find("[\"'%[%]{};:]") then
return nil, "cost_expr contains disallowed characters"
end
if expr_str:find("%-%-") or expr_str:find("%.%.") then
return nil, "cost_expr contains disallowed tokens"
end
if expr_str:find("[^%w_%s%+%-%*/%%%^%(%)%,%.]") then
return nil, "cost_expr may only contain identifiers, numbers,
arithmetic operators, parentheses, commas, dots, and whitespace"
end
for _, keyword in ipairs(cost_expr_disallowed_keywords) do
if expr_str:find("%f[%a_]" .. keyword .. "%f[^%a_]") then
return nil, "cost_expr contains disallowed keyword: " .. keyword
end
end
return true
end
local function compile_cost_expr(expr_str)
local ok, err = validate_cost_expr(expr_str)
if not ok then
return nil, err
end
local fn_code = "return " .. expr_str
-- validate syntax after enforcing the arithmetic-only subset
local fn, load_err = load(fn_code, "cost_expr", "t", expr_safe_env)
if not fn then
return nil, load_err
```
##########
apisix/plugins/ai-rate-limiting.lua:
##########
@@ -264,7 +311,51 @@ function _M.check_instance_status(conf, ctx, instance_name)
end
+local function eval_cost_expr(conf_cost_expr, raw)
+ local fn_code = "return " .. conf_cost_expr
+ -- build environment: safe math + usage variables (missing vars default to
0)
+ local env = setmetatable({}, {
+ __index = function(_, k)
+ local v = expr_safe_env[k]
+ if v ~= nil then
+ return v
+ end
+ return 0
+ end
+ })
+ for k, v in pairs(raw) do
+ if type(v) == "number" then
+ env[k] = v
+ end
+ end
+ local fn, err = load(fn_code, "cost_expr", "t", env)
+ if not fn then
+ return nil, "failed to compile cost_expr: " .. err
+ end
+ local ok, result = pcall(fn)
+ if not ok then
+ return nil, "failed to evaluate cost_expr: " .. result
+ end
+ if type(result) ~= "number" then
+ return nil, "cost_expr must return a number, got: " .. type(result)
+ end
+ return math_floor(result + 0.5)
Review Comment:
The expression result is accepted as any Lua number, including negative
values and non-finite values (NaN/inf). A negative cost can effectively
'credit' tokens back (depending on how the counter is applied), and NaN/inf can
break downstream logic. Consider enforcing `result` to be finite and >= 0
(clamp or treat as an error), and document/choose a rounding strategy that
avoids undercounting (often `ceil` is safer for rate limiting when fractional
weights are allowed).
##########
apisix/plugins/ai-rate-limiting.lua:
##########
@@ -61,10 +65,19 @@ local schema = {
show_limit_quota_header = {type = "boolean", default = true},
limit_strategy = {
type = "string",
- enum = {"total_tokens", "prompt_tokens", "completion_tokens"},
+ enum = {"total_tokens", "prompt_tokens", "completion_tokens",
"expression"},
default = "total_tokens",
description = "The strategy to limit the tokens"
},
+ cost_expr = {
+ type = "string",
+ minLength = 1,
+ description = "Lua arithmetic expression for dynamic token cost
calculation. "
+ .. "Variables are injected from the LLM API raw usage response
fields. "
+ .. "Missing variables default to 0. "
+ .. "Only valid when limit_strategy is 'expression'. "
+ .. "Example: input_tokens + cache_creation_input_tokens +
output_tokens",
+ },
Review Comment:
`cost_expr` is described as 'Only valid when limit_strategy is expression',
but the schema/check currently does not reject configs that set `cost_expr`
with a different `limit_strategy`. Consider explicitly rejecting `cost_expr`
unless `limit_strategy == 'expression'` to avoid silent
misconfiguration/confusion.
##########
apisix/plugins/ai-rate-limiting.lua:
##########
@@ -136,8 +149,42 @@ local limit_conf_cache = core.lrucache.new({
})
+-- safe math functions allowed in cost expressions
+local expr_safe_env = {
+ math = math,
+ abs = math.abs,
+ ceil = math.ceil,
+ floor = math.floor,
+ max = math.max,
+ min = math.min,
+}
+
+local function compile_cost_expr(expr_str)
+ local fn_code = "return " .. expr_str
+ -- validate syntax by loading first
+ local fn, err = load(fn_code, "cost_expr", "t", expr_safe_env)
+ if not fn then
+ return nil, err
+ end
+ return fn_code
+end
+
+
function _M.check_schema(conf)
Review Comment:
`cost_expr` is described as 'Only valid when limit_strategy is expression',
but the schema/check currently does not reject configs that set `cost_expr`
with a different `limit_strategy`. Consider explicitly rejecting `cost_expr`
unless `limit_strategy == 'expression'` to avoid silent
misconfiguration/confusion.
##########
apisix/plugins/ai-rate-limiting.lua:
##########
@@ -264,7 +311,51 @@ function _M.check_instance_status(conf, ctx, instance_name)
end
+local function eval_cost_expr(conf_cost_expr, raw)
+ local fn_code = "return " .. conf_cost_expr
+ -- build environment: safe math + usage variables (missing vars default to
0)
+ local env = setmetatable({}, {
+ __index = function(_, k)
+ local v = expr_safe_env[k]
+ if v ~= nil then
+ return v
+ end
+ return 0
+ end
+ })
+ for k, v in pairs(raw) do
+ if type(v) == "number" then
+ env[k] = v
+ end
+ end
+ local fn, err = load(fn_code, "cost_expr", "t", env)
+ if not fn then
+ return nil, "failed to compile cost_expr: " .. err
+ end
+ local ok, result = pcall(fn)
Review Comment:
Using `load()` on an arbitrary expression string allows constructs that are
not 'arithmetic' (e.g., `(function() while true do end end)()`), which can hang
a worker or consume excessive CPU/memory if configured incorrectly. If the
feature is intended to support arithmetic only, add a strict validator at
config-time (e.g., deny Lua keywords like `function/while/for/repeat/do/end`,
quotes, braces, semicolons, etc., and allow only
identifiers/numbers/operators/parentheses/whitespace) before calling `load()`.
##########
apisix/plugins/ai-rate-limiting.lua:
##########
@@ -136,8 +149,42 @@ local limit_conf_cache = core.lrucache.new({
})
+-- safe math functions allowed in cost expressions
+local expr_safe_env = {
+ math = math,
+ abs = math.abs,
+ ceil = math.ceil,
+ floor = math.floor,
+ max = math.max,
+ min = math.min,
+}
+
+local function compile_cost_expr(expr_str)
+ local fn_code = "return " .. expr_str
+ -- validate syntax by loading first
+ local fn, err = load(fn_code, "cost_expr", "t", expr_safe_env)
+ if not fn then
+ return nil, err
+ end
+ return fn_code
+end
+
+
function _M.check_schema(conf)
- return core.schema.check(schema, conf)
+ local ok, err = core.schema.check(schema, conf)
+ if not ok then
+ return false, err
+ end
+ if conf.limit_strategy == "expression" then
Review Comment:
`cost_expr` is described as 'Only valid when limit_strategy is expression',
but the schema/check currently does not reject configs that set `cost_expr`
with a different `limit_strategy`. Consider explicitly rejecting `cost_expr`
unless `limit_strategy == 'expression'` to avoid silent
misconfiguration/confusion.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
