feat(limit-count): make rate-limit response headers configurable

apisix/plugins/limit-count.lua

@@ -23,11 +23,12 @@ local _M = {
     priority = 1002,
     name = plugin_name,
     schema = limit_count.schema,
+    metadata_schema = limit_count.metadata_schema,
 }
 
 
-function _M.check_schema(conf)
-    return limit_count.check_schema(conf)
+function _M.check_schema(conf, schema_type)
+    return limit_count.check_schema(conf, schema_type)
 end
 
 

apisix/plugins/limit-count/init.lua

@@ -42,6 +42,30 @@ local group_conf_lru = core.lrucache.new({
     type = 'plugin',
 })
 
+local metadata_defaults = {
+    limit_header = "X-RateLimit-Limit",
+    remaining_header = "X-RateLimit-Remaining",
+    reset_header = "X-RateLimit-Reset",
+}
+
+local metadata_schema = {
+    type = "object",
+    properties = {
+        limit_header = {
+            type = "string",
+            default = metadata_defaults.limit_header,
+        },
+        remaining_header = {
+            type = "string",
+            default = metadata_defaults.remaining_header,
+        },
+        reset_header = {
+            type = "string",
+            default = metadata_defaults.reset_header,
+        },
+    },
+}
+
 local schema = {
     type = "object",
     properties = {
@@ -91,7 +115,8 @@ local schema = {
 local schema_copy = core.table.deepcopy(schema)
 
 local _M = {
-    schema = schema
+    schema = schema,
+    metadata_schema = metadata_schema,
 }
 
 
@@ -100,7 +125,12 @@ local function group_conf(conf)
 end
 
 
-function _M.check_schema(conf)
+
+function _M.check_schema(conf, schema_type)
+    if schema_type == core.schema.TYPE_METADATA then
+        return core.schema.check(metadata_schema, conf)
+    end
+
     local ok, err = core.schema.check(schema, conf)
     if not ok then
         return false, err
@@ -250,14 +280,22 @@ function _M.rate_limit(conf, ctx, name, cost)
         delay, remaining, reset = lim:incoming(key, cost)
     end
 
+    local metadata = apisix_plugin.plugin_metadata("limit-count")
+    if metadata then
+        metadata = metadata.value
+    else
+        metadata = metadata_defaults
+    end
+    core.log.info("limit-count plugin-metadata: ", core.json.delay_encode(metadata))
+
     if not delay then
         local err = remaining
         if err == "rejected" then
             -- show count limit header when rejected
             if conf.show_limit_quota_header then
-                core.response.set_header("X-RateLimit-Limit", conf.count,
-                    "X-RateLimit-Remaining", 0,
-                    "X-RateLimit-Reset", reset)
+                core.response.set_header(metadata.limit_header, conf.count,
+                    metadata.remaining_header, 0,
+                    metadata.reset_header, reset)
             end
 
             if conf.rejected_msg then
@@ -274,9 +312,9 @@ function _M.rate_limit(conf, ctx, name, cost)
     end
 
     if conf.show_limit_quota_header then
-        core.response.set_header("X-RateLimit-Limit", conf.count,
-            "X-RateLimit-Remaining", remaining,
-            "X-RateLimit-Reset", reset)
+        core.response.set_header(metadata.limit_header, conf.count,
+            metadata.remaining_header, remaining,
+            metadata.reset_header, reset)
     end
 end
 

t/plugin/limit-count5.t

@@ -137,3 +137,66 @@ passed
 ["GET /hello", "GET /hello", "GET /hello", "GET /hello"]
 --- error_code eval
 [200, 200, 503, 503]
+
+
+
+=== TEST 4: customize rate limit headers by plugin metadata
+--- config
+    location /t {
+        content_by_lua_block {
+            local t = require("lib.test_admin").test
+            local code, body = t('/apisix/admin/routes/1',
+                 ngx.HTTP_PUT,
+                 [[{
+                        "methods": ["GET"],
+                        "plugins": {
+                            "limit-count": {
+                                "count": 10,
+                                "time_window": 60,
+                                "rejected_code": 503,
+                                "key_type": "var",
+                                "key": "remote_addr"
+                            }
+                        },
+                        "upstream": {
+                            "nodes": {
+                                "127.0.0.1:1980": 1
+                            },
+                            "type": "roundrobin"
+                        },
+                        "uri": "/hello"
+                }]]
+                )
+            if code >= 300 then
+                ngx.status = code
+                ngx.say("fail")
+                return
+            end
+            local code, meta_body = t('/apisix/admin/plugin_metadata/limit-count',
+                 ngx.HTTP_PUT,
+                 [[{
+                        "limit_header":"APISIX-RATELIMIT-QUOTA",
+                        "remaining_header":"APISIX-RATELIMIT-REMAINING",
+                        "reset_header":"APISIX-RATELIMIT-RESET"
+                }]]
+                )
+            if code >= 300 then
+                ngx.status = code
+                ngx.say("fail")
+                return
+            end
+            ngx.say("passed")
+        }
+    }
+--- response_body
+passed
+
+
+
+=== TEST 5: check rate limit headers
+--- request
+GET /hello
+--- response_headers_like
+APISIX-RATELIMIT-QUOTA: 10
+APISIX-RATELIMIT-REMAINING: 9
+APISIX-RATELIMIT-RESET: \d+