From 4edbd8c9ea191a4ded46c8b370a4d61fb03d61eb Mon Sep 17 00:00:00 2001
From: levy liu <monkeydluffy6017@gmail.com>
Date: Mon, 20 Mar 2023 14:23:47 +0800
Subject: [PATCH 1/4] feat: support variable when rewrite header in proxy
 rewrite plugin

---
 apisix/core/utils.lua                   |  41 +++++++
 apisix/plugins/proxy-rewrite.lua        |  27 ++++-
 docs/en/latest/plugins/proxy-rewrite.md |   4 +-
 docs/zh/latest/plugins/proxy-rewrite.md |   4 +-
 t/plugin/proxy-rewrite3.t               | 149 ++++++++++++++++++++++++
 5 files changed, 216 insertions(+), 9 deletions(-)

diff --git a/apisix/core/utils.lua b/apisix/core/utils.lua
index 01c8b34c8503..9e1658753b5f 100644
--- a/apisix/core/utils.lua
+++ b/apisix/core/utils.lua
@@ -335,4 +335,45 @@ end
 _M.resolve_var = resolve_var
 
 
+local resolve_var_with_captures
+do
+    local _captures
+    -- escape is not supported very well, like there si a redundant '\' after escape "$1"
+    local pat = [[ (?<! \\) \$ \{? (\d+) \}? ]]
+
+    local function resolve(m)
+        local v = _captures[tonumber(m[1])]
+        if not v then
+            v = ""
+        end
+        return v
+    end
+
+    -- captures is the match result of regex uri in proxy-rewrite plugin
+    function resolve_var_with_captures(tpl, captures)
+        if not tpl then
+            return tpl, nil
+        end
+
+        local from = core_str.find(tpl, "$")
+        if not from then
+            return tpl, nil
+        end
+
+        captures = captures or {}
+
+        _captures = captures
+        local res, _, err = re_gsub(tpl, pat, resolve, "jox")
+        _captures = nil
+        if not res then
+            return nil, err
+        end
+
+        return res, nil
+    end
+end
+-- Resolve {$1, $2, ...} in the given string
+_M.resolve_var_with_captures = resolve_var_with_captures
+
+
 return _M
diff --git a/apisix/plugins/proxy-rewrite.lua b/apisix/plugins/proxy-rewrite.lua
index 0308557ee44c..38ad0c62e413 100644
--- a/apisix/plugins/proxy-rewrite.lua
+++ b/apisix/plugins/proxy-rewrite.lua
@@ -21,6 +21,7 @@ local ipairs      = ipairs
 local ngx         = ngx
 local type        = type
 local re_sub      = ngx.re.sub
+local re_match    = ngx.re.match
 local sub_str     = string.sub
 local str_find    = core.string.find
 
@@ -41,6 +42,10 @@ local lrucache = core.lrucache.new({
     type = "plugin",
 })
 
+core.ctx.register_var("proxy_rewrite_regex_uri_captures", function(ctx)
+    return ctx.proxy_rewrite_regex_uri_captures
+end)
+
 local schema = {
     type = "object",
     properties = {
@@ -257,6 +262,7 @@ do
         return re_sub(s, [[\?]], "%3F", "jo")
     end
 
+
 function _M.rewrite(conf, ctx)
     for _, name in ipairs(upstream_names) do
         if conf[name] then
@@ -278,15 +284,24 @@ function _M.rewrite(conf, ctx)
 
         local uri, _, err = re_sub(upstream_uri, conf.regex_uri[1],
                                    conf.regex_uri[2], "jo")
-        if uri then
-            upstream_uri = uri
-        else
+        if not uri then
             local msg = "failed to substitute the uri " .. ctx.var.uri ..
                         " (" .. conf.regex_uri[1] .. ") with " ..
                         conf.regex_uri[2] .. " : " .. err
             core.log.error(msg)
             return 500, {message = msg}
         end
+
+        local m, err = re_match(upstream_uri, conf.regex_uri[1], "jo")
+        if not m then
+            if err then
+                core.log.error("match error in proxy-rewrite plugin, please check: ", err)
+                return
+            end
+        end
+        ctx.proxy_rewrite_regex_uri_captures = m
+
+        upstream_uri = uri
     end
 
     if not conf.use_real_request_uri_unsafe then
@@ -325,14 +340,16 @@ function _M.rewrite(conf, ctx)
 
         local field_cnt = #hdr_op.add
         for i = 1, field_cnt, 2 do
-            local val = core.utils.resolve_var(hdr_op.add[i + 1], ctx.var)
+            local val = core.utils.resolve_var_with_captures(hdr_op.add[i + 1], ctx.proxy_rewrite_regex_uri_captures)
+            val = core.utils.resolve_var(val, ctx.var)
             local header = hdr_op.add[i]
             core.request.add_header(ctx, header, val)
         end
 
         local field_cnt = #hdr_op.set
         for i = 1, field_cnt, 2 do
-            local val = core.utils.resolve_var(hdr_op.set[i + 1], ctx.var)
+            local val = core.utils.resolve_var_with_captures(hdr_op.set[i + 1], ctx.proxy_rewrite_regex_uri_captures)
+            val = core.utils.resolve_var(val, ctx.var)
             core.request.set_header(ctx, hdr_op.set[i], val)
         end
 
diff --git a/docs/en/latest/plugins/proxy-rewrite.md b/docs/en/latest/plugins/proxy-rewrite.md
index bedd246d976e..0e18d9c5d79c 100644
--- a/docs/en/latest/plugins/proxy-rewrite.md
+++ b/docs/en/latest/plugins/proxy-rewrite.md
@@ -41,8 +41,8 @@ The `proxy-rewrite` Plugin rewrites Upstream proxy information such as `scheme`,
 | regex_uri                   | array[string] | False    |         |                                                                                                                                        | New upstream forwarding address. Regular expressions can be used to match the URL from client. If it matches, the URL template is forwarded to the Upstream otherwise, the URL from the client is forwarded. When both `uri` and `regex_uri` are configured, `uri` is used first. For example, `[" ^/iresty/(.*)/(.*)/(.*)", "/$1-$2-$3"]`. Here, the first element is the regular expression to match and the second element is the URL template forwarded to the Upstream. |
 | host                        | string        | False    |         |                                                                                                                                        | New Upstream host address.                                                                                                                                                                                                                                                                                                                                                                                                                                                   |
 | headers                     | object        | False    |         |                                                                                                                                   |                   |
-| headers.add     | object   | false     |        |                 | Append the new headers. The format is `{"name: value",...}`. The values in the header can contain Nginx variables like $remote_addr and $balancer_ip.                                                                                              |
-| headers.set     | object  | false     |        |                 | Overwrite the headers. If the header does not exist, it will be added. The format is  `{"name": "value", ...}`. The values in the header can contain Nginx variables like $remote_addr and $balancer_ip.                                                                                                |
+| headers.add     | object   | false     |        |                 | Append the new headers. The format is `{"name: value",...}`. The values in the header can contain Nginx variables like `$remote_addr` and `$balancer_ip`. It also supports referencing the match result of `regex_uri` as a variable like `$1-$2-$3`.                                                                                              |
+| headers.set     | object  | false     |        |                 | Overwrite the headers. If the header does not exist, it will be added. The format is  `{"name": "value", ...}`. The values in the header can contain Nginx variables like `$remote_addr` and `$balancer_ip`. It also supports referencing the match result of `regex_uri` as a variable like `$1-$2-$3`.                                                                                        |
 | headers.remove  | array   | false     |        |                 | Remove the headers. The format is `["name", ...]`.
 | use_real_request_uri_unsafe | boolean       | False    | false   |                                                                                                                                        | Use real_request_uri (original $request_uri in nginx) to bypass URI normalization. **Enabling this is considered unsafe as it bypasses all URI normalization steps**.                                                                                                                                                                                                                                                                                                     |
 
diff --git a/docs/zh/latest/plugins/proxy-rewrite.md b/docs/zh/latest/plugins/proxy-rewrite.md
index 881be42ace1d..ff6feaffded3 100644
--- a/docs/zh/latest/plugins/proxy-rewrite.md
+++ b/docs/zh/latest/plugins/proxy-rewrite.md
@@ -41,8 +41,8 @@ description: 本文介绍了关于 Apache APISIX `proxy-rewrite` 插件的基本
 | regex_uri | array[string] | 否    |         |                                                                                                                                        | 转发到上游的新 `uri` 地址。使用正则表达式匹配来自客户端的 `uri`，如果匹配成功，则使用模板替换转发到上游的 `uri`，如果没有匹配成功，则将客户端请求的 `uri` 转发至上游。当同时配置 `uri` 和 `regex_uri` 属性时，优先使用 `uri`。例如：["^/iresty/(.*)/(.*)/(.*)","/$1-$2-$3"] 第一个元素代表匹配来自客户端请求的 `uri` 正则表达式，第二个元素代表匹配成功后转发到上游的 `uri` 模板。但是目前 APISIX 仅支持一个 `regex_uri`，所以 `regex_uri` 数组的长度是 `2`。 |
 | host      | string        | 否    |         |                   | 转发到上游的新 `host` 地址，例如：`iresty.com`。|
 | headers   | object        | 否    |         |                   |   |
-| headers.add     | object   | 否     |        |                 | 添加新的请求头，如果头已经存在，会追加到末尾。格式为 `{"name: value", ...}`。这个值能够以 `$var` 的格式包含 NGINX 变量，比如 `$remote_addr $balancer_ip`。                                                                                              |
-| headers.set     | object  | 否     |        |                 | 改写请求头，如果请求头不存在，则会添加这个请求头。格式为 `{"name": "value", ...}`。这个值能够以 `$var` 的格式包含 NGINX 变量，比如 `$remote_addr $balancer_ip`。                                                                                                |
+| headers.add     | object   | 否     |        |                 | 添加新的请求头，如果头已经存在，会追加到末尾。格式为 `{"name: value", ...}`。这个值能够以 `$var` 的格式包含 NGINX 变量，比如 `$remote_addr $balancer_ip`。也支持以变量的形式引用 `regex_uri` 的匹配结果，比如 `$1-$2-$3`。                                                                                              |
+| headers.set     | object  | 否     |        |                 | 改写请求头，如果请求头不存在，则会添加这个请求头。格式为 `{"name": "value", ...}`。这个值能够以 `$var` 的格式包含 NGINX 变量，比如 `$remote_addr $balancer_ip`。也支持以变量的形式引用 `regex_uri` 的匹配结果，比如 `$1-$2-$3`。                                                                                           |
 | headers.remove  | array   | 否     |        |                 | 移除响应头。格式为 `["name", ...]`。
 
 ## Header 优先级
diff --git a/t/plugin/proxy-rewrite3.t b/t/plugin/proxy-rewrite3.t
index eaea849cd432..45bcd6f5c918 100644
--- a/t/plugin/proxy-rewrite3.t
+++ b/t/plugin/proxy-rewrite3.t
@@ -622,3 +622,152 @@ GET /test/plugin/proxy/rewrite HTTP/1.1
     }
 --- response_body
 /plugin_proxy_rewrite?a=c
+
+
+
+=== TEST 27: use variables in headers when captured by regex_uri
+--- 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": "/test/*",
+                      "plugins": {
+                        "proxy-rewrite": {
+                            "regex_uri": ["^/test/(.*)/(.*)/(.*)", "/echo"],
+                            "headers": {
+                                "add": {
+                                    "X-Request-ID": "$1/$2/$3"
+                                }
+                            }
+                        }
+                      },
+                      "upstream": {
+                          "nodes": {
+                              "127.0.0.1:1980": 1
+                          },
+                          "type": "roundrobin"
+                      }
+                 }]]
+                 )
+
+            if code >= 300 then
+                ngx.status = code
+            end
+            ngx.say(body)
+        }
+    }
+--- request
+GET /t
+--- response_body
+passed
+
+
+
+=== TEST 28: hit
+--- request
+GET /test/plugin/proxy/rewrite HTTP/1.1
+--- response_headers
+X-Request-ID: plugin/proxy/rewrite
+
+
+
+=== TEST 29: use variables in header when not matched regex_uri
+--- 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": "/echo*",
+                      "plugins": {
+                        "proxy-rewrite": {
+                            "regex_uri": ["^/test/(.*)/(.*)/(.*)", "/echo"],
+                            "headers": {
+                                "add": {
+                                    "X-Request-ID": "$1/$2/$3"
+                                }
+                            }
+                        }
+                      },
+                      "upstream": {
+                          "nodes": {
+                              "127.0.0.1:1980": 1
+                          },
+                          "type": "roundrobin"
+                      }
+                 }]]
+                 )
+
+            if code >= 300 then
+                ngx.status = code
+            end
+            ngx.say(body)
+        }
+    }
+--- request
+GET /t
+--- response_body
+passed
+
+
+
+=== TEST 30: hit
+--- request
+GET /echo HTTP/1.1
+--- more_headers
+X-Foo: Foo
+--- response_headers
+X-Foo: Foo
+
+
+
+=== TEST 31: use variables in headers when captured by regex_uri
+--- 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": "/test/*",
+                      "plugins": {
+                        "proxy-rewrite": {
+                            "regex_uri": ["^/test/(not_matched)?.*", "/echo"],
+                            "headers": {
+                                "add": {
+                                    "X-Request-ID": "test1/$1/$2/test2"
+                                }
+                            }
+                        }
+                      },
+                      "upstream": {
+                          "nodes": {
+                              "127.0.0.1:1980": 1
+                          },
+                          "type": "roundrobin"
+                      }
+                 }]]
+                 )
+
+            if code >= 300 then
+                ngx.status = code
+            end
+            ngx.say(body)
+        }
+    }
+--- request
+GET /t
+--- response_body
+passed
+
+
+
+=== TEST 32: hit
+--- request
+GET /test/plugin/proxy/rewrite HTTP/1.1
+--- response_headers
+X-Request-ID: test1///test2

From a50cd9c4ca2acafab365a6436227dabff9f12191 Mon Sep 17 00:00:00 2001
From: levy liu <monkeydluffy6017@gmail.com>
Date: Mon, 20 Mar 2023 14:35:59 +0800
Subject: [PATCH 2/4] style: minor tweaks

---
 apisix/plugins/proxy-rewrite.lua | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/apisix/plugins/proxy-rewrite.lua b/apisix/plugins/proxy-rewrite.lua
index 38ad0c62e413..7a2c4045c7ad 100644
--- a/apisix/plugins/proxy-rewrite.lua
+++ b/apisix/plugins/proxy-rewrite.lua
@@ -340,7 +340,8 @@ function _M.rewrite(conf, ctx)
 
         local field_cnt = #hdr_op.add
         for i = 1, field_cnt, 2 do
-            local val = core.utils.resolve_var_with_captures(hdr_op.add[i + 1], ctx.proxy_rewrite_regex_uri_captures)
+            local val = core.utils.resolve_var_with_captures(hdr_op.add[i + 1],
+                                            ctx.proxy_rewrite_regex_uri_captures)
             val = core.utils.resolve_var(val, ctx.var)
             local header = hdr_op.add[i]
             core.request.add_header(ctx, header, val)
@@ -348,7 +349,8 @@ function _M.rewrite(conf, ctx)
 
         local field_cnt = #hdr_op.set
         for i = 1, field_cnt, 2 do
-            local val = core.utils.resolve_var_with_captures(hdr_op.set[i + 1], ctx.proxy_rewrite_regex_uri_captures)
+            local val = core.utils.resolve_var_with_captures(hdr_op.set[i + 1],
+                                            ctx.proxy_rewrite_regex_uri_captures)
             val = core.utils.resolve_var(val, ctx.var)
             core.request.set_header(ctx, hdr_op.set[i], val)
         end

From 79f924199f6b04aa1797a15246810370ecdd72b3 Mon Sep 17 00:00:00 2001
From: levy liu <monkeydluffy6017@gmail.com>
Date: Thu, 23 Mar 2023 00:02:21 +0800
Subject: [PATCH 3/4] refactor: minor tweaks

---
 apisix/core/utils.lua            | 2 +-
 apisix/plugins/proxy-rewrite.lua | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/apisix/core/utils.lua b/apisix/core/utils.lua
index 9e1658753b5f..70531e7c1fa9 100644
--- a/apisix/core/utils.lua
+++ b/apisix/core/utils.lua
@@ -338,7 +338,7 @@ _M.resolve_var = resolve_var
 local resolve_var_with_captures
 do
     local _captures
-    -- escape is not supported very well, like there si a redundant '\' after escape "$1"
+    -- escape is not supported very well, like there is a redundant '\' after escape "$1"
     local pat = [[ (?<! \\) \$ \{? (\d+) \}? ]]
 
     local function resolve(m)
diff --git a/apisix/plugins/proxy-rewrite.lua b/apisix/plugins/proxy-rewrite.lua
index 7a2c4045c7ad..0fb49842c9d4 100644
--- a/apisix/plugins/proxy-rewrite.lua
+++ b/apisix/plugins/proxy-rewrite.lua
@@ -296,7 +296,7 @@ function _M.rewrite(conf, ctx)
         if not m then
             if err then
                 core.log.error("match error in proxy-rewrite plugin, please check: ", err)
-                return
+                return 500
             end
         end
         ctx.proxy_rewrite_regex_uri_captures = m

From 652873970e69d3631dfcf866de934bdbf8365ecd Mon Sep 17 00:00:00 2001
From: levy liu <monkeydluffy6017@gmail.com>
Date: Thu, 23 Mar 2023 09:41:01 +0800
Subject: [PATCH 4/4] refactor: minor tweaks

---
 apisix/plugins/proxy-rewrite.lua | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/apisix/plugins/proxy-rewrite.lua b/apisix/plugins/proxy-rewrite.lua
index 0fb49842c9d4..8b1df5aaa876 100644
--- a/apisix/plugins/proxy-rewrite.lua
+++ b/apisix/plugins/proxy-rewrite.lua
@@ -293,11 +293,9 @@ function _M.rewrite(conf, ctx)
         end
 
         local m, err = re_match(upstream_uri, conf.regex_uri[1], "jo")
-        if not m then
-            if err then
-                core.log.error("match error in proxy-rewrite plugin, please check: ", err)
-                return 500
-            end
+        if not m and err then
+            core.log.error("match error in proxy-rewrite plugin, please check: ", err)
+            return 500
         end
         ctx.proxy_rewrite_regex_uri_captures = m