feat: support register custom var prefix

apisix/core/ctx.lua

@@ -204,6 +204,12 @@ do
         var_x_forwarded_host       = true,
     }
 
+    local no_cacheable_var_prefix_names = {
+    }
+
+    local apisix_var_prefix_names = {
+    }
+
     -- sort in alphabetical
     local apisix_var_names = {
         balancer_ip = true,
@@ -232,6 +238,7 @@ do
             end
 
             local val
+            local matched_prefix_name
             local method = var_methods[key]
             if method then
                 val = method()
@@ -302,13 +309,26 @@ do
                         -- the getter is registered by ctx.register_var
                         val = getter(ctx)
                     end
-
                 else
-                    val = get_var(key, t._request)
+                    local matched_prefix_getter = nil
+                    for prefix_name, prefix_getter in pairs(apisix_var_prefix_names) do
+                        if core_str.has_prefix(key, prefix_name) then
+                            matched_prefix_name = prefix_name
+                            matched_prefix_getter = prefix_getter
+                            break
+                        end
+                    end
+                    if matched_prefix_getter then
+                        key = sub_str(key, #matched_prefix_name + 1)
+                        val = matched_prefix_getter(t._ctx, key)
+                    else
+                        val = get_var(key, t._request)
+                    end
                 end
             end
 
-            if val ~= nil and not no_cacheable_var_names[key] then
+            if val ~= nil and not no_cacheable_var_names[key] and
+               not no_cacheable_var_prefix_names[matched_prefix_name] then
                 t._cache[key] = val
             end
 
@@ -348,16 +368,23 @@ do
 --
 -- We support the options below in the `opts`:
 -- * no_cacheable: if the result of getter is cacheable or not. Default to `false`.
+-- * prefix: if the variable is a prefix variable or not. Default to `false`.
 function _M.register_var(name, getter, opts)
     if type(getter) ~= "function" then
         error("the getter of registered var should be a function")
     end
 
-    apisix_var_names[name] = getter
-
-    if opts then
+    if opts and opts.prefix then
+        apisix_var_prefix_names[name] = getter
         if opts.no_cacheable then
-            no_cacheable_var_names[name] = true
+            no_cacheable_var_prefix_names[name] = true
+        end
+    else
+        apisix_var_names[name] = getter
+        if opts then
+            if opts.no_cacheable then
+                no_cacheable_var_names[name] = true
+            end
         end
     end
 end

t/core/ctx2.t

@@ -447,3 +447,130 @@ find ctx.var.a6_labels_zone: Singapore
 --- response_body
 find ctx.var.a6_count: 1
 find ctx.var.a6_count: 2
+
+
+
+=== TEST 20: register custom variable with no cacheable
+--- 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,
+                 [=[{
+                        "upstream": {
+                            "nodes": {
+                                "127.0.0.1:1980": 1
+                            },
+                            "type": "roundrobin"
+                        },
+                        "plugins": {
+                            "serverless-pre-function": {
+                                "phase": "rewrite",
+                                "functions" : ["return function(conf, ctx) ngx.say('find ctx.var.a6_count1: ', ctx.var.a6_count1) ngx.say('find ctx.var.a6_count2: ', ctx.var.a6_count2) end"]
+                            },
+                            "serverless-post-function": {
+                                "phase": "rewrite",
+                                "functions" : ["return function(conf, ctx) ngx.say('find ctx.var.a6_count1: ', ctx.var.a6_count1) ngx.say('find ctx.var.a6_count2: ', ctx.var.a6_count2) end"]
+                            }
+                        },
+                        "uri": "/hello"
+                }]=]
+                )
+
+            if code >= 300 then
+                ngx.status = code
+            end
+            ngx.say(body)
+        }
+    }
+
+
+
+=== TEST 21: hit
+--- config
+    location /t {
+        content_by_lua_block {
+            local http = require "resty.http"
+            local core = require "apisix.core"
+            core.ctx.register_var("a6_", function(ctx, key)
+                if not ctx[key] then
+                    ctx[key] = 0
+                end
+                ctx[key] = ctx[key] + 1
+                return ctx[key]
+            end, {no_cacheable = true, prefix = true})
+            local uri = "http://127.0.0.1:" .. ngx.var.server_port .. "/hello"
+            local httpc = http.new()
+            local res = assert(httpc:request_uri(uri))
+            ngx.print(res.body)
+        }
+    }
+--- response_body
+find ctx.var.a6_count1: 1
+find ctx.var.a6_count2: 1
+find ctx.var.a6_count1: 2
+find ctx.var.a6_count2: 2
+
+
+
+=== TEST 22: register custom variable with cacheable
+--- 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,
+                 [=[{
+                        "upstream": {
+                            "nodes": {
+                                "127.0.0.1:1980": 1
+                            },
+                            "type": "roundrobin"
+                        },
+                        "plugins": {
+                            "serverless-pre-function": {
+                                "phase": "rewrite",
+                                "functions" : ["return function(conf, ctx) ngx.say('find ctx.var.a6_count1: ', ctx.var.a6_count1) ngx.say('find ctx.var.a6_count2: ', ctx.var.a6_count2) end"]
+                            },
+                            "serverless-post-function": {
+                                "phase": "rewrite",
+                                "functions" : ["return function(conf, ctx) ngx.say('find ctx.var.a6_count1: ', ctx.var.a6_count1) ngx.say('find ctx.var.a6_count2: ', ctx.var.a6_count2) end"]
+                            }
+                        },
+                        "uri": "/hello"
+                }]=]
+                )
+
+            if code >= 300 then
+                ngx.status = code
+            end
+            ngx.say(body)
+        }
+    }
+
+
+
+=== TEST 23: hit
+--- config
+    location /t {
+        content_by_lua_block {
+            local http = require "resty.http"
+            local core = require "apisix.core"
+            core.ctx.register_var("a6_", function(ctx, key)
+                if not ctx[key] then
+                    ctx[key] = 0
+                end
+                ctx[key] = ctx[key] + 1
+                return ctx[key]
+            end, {no_cacheable = false, prefix = true})
+            local uri = "http://127.0.0.1:" .. ngx.var.server_port .. "/hello"
+            local httpc = http.new()
+            local res = assert(httpc:request_uri(uri))
+            ngx.print(res.body)
+        }
+    }
+--- response_body
+find ctx.var.a6_count1: 1
+find ctx.var.a6_count2: 1
+