feat: enhance rate limiting rules support

Signed-off-by: Abhishek Choudhary <shreemaan.abhishek@gmail.com>

Author: shreemaan-abhishek

apisix/core/utils.lua

@@ -310,7 +310,10 @@ do
         local v = _ctx[segs[1]]
 
         if v == nil then
-            return segs[2] or ""
+            if #segs == 1 then
+                return ""
+            end
+            v = segs[2]
         end
         n_resolved = n_resolved + 1
         if _escaper then

apisix/plugins/ai-rate-limiting.lua

@@ -76,17 +76,50 @@ local schema = {
         rejected_msg = {
             type = "string", minLength = 1
         },
+        rules = {
+            type = "array",
+            items = {
+                type = "object",
+                properties = {
+                    count = {
+                        oneOf = {
+                            {type = "integer", exclusiveMinimum = 0},
+                            {type = "string"},
+                        },
+                    },
+                    time_window = {
+                        oneOf = {
+                            {type = "integer", exclusiveMinimum = 0},
+                            {type = "string"},
+                        },
+                    },
+                    key = {type = "string"},
+                    header_prefix = {
+                        type = "string",
+                        description = "prefix for rate limit headers"
+                    },
+                },
+                required = {"count", "time_window", "key"},
+            },
+        },
     },
     dependencies = {
         limit = {"time_window"},
         time_window = {"limit"}
     },
-    anyOf = {
+    oneOf = {
         {
-            required = {"limit", "time_window"}
+            anyOf = {
+                {
+                    required = {"limit", "time_window"}
+                },
+                {
+                    required = {"instances"}
+                }
+            }
         },
         {
-            required = {"instances"}
+            required = {"rules"},
         }
     }
 }
@@ -109,6 +142,26 @@ end
 
 
 local function transform_limit_conf(plugin_conf, instance_conf, instance_name)
+    local limit_conf = {
+        rejected_code = plugin_conf.rejected_code,
+        rejected_msg = plugin_conf.rejected_msg,
+        show_limit_quota_header = plugin_conf.show_limit_quota_header,
+
+        -- we may expose those fields to ai-rate-limiting later
+        policy = "local",
+        key_type = "constant",
+        allow_degradation = false,
+        sync_interval = -1,
+        limit_header = "X-AI-RateLimit-Limit",
+        remaining_header = "X-AI-RateLimit-Remaining",
+        reset_header = "X-AI-RateLimit-Reset",
+    }
+    if plugin_conf.rules and #plugin_conf.rules > 0 then
+        limit_conf.rules = plugin_conf.rules
+        limit_conf._meta = plugin_conf._meta
+        return limit_conf
+    end
+
     local key = plugin_name .. "#global"
     local limit = plugin_conf.limit
     local time_window = plugin_conf.time_window
@@ -119,26 +172,15 @@ local function transform_limit_conf(plugin_conf, instance_conf, instance_name)
         limit = instance_conf.limit
         time_window = instance_conf.time_window
     end
-    return {
-        _vid = key,
-
-        key = key,
-        _meta = plugin_conf._meta,
-        count = limit,
-        time_window = time_window,
-        rejected_code = plugin_conf.rejected_code,
-        rejected_msg = plugin_conf.rejected_msg,
-        show_limit_quota_header = plugin_conf.show_limit_quota_header,
-        -- limit-count need these fields
-        policy = "local",
-        key_type = "constant",
-        allow_degradation = false,
-        sync_interval = -1,
-
-        limit_header = "X-AI-RateLimit-Limit-" .. name,
-        remaining_header = "X-AI-RateLimit-Remaining-" .. name,
-        reset_header = "X-AI-RateLimit-Reset-" .. name,
-    }
+    limit_conf._vid = key
+    limit_conf.key = key
+    limit_conf._meta = plugin_conf._meta
+    limit_conf.count = limit
+    limit_conf.time_window = time_window
+    limit_conf.limit_header = "X-AI-RateLimit-Limit-" .. name
+    limit_conf.remaining_header = "X-AI-RateLimit-Remaining-" .. name
+    limit_conf.reset_header = "X-AI-RateLimit-Reset-" .. name
+    return limit_conf
 end
 
 
@@ -169,8 +211,13 @@ function _M.access(conf, ctx)
         return
     end
 
-    local limit_conf_kvs = limit_conf_cache(conf, nil, fetch_limit_conf_kvs, conf)
-    local limit_conf = limit_conf_kvs[ai_instance_name]
+    local limit_conf
+    if conf.rules and #conf.rules > 0 then
+        limit_conf = transform_limit_conf(conf)
+    else
+        local limit_conf_kvs = limit_conf_cache(conf, nil, fetch_limit_conf_kvs, conf)
+        limit_conf = limit_conf_kvs[ai_instance_name]
+    end
     if not limit_conf then
         return
     end
@@ -244,8 +291,13 @@ function _M.log(conf, ctx)
 
     core.log.info("instance name: ", instance_name, " used tokens: ", used_tokens)
 
-    local limit_conf_kvs = limit_conf_cache(conf, nil, fetch_limit_conf_kvs, conf)
-    local limit_conf = limit_conf_kvs[instance_name]
+    local limit_conf
+    if conf.rules and #conf.rules > 0 then
+        limit_conf = transform_limit_conf(conf)
+    else
+        local limit_conf_kvs = limit_conf_cache(conf, nil, fetch_limit_conf_kvs, conf)
+        limit_conf = limit_conf_kvs[instance_name]
+    end
     if limit_conf then
         limit_count.rate_limit(limit_conf, ctx, plugin_name, used_tokens)
     end

apisix/plugins/limit-conn.lua

@@ -54,9 +54,38 @@ local schema = {
         rejected_msg = {
             type = "string", minLength = 1
         },
-        allow_degradation = {type = "boolean", default = false}
+        allow_degradation = {type = "boolean", default = false},
+        rules = {
+            type = "array",
+            items = {
+                type = "object",
+                properties = {
+                    conn = {
+                        oneOf = {
+                            {type = "integer", exclusiveMinimum = 0},
+                            {type = "string"},
+                        },
+                    },
+                    burst = {
+                        oneOf = {
+                            {type = "integer", minimum = 0},
+                            {type = "string"},
+                        },
+                    },
+                    key = {type = "string"},
+                },
+                required = {"conn", "burst", "key"},
+            },
+        },
+    },
+    oneOf = {
+        {
+            required = {"conn", "burst", "default_conn_delay", "key"},
+        },
+        {
+            required = {"default_conn_delay", "rules"},
+        }
     },
-    required = {"conn", "burst", "default_conn_delay", "key"},
     ["if"] = {
         properties = {
             policy = {

apisix/plugins/limit-conn/init.lua

@@ -40,60 +40,100 @@ end
 local _M = {}
 
 
-local function create_limit_obj(ctx, conf)
-    core.log.info("create new limit-conn plugin instance")
-
-    local conn = conf.conn
-    if type(conn) == "string" then
+local function resolve_var(ctx, value)
+    if type(value) == "string" then
         local err, _
-        conn, err, _ = core.utils.resolve_var(conn, ctx.var)
+        value, err, _ = core.utils.resolve_var(value, ctx.var)
         if err then
-            return nil, "could not resolve vars in conn: " .. err
+            return nil, "could not resolve var for value: " .. value .. ", err: " .. err
         end
-        conn = tonumber(conn)
-        if not conn then
-            return nil, "resolved conn is not a number: " .. tostring(conn)
+        value = tonumber(value)
+        if not value then
+            return nil, "resolved value is not a number: " .. tostring(value)
         end
     end
+    return value
+end
 
-    local burst = conf.burst
-    if type(burst) == "string" then
-        local err, _
-        burst, err, _ = core.utils.resolve_var(burst, ctx.var)
+
+local function get_rules(ctx, conf)
+    if not conf.rules then
+        local conn, err = resolve_var(ctx, conf.conn)
+        if err then
+            return nil, err
+        end
+        local burst, err2 = resolve_var(ctx, conf.burst)
+        if err2 then
+            return nil, err2
+        end
+        return {
+            {
+                conn = conn,
+                burst = burst,
+                key = conf.key,
+                key_type = conf.key_type,
+            }
+        }
+    end
+
+    local rules = {}
+    for _, rule in ipairs(conf.rules) do
+        local conn, err = resolve_var(ctx, rule.conn)
         if err then
-            return nil, "could not resolve vars in burst: " .. err
+            goto CONTINUE
+        end
+        local burst, err2 = resolve_var(ctx, rule.burst)
+        if err2 then
+            goto CONTINUE
         end
-        burst = tonumber(burst)
-        if not burst then
-            return nil, "resolved burst is not a number: " .. tostring(burst)
+
+        local key, _, n_resolved = core.utils.resolve_var(rule.key, ctx.var)
+        if n_resolved == 0 then
+            goto CONTINUE
         end
+        core.table.insert(rules, {
+            conn = conn,
+            burst = burst,
+            key_type = "constant",
+            key = key,
+        })
+
+        ::CONTINUE::
     end
+    return rules
+end
+
+
+local function create_limit_obj(conf, rule, default_conn_delay)
+    core.log.info("create new limit-conn plugin instance")
+
+    local conn = rule.conn
+    local burst = rule.burst
 
     core.log.info("limit conn: ", conn, ", burst: ", burst)
 
     if conf.policy == "redis" then
         core.log.info("create new limit-conn redis plugin instance")
 
         return redis_single_new("plugin-limit-conn", conf, conn, burst,
-                                conf.default_conn_delay)
+                                default_conn_delay)
 
     elseif conf.policy == "redis-cluster" then
 
         core.log.info("create new limit-conn redis-cluster plugin instance")
 
         return redis_cluster_new("plugin-limit-conn", conf, conn, burst,
-                                 conf.default_conn_delay)
+                                 default_conn_delay)
     else
         core.log.info("create new limit-conn plugin instance")
         return limit_conn_new(shdict_name, conn, burst,
-                              conf.default_conn_delay)
+                              default_conn_delay)
     end
 end
 
 
-function _M.increase(conf, ctx)
-    core.log.info("ver: ", ctx.conf_version)
-    local lim, err = create_limit_obj(ctx, conf)
+local function run_limit_conn(conf, rule, ctx)
+    local lim, err = create_limit_obj(conf, rule, conf.default_conn_delay)
     if not lim then
         core.log.error("failed to instantiate a resty.limit.conn object: ", err)
         if conf.allow_degradation then
@@ -102,9 +142,9 @@ function _M.increase(conf, ctx)
         return 500
     end
 
-    local conf_key = conf.key
+    local conf_key = rule.key
     local key
-    if conf.key_type == "var_combination" then
+    if rule.key_type == "var_combination" then
         local err, n_resolved
         key, err, n_resolved = core.utils.resolve_var(conf_key, ctx.var)
         if err then
@@ -114,6 +154,8 @@ function _M.increase(conf, ctx)
         if n_resolved == 0 then
             key = nil
         end
+    elseif rule.key_type == "constant" then
+        key = conf_key
     else
         key = ctx.var[conf_key]
     end
@@ -157,6 +199,27 @@ function _M.increase(conf, ctx)
 end
 
 
+function _M.increase(conf, ctx)
+    core.log.info("ver: ", ctx.conf_version)
+
+    local rules, err = get_rules(ctx, conf)
+    if not rules or #rules == 0 then
+        core.log.error("failed to get limit conn rules: ", err)
+        if conf.allow_degradation then
+            return
+        end
+        return 500
+    end
+
+    for _, rule in ipairs(rules) do
+        local code, msg = run_limit_conn(conf, rule, ctx)
+        if code then
+            return code, msg
+        end
+    end
+end
+
+
 function _M.decrease(conf, ctx)
     local limit_conn = ctx.limit_conn
     if not limit_conn then