feat(limit-count): support configuring multiple rules

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

Author: shreemaan-abhishek

apisix/plugins/limit-count/init.lua

@@ -25,6 +25,7 @@ local get_phase = ngx.get_phase
 local tonumber = tonumber
 local type = type
 local tostring = tostring
+local str_format = string.format
 
 local limit_redis_cluster_new
 local limit_redis_new
@@ -82,6 +83,28 @@ local schema = {
                 {type = "string"},
             },
         },
+        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"},
+                },
+                required = {"count", "time_window", "key"},
+            },
+        },
         group = {type = "string"},
         key = {type = "string", default = "remote_addr"},
         key_type = {type = "string",
@@ -102,7 +125,14 @@ local schema = {
         allow_degradation = {type = "boolean", default = false},
         show_limit_quota_header = {type = "boolean", default = true}
     },
-    required = {"count", "time_window"},
+    oneOf = {
+        {
+            required = {"count", "time_window"},
+        },
+        {
+            required = {"rules"},
+        }
+    },
     ["if"] = {
         properties = {
             policy = {
@@ -180,51 +210,34 @@ function _M.check_schema(conf, schema_type)
         end
     end
 
-    return true
-end
-
-
-local function create_limit_obj(conf, ctx, plugin_name)
-    core.log.info("create new " .. plugin_name .. " plugin instance")
-
-    local count = conf.count
-    if type(count) == "string" then
-        local err, _
-        count, err, _ = core.utils.resolve_var(count, ctx.var)
-        if err then
-            return nil, "could not resolve vars in count: " .. err
-        end
-        count = tonumber(count)
-        if not count then
-            return nil, "resolved count is not a number: " .. tostring(count)
+    local keys = {}
+    for _, rule in ipairs(conf.rules or {}) do
+        if keys[rule.key] then
+            return false, str_format("duplicate key '%s' in rules", rule.key)
         end
+        keys[rule.key] = true
     end
 
-    local time_window = conf.time_window
-    if type(time_window) == "string" then
-        local err, _
-        time_window, err, _ = core.utils.resolve_var(time_window, ctx.var)
-        if err then
-            return nil, "could not resolve vars in time_window: " .. err
-        end
-        time_window = tonumber(time_window)
-        if not time_window then
-            return nil, "resolved time_window is not a number: " .. tostring(time_window)
-        end
-    end
+    return true
+end
+
 
-    core.log.info("limit count: ", count, ", time_window: ", time_window)
+local function create_limit_obj(conf, rule, plugin_name)
+    core.log.info("create new " .. plugin_name .. " plugin instance",
+        ", rule: ", core.json.delay_encode(rule, true))
 
     if not conf.policy or conf.policy == "local" then
-        return limit_local_new("plugin-" .. plugin_name, count, time_window)
+        return limit_local_new("plugin-" .. plugin_name, rule.count,
+                               rule.time_window)
     end
 
     if conf.policy == "redis" then
-        return limit_redis_new("plugin-" .. plugin_name, count, time_window, conf)
+        return limit_redis_new("plugin-" .. plugin_name, rule.count, rule.time_window, conf)
     end
 
     if conf.policy == "redis-cluster" then
-        return limit_redis_cluster_new("plugin-" .. plugin_name, count, time_window, conf)
+        return limit_redis_cluster_new("plugin-" .. plugin_name, rule.count,
+                                       rule.time_window, conf)
     end
 
     return nil
@@ -258,11 +271,71 @@ local function gen_limit_key(conf, ctx, key)
 end
 
 
-function _M.rate_limit(conf, ctx, name, cost, dry_run)
-    core.log.info("ver: ", ctx.conf_version)
-    core.log.info("conf: ", core.json.delay_encode(conf, true))
+local function resolve_var(ctx, value)
+    if type(value) == "string" then
+        local err, _
+        value, err, _ = core.utils.resolve_var(value, ctx.var)
+        if err then
+            return nil, "could not resolve var for value: " .. value .. ", err: " .. err
+        end
+        value = tonumber(value)
+        if not value then
+            return nil, "resolved value is not a number: " .. tostring(value)
+        end
+    end
+    return value
+end
 
-    local lim, err = create_limit_obj(conf, ctx, name)
+
+local function get_rules(ctx, conf)
+    if not conf.rules then
+        local count, err = resolve_var(ctx, conf.count)
+        if err then
+            return nil, err
+        end
+        local time_window, err2 = resolve_var(ctx, conf.time_window)
+        if err2 then
+            return nil, err2
+        end
+        return {
+            {
+                count = count,
+                time_window = time_window,
+                key = conf.key,
+                key_type = conf.key_type,
+            }
+        }
+    end
+
+    local rules = {}
+    for _, rule in ipairs(conf.rules) do
+        local count, err = resolve_var(ctx, rule.count)
+        if err then
+            goto CONTINUE
+        end
+        local time_window, err2 = resolve_var(ctx, rule.time_window)
+        if err2 then
+            goto CONTINUE
+        end
+        local key, _, n_resolved = core.utils.resolve_var(rule.key, ctx.var)
+        if n_resolved == 0 then
+            goto CONTINUE
+        end
+        core.table.insert(rules, {
+            count = count,
+            time_window = time_window,
+            key_type = "constant",
+            key = key,
+        })
+
+        ::CONTINUE::
+    end
+    return rules
+end
+
+
+local function run_rate_limit(conf, rule, ctx, name, cost, dry_run)
+    local lim, err = create_limit_obj(conf, rule, name)
 
     if not lim then
         core.log.error("failed to fetch limit.count object: ", err)
@@ -272,9 +345,9 @@ function _M.rate_limit(conf, ctx, name, cost, dry_run)
         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
@@ -284,7 +357,7 @@ function _M.rate_limit(conf, ctx, name, cost, dry_run)
         if n_resolved == 0 then
             key = nil
         end
-    elseif conf.key_type == "constant" then
+    elseif rule.key_type == "constant" then
         key = conf_key
     else
         key = ctx.var[conf_key]
@@ -353,4 +426,25 @@ function _M.rate_limit(conf, ctx, name, cost, dry_run)
 end
 
 
+function _M.rate_limit(conf, ctx, name, cost, dry_run)
+    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 rate limit rules: ", err)
+        if conf.allow_degradation then
+            return
+        end
+        return 500
+    end
+
+    for _, rule in ipairs(rules) do
+        local code, msg = run_rate_limit(conf, rule, ctx, name, cost, dry_run)
+        if code then
+            return code, msg
+        end
+    end
+end
+
+
 return _M