feat:support limit_size field to the hold_body_chunk method to truncate response_body

apisix/core/response.lua

@@ -153,6 +153,19 @@ function _M.clear_header_as_body_modified()
     ngx.header.etag = nil
 end
 
+local function check_limit_size(ctx, body_buffer_size, limit_size)
+    if not limit_size then
+        return
+    end
+
+    if body_buffer_size > limit_size then
+        ctx._body_buffer = nil
+        ctx._body_hold_flag = true
+    else
+        ctx._body_buffer_size = body_buffer_size
+        ctx._body_hold_flag = false
+    end
+end
 
 -- Hold body chunks and return the final body once all chunks have been read.
 -- Usage:
@@ -164,11 +177,17 @@ end
 --  final_body = transform(final_body)
 --  ngx.arg[1] = final_body
 --  ...
-function _M.hold_body_chunk(ctx, hold_the_copy)
+function _M.hold_body_chunk(ctx, hold_the_copy, limit_size)
+    local body_buffer_size, body_hold_flag
     local body_buffer
     local chunk, eof = arg[1], arg[2]
     if eof then
         body_buffer = ctx._body_buffer
+        body_hold_flag = ctx._body_hold_flag
+        if body_hold_flag then
+            return nil
+        end
+
         if not body_buffer then
             return chunk
         end
@@ -179,17 +198,26 @@ function _M.hold_body_chunk(ctx, hold_the_copy)
     end
 
     if type(chunk) == "string" and chunk ~= "" then
+        body_hold_flag = ctx._body_hold_flag
+        if body_hold_flag then
+            return nil
+        end
+
         body_buffer = ctx._body_buffer
         if not body_buffer then
             body_buffer = {
                 chunk,
                 n = 1
             }
+            ctx._body_buffer_size = #chunk
             ctx._body_buffer = body_buffer
+            check_limit_size(ctx, #chunk, limit_size)
         else
             local n = body_buffer.n + 1
             body_buffer.n = n
             body_buffer[n] = chunk
+            body_buffer_size = ctx._body_buffer_size + #chunk
+            check_limit_size(ctx, body_buffer_size, limit_size)
         end
     end
 

apisix/plugins/limit-count.lua

@@ -90,7 +90,7 @@ local schema = {
         group = {type = "string"},
         key = {type = "string", default = "remote_addr"},
         key_type = {type = "string",
-            enum = {"var", "var_combination"},
+            enum = {"var", "var_combination", "constant"},
             default = "var",
         },
         rejected_code = {
@@ -238,6 +238,8 @@ function _M.access(conf, ctx)
         if n_resolved == 0 then
             key = nil
         end
+    elseif conf.key_type == "constant" then
+        key = conf_key
     else
         key = ctx.var[conf_key]
     end
@@ -248,10 +250,11 @@ function _M.access(conf, ctx)
         key = ctx.var["remote_addr"]
     end
 
+    -- here we add a separator ':' to mark the boundary of the prefix and the key itself
     if not conf.group then
-        key = key .. ctx.conf_type .. ctx.conf_version
+        key = ctx.conf_type .. ctx.conf_version .. ':' .. key
     else
-        key = key .. conf.group
+        key = conf.group .. ':' .. key
     end
 
     core.log.info("limit key: ", key)

docs/en/latest/plugins/limit-count.md

@@ -39,8 +39,8 @@ Limit request rate by a fixed number of requests in a given time window.
 | ------------------- | ------- | --------------------------------------- | ------------- | ------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
 | count               | integer | required                                |               | count > 0                                                                                               | the specified number of requests threshold.                                                                                                                                                                                                                                                                |
 | time_window         | integer | required                    |               | time_window > 0                                                                                         | the time window in seconds before the request count is reset.                                                                                                                                                                                                                                              |
-| key_type      | string  | optional    |   "var"   | ["var", "var_combination"] | the type of key. |
-| key           | string  | optional    |     "remote_addr"    |  | the user specified key to limit the rate. If the `key_type` is "var", the key will be treated as a name of variable. If the `key_type` is "var_combination", the key will be a combination of variables. For example, if we use "$remote_addr $consumer_name" as keys, plugin will be restricted by two keys which are "remote_addr" and "consumer_name". If the value of the key is empty, `remote_addr` will be set as the default key.|
+| key_type      | string  | optional    |   "var"   | ["var", "var_combination", "constant"] | the type of key. |
+| key           | string  | optional    |     "remote_addr"    |  | the user specified key to limit the rate. If the `key_type` is "constant", the key will be treated as a constant. If the `key_type` is "var", the key will be treated as a name of variable. If the `key_type` is "var_combination", the key will be a combination of variables. For example, if we use "$remote_addr $consumer_name" as key, plugin will be restricted by two variables which are "remote_addr" and "consumer_name". If the value of the key is empty, `remote_addr` will be set as the default key.|
 | rejected_code       | integer | optional                                | 503           | [200,...,599]                                                                                           | The HTTP status code returned when the request exceeds the threshold is rejected, default 503.                                                                                                                                                                                                             |
 | rejected_msg       | string | optional                                |            | non-empty                                                                                           | The response body returned when the request exceeds the threshold is rejected.                                                                                                                                                                                                             |
 | policy              | string  | optional                                | "local"       | ["local", "redis", "redis-cluster"]                                                                     | The rate-limiting policies to use for retrieving and incrementing the limits. Available values are `local`(the counters will be stored locally in-memory on the node), `redis`(counters are stored on a Redis server and will be shared across the nodes, usually use it to do the global speed limit), and `redis-cluster` which works the same as `redis` but with redis cluster. |
@@ -110,7 +110,7 @@ You also can complete the above operation through the web interface, first add a
 
 It is possible to share the same limit counter across different routes. For example,
 
-```
+```shell
 curl -i http://127.0.0.1:9080/apisix/admin/services/1 -H 'X-API-KEY: edd1c9f034335f136f87ad84b625c8f1' -X PUT -d '
 {
     "plugins": {
@@ -133,7 +133,7 @@ curl -i http://127.0.0.1:9080/apisix/admin/services/1 -H 'X-API-KEY: edd1c9f0343
 
 Every route which group name is "services_1#1640140620" will share the same count limitation `1` in one minute per remote_addr.
 
-```
+```shell
 $ curl -i http://127.0.0.1:9080/apisix/admin/routes/1 -H 'X-API-KEY: edd1c9f034335f136f87ad84b625c8f1' -X PUT -d '
 {
     "service_id": "1",
@@ -156,6 +156,36 @@ HTTP/1.1 503 ...
 Note that every limit-count configuration of the same group must be the same.
 Therefore, once update the configuration, we also need to update the group name.
 
+It is also possible to share the same limit counter in all requests. For example,
+
+```shell
+curl -i http://127.0.0.1:9080/apisix/admin/services/1 -H 'X-API-KEY: edd1c9f034335f136f87ad84b625c8f1' -X PUT -d '
+{
+    "plugins": {
+        "limit-count": {
+            "count": 1,
+            "time_window": 60,
+            "rejected_code": 503,
+            "key": "remote_addr",
+            "key_type": "constant",
+            "group": "services_1#1640140621"
+        }
+    },
+    "upstream": {
+        "type": "roundrobin",
+        "nodes": {
+            "127.0.0.1:1980": 1
+        }
+    }
+}'
+```
+
+Compared with the previous configuration, we set the `key_type` to `constant`.
+By setting `key_type` to `constant`, we don't evaluate the value of `key` but treat it as a constant.
+
+Now every route which group name is "services_1#1640140621" will share the same count limitation `1` in one minute among all the requests,
+even these requests are from different remote_addr.
+
 If you need a cluster-level precision traffic limit, then we can do it with the redis server. The rate limit of the traffic will be shared between different APISIX nodes to limit the rate of cluster traffic.
 
 Here is the example if we use single `redis` policy:

docs/zh/latest/plugins/limit-count.md

@@ -42,8 +42,8 @@ title: limit-count
 | ------------------- | ------- | --------------------------------- | ------------- | ------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
 | count               | integer | 必须                               |               | count > 0                                                                                               | 指定时间窗口内的请求数量阈值                                                                                                                                                                                                                                                                                                                                                                                          |
 | time_window         | integer | 必须                               |               | time_window > 0                                                                                         | 时间窗口的大小（以秒为单位），超过这个时间就会重置                                                                                                                                                                                                                                                                                                                                                                    |
-| key_type      | string | 可选   |  "var"      | ["var", "var_combination"]                                          | key 的类型 |
-| key           | string  | 可选   |    "remote_addr"    |  | 用来做请求计数的依据。如果 `key_type` 为 "var"，那么 key 会被当作变量名称。如果 `key_type` 为 "var_combination"，那么 key 会当作变量组。比如如果设置 "$remote_addr $consumer_name" 作为 keys，那么插件会同时受 remote_addr 和 consumer_name 两个 key 的约束。如果 key 的值为空，$remote_addr 会被作为默认 key。 |
+| key_type      | string | 可选   |  "var"      | ["var", "var_combination", "constant"]                                          | key 的类型 |
+| key           | string  | 可选   |    "remote_addr"    |  | 用来做请求计数的依据。如果 `key_type` 为 "constant"，那么 key 会被当作常量。如果 `key_type` 为 "var"，那么 key 会被当作变量名称。如果 `key_type` 为 "var_combination"，那么 key 会当作变量组。比如如果设置 "$remote_addr $consumer_name" 作为 key，那么插件会同时受 remote_addr 和 consumer_name 两个变量的约束。如果 key 的值为空，$remote_addr 会被作为默认 key。 |
 | rejected_code       | integer | 可选                               | 503           | [200,...,599]                                                                                           | 当请求超过阈值被拒绝时，返回的 HTTP 状态码                                                                                                                                                                                                                                                                                                                                                                            |
 | rejected_msg       | string | 可选                                |            | 非空                                                                                           | 当请求超过阈值被拒绝时，返回的响应体。                                                                                                                                                                                                             |
 | policy              | string  | 可选                               | "local"       | ["local", "redis", "redis-cluster"]                                                                     | 用于检索和增加限制的速率限制策略。可选的值有：`local`(计数器被以内存方式保存在节点本地，默认选项) 和 `redis`(计数器保存在 Redis 服务节点上，从而可以跨节点共享结果，通常用它来完成全局限速)；以及`redis-cluster`，跟 redis 功能一样，只是使用 redis 集群方式。                                                                                                                                                        |
@@ -115,7 +115,7 @@ curl -i http://127.0.0.1:9080/apisix/admin/routes/1 -H 'X-API-KEY: edd1c9f034335
 
 我们也支持在多个 Route 间共享同一个限流计数器。举个例子，
 
-```
+```shell
 curl -i http://127.0.0.1:9080/apisix/admin/services/1 -H 'X-API-KEY: edd1c9f034335f136f87ad84b625c8f1' -X PUT -d '
 {
     "plugins": {
@@ -138,7 +138,7 @@ curl -i http://127.0.0.1:9080/apisix/admin/services/1 -H 'X-API-KEY: edd1c9f0343
 
 每个配置了 `group` 为 `services_1#1640140620` 的 Route 都将共享同一个每个 IP 地址每分钟只能访问一次的计数器。
 
-```
+```shell
 $ curl -i http://127.0.0.1:9080/apisix/admin/routes/1 -H 'X-API-KEY: edd1c9f034335f136f87ad84b625c8f1' -X PUT -d '
 {
     "service_id": "1",
@@ -161,6 +161,35 @@ HTTP/1.1 503 ...
 注意同一个 group 里面的 limit-count 配置必须一样。
 所以，一旦修改了配置，我们需要更新对应的 group 的值。
 
+我们也支持在所有请求间共享同一个限流计数器。举个例子，
+
+```shell
+curl -i http://127.0.0.1:9080/apisix/admin/services/1 -H 'X-API-KEY: edd1c9f034335f136f87ad84b625c8f1' -X PUT -d '
+{
+    "plugins": {
+        "limit-count": {
+            "count": 1,
+            "time_window": 60,
+            "rejected_code": 503,
+            "key": "remote_addr",
+            "key_type": "constant",
+            "group": "services_1#1640140621"
+        }
+    },
+    "upstream": {
+        "type": "roundrobin",
+        "nodes": {
+            "127.0.0.1:1980": 1
+        }
+    }
+}'
+```
+
+在上面的例子中，我们将 `key_type` 设置为 `constant`。
+通过设置 `key_type` 为 `constant`，`key` 的值将会直接作为常量来处理。
+
+现在每个配置了 `group` 为 `services_1#1640140620` 的 Route 上的所有请求，都将共享同一个每分钟只能访问一次的计数器，即使它们来自不同的 IP 地址。
+
 如果你需要一个集群级别的流量控制，我们可以借助 redis server 来完成。不同的 APISIX 节点之间将共享流量限速结果，实现集群流量限速。
 
 如果启用单 redis 策略，请看下面例子：

rockspec/apisix-master-0.rockspec

@@ -32,7 +32,7 @@ description = {
 
 dependencies = {
     "lua-resty-ctxdump = 0.1-0",
-    "lua-resty-dns-client = 5.2.0",
+    "lua-resty-dns-client = 5.2.3",
     "lua-resty-template = 2.0",
     "lua-resty-etcd = 1.6.0",
     "api7-lua-resty-http = 0.2.0",

t/node/upstream-domain.t

@@ -399,3 +399,51 @@ GET /t
 1980, 1981, 1981
 --- no_error_log
 [error]
+
+
+
+=== TEST 15: set route(with upstream)
+--- 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,
+                [[{
+                    "uri": "/hello",
+                    "upstream": {
+                        "nodes": {
+                            "foo.com.": 0,
+                            "127.0.0.1:1980": 1
+                        },
+                        "type": "roundrobin",
+                        "desc": "new upstream"
+                    },
+                    "service_id": "1"
+                }]]
+                )
+
+            if code >= 300 then
+                ngx.status = code
+            end
+            ngx.say(body)
+        }
+    }
+--- request
+GET /t
+--- response_body
+passed
+--- no_error_log
+[error]
+
+
+
+=== TEST 16: hit routes, parse the domain of upstream node
+--- request
+GET /hello
+--- response_body
+hello world
+--- error_log eval
+qr/dns resolver domain: foo.com. to \d+.\d+.\d+.\d+/
+--- no_error_log
+[error]