From 6ec08c77a57ff35173dc5491870574281e41b98a Mon Sep 17 00:00:00 2001
From: spacewander <spacewanderlzx@gmail.com>
Date: Mon, 7 Dec 2020 19:42:58 +0800
Subject: [PATCH] feat(fault-injection): support Nginx variable in abort.body

Fix #2827
---
 apisix/core/utils.lua                |  41 +++++++++++
 apisix/plugins/fault-injection.lua   |  16 ++--
 doc/plugins/fault-injection.md       |   6 +-
 doc/zh-cn/plugins/fault-injection.md |   6 +-
 t/core/utils.t                       |  40 ++++++++++
 t/plugin/fault-injection.t           | 106 +++++++++++++++++++++++++++
 6 files changed, 198 insertions(+), 17 deletions(-)

diff --git a/apisix/core/utils.lua b/apisix/core/utils.lua
index a7cad379e629..041bd7e8be46 100644
--- a/apisix/core/utils.lua
+++ b/apisix/core/utils.lua
@@ -14,6 +14,7 @@
 -- See the License for the specific language governing permissions and
 -- limitations under the License.
 --
+local core_str = require("apisix.core.string")
 local table    = require("apisix.core.table")
 local ngx_re   = require("ngx.re")
 local resolver = require("resty.dns.resolver")
@@ -25,6 +26,8 @@ local math     = math
 local sub_str  = string.sub
 local str_byte = string.byte
 local tonumber = tonumber
+local tostring = tostring
+local re_gsub  = ngx.re.gsub
 local type     = type
 local C        = ffi.C
 local ffi_string = ffi.string
@@ -219,4 +222,42 @@ end
 _M.sleep = sleep
 
 
+local resolve_var
+do
+    local _ctx
+    local pat = [[(?<!\\)\$(\w+)]]
+
+    local function resolve(m)
+        local v = _ctx[m[1]]
+        if v == nil then
+            return ""
+        end
+        return tostring(v)
+    end
+
+    function resolve_var(tpl, ctx)
+        if not tpl then
+            return tpl
+        end
+
+        local from = core_str.find(tpl, "$")
+        if not from then
+            return tpl
+        end
+
+        -- avoid creating temporary function
+        _ctx = ctx
+        local res, _, err = re_gsub(tpl, pat, resolve, "jo")
+        _ctx = nil
+        if not res then
+            return nil, err
+        end
+
+        return res
+    end
+end
+-- Resolve ngx.var in the given string
+_M.resolve_var = resolve_var
+
+
 return _M
diff --git a/apisix/plugins/fault-injection.lua b/apisix/plugins/fault-injection.lua
index aace3e249941..df4ab6abb1fa 100644
--- a/apisix/plugins/fault-injection.lua
+++ b/apisix/plugins/fault-injection.lua
@@ -32,7 +32,7 @@ local schema = {
                 body = {type = "string", minLength = 0},
                 percentage = {type = "integer", minimum = 0, maximum = 100}
             },
-            minProperties = 1,
+            required = {"http_status"},
         },
         delay = {
             type = "object",
@@ -40,7 +40,7 @@ local schema = {
                 duration = {type = "number", minimum = 0},
                 percentage = {type = "integer", minimum = 0, maximum = 100}
             },
-            minProperties = 1,
+            required = {"duration"},
         }
     },
     minProperties = 1,
@@ -77,18 +77,12 @@ end
 function _M.rewrite(conf, ctx)
     core.log.info("plugin rewrite phase, conf: ", core.json.delay_encode(conf))
 
-    if conf.delay
-       and conf.delay.duration ~= nil
-       and sample_hit(conf.delay.percentage)
-    then
+    if conf.delay and sample_hit(conf.delay.percentage) then
         sleep(conf.delay.duration)
     end
 
-    if conf.abort
-       and conf.abort.http_status ~= nil
-       and sample_hit(conf.abort.percentage)
-    then
-        return conf.abort.http_status, conf.abort.body
+    if conf.abort and sample_hit(conf.abort.percentage) then
+        return conf.abort.http_status, core.utils.resolve_var(conf.abort.body, ctx.var)
     end
 end
 
diff --git a/doc/plugins/fault-injection.md b/doc/plugins/fault-injection.md
index ff09cfc07b78..f2f4039919fe 100644
--- a/doc/plugins/fault-injection.md
+++ b/doc/plugins/fault-injection.md
@@ -27,10 +27,10 @@ Fault injection plugin, this plugin can be used with other plugins and will be e
 
 | Name              | Type    | Requirement | Default | Valid      | Description                                      |
 | ----------------- | ------- | ----------- | ------- | ---------- | ------------------------------------------------ |
-| abort.http_status | integer | optional    |         | [200, ...] | user-specified http code returned to the client. |
-| abort.body        | string  | optional    |         |            | response data returned to the client.            |
+| abort.http_status | integer | required    |         | [200, ...] | user-specified http code returned to the client. |
+| abort.body        | string  | optional    |         |            | response data returned to the client. Nginx varialbe can be used inside, like `client addr: $remote_addr\n`           |
 | abort.percentage  | integer | optional    |         | [0, 100]   | percentage of requests to be aborted.            |
-| delay.duration    | number  | optional    |         |            | delay time (can be decimal).                     |
+| delay.duration    | number  | required    |         |            | delay time (can be decimal).                     |
 | delay.percentage  | integer | optional    |         | [0, 100]   | percentage of requests to be delayed.            |
 
 Note: One of `abort` and `delay` must be specified.
diff --git a/doc/zh-cn/plugins/fault-injection.md b/doc/zh-cn/plugins/fault-injection.md
index ad404922d328..45e5936279d3 100644
--- a/doc/zh-cn/plugins/fault-injection.md
+++ b/doc/zh-cn/plugins/fault-injection.md
@@ -27,10 +27,10 @@
 
 | 名称              | 类型    | 必选项 | 默认值 | 有效值     | 描述                       |
 | ----------------- | ------- | ------ | ------ | ---------- | -------------------------- |
-| abort.http_status | integer | 可选   |        | [200, ...] | 返回给客户端的 http 状态码 |
-| abort.body        | string  | 可选   |        |            | 返回给客户端的响应数据     |
+| abort.http_status | integer | 必需   |        | [200, ...] | 返回给客户端的 http 状态码 |
+| abort.body        | string  | 可选   |        |            | 返回给客户端的响应数据。支持使用 Nginx 变量，如 `client addr: $remote_addr\n`|
 | abort.percentage  | integer | 可选   |        | [0, 100]   | 将被中断的请求占比         |
-| delay.duration    | number  | 可选   |        |            | 延迟时间，可以指定小数     |
+| delay.duration    | number  | 必需   |        |            | 延迟时间，可以指定小数     |
 | delay.percentage  | integer | 可选   |        | [0, 100]   | 将被延迟的请求占比         |
 
 注：参数 abort 和 delay 至少要存在一个。
diff --git a/t/core/utils.t b/t/core/utils.t
index f4904da789d7..4006b6a71d7a 100644
--- a/t/core/utils.t
+++ b/t/core/utils.t
@@ -199,3 +199,43 @@ received: \nreceived: hello world
 close: 1 nil}
 --- no_error_log
 [error]
+
+
+
+=== TEST 6: resolve_var
+--- config
+    location /t {
+        content_by_lua_block {
+            local resolve_var = require("apisix.core.utils").resolve_var
+            local cases = {
+                "",
+                "xx",
+                "$me",
+                "$me run",
+                "talk with $me",
+                "tell $me to",
+                "$you and $me",
+                "$eva and $me",
+                "$you and \\$me",
+            }
+            local ctx = {
+                you = "John",
+                me = "David",
+            }
+            for _, case in ipairs(cases) do
+                ngx.say("res:", resolve_var(case, ctx))
+            end
+        }
+    }
+--- request
+GET /t
+--- response_body
+res:
+res:xx
+res:David
+res:David run
+res:talk with David
+res:tell David to
+res:John and David
+res: and David
+res:John and \$me
diff --git a/t/plugin/fault-injection.t b/t/plugin/fault-injection.t
index c427714e09d8..439381140b20 100644
--- a/t/plugin/fault-injection.t
+++ b/t/plugin/fault-injection.t
@@ -635,3 +635,109 @@ GET /hello HTTP/1.1
 hello1 world
 --- no_error_log
 [error]
+
+
+
+=== TEST 20: set route(body with var)
+--- 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,
+                    [[{
+                    "plugins": {
+                        "fault-injection": {
+                            "abort": {
+                                "http_status": 200,
+                                "body": "client addr: $remote_addr\n"
+                            }
+                        },
+                        "proxy-rewrite": {
+                            "uri": "/hello"
+                            }
+                    },
+                    "upstream": {
+                        "nodes": {
+                            "127.0.0.1:1980": 1
+                        },
+                        "type": "roundrobin"
+                    },
+                    "uri": "/hello"
+                   }]]
+                   )
+
+               if code >= 300 then
+                   ngx.status = code
+               end
+               ngx.say(body)
+           }
+       }
+--- request
+GET /t
+--- response_body
+passed
+--- no_error_log
+[error]
+
+
+
+=== TEST 21: hit route(body with var)
+--- request
+GET /hello
+--- response_body
+client addr: 127.0.0.1
+--- no_error_log
+[error]
+
+
+
+=== TEST 22: set route(abort without body)
+--- 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,
+                    [[{
+                    "plugins": {
+                        "fault-injection": {
+                            "abort": {
+                                "http_status": 200
+                            }
+                        },
+                        "proxy-rewrite": {
+                            "uri": "/hello"
+                            }
+                    },
+                    "upstream": {
+                        "nodes": {
+                            "127.0.0.1:1980": 1
+                        },
+                        "type": "roundrobin"
+                    },
+                    "uri": "/hello"
+                   }]]
+                   )
+
+               if code >= 300 then
+                   ngx.status = code
+               end
+               ngx.say(body)
+           }
+       }
+--- request
+GET /t
+--- response_body
+passed
+--- no_error_log
+[error]
+
+
+
+=== TEST 23: hit route(abort without body)
+--- request
+GET /hello
+--- response_body
+--- no_error_log
+[error]