From bb5b8674c08e0f1fcb43984dd0a71d25e50f5aa4 Mon Sep 17 00:00:00 2001 From: louxiu Date: Wed, 13 Nov 2024 10:26:52 +0800 Subject: [PATCH] feat: support register custom var prefix --- apisix/core/ctx.lua | 41 +++++++++++--- t/core/ctx2.t | 130 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 164 insertions(+), 7 deletions(-) diff --git a/apisix/core/ctx.lua b/apisix/core/ctx.lua index 36b8788bd02d..88fc0775aa43 100644 --- a/apisix/core/ctx.lua +++ b/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 + local arg_key = sub_str(key, #matched_prefix_name + 1) + val = matched_prefix_getter(t._ctx, arg_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 diff --git a/t/core/ctx2.t b/t/core/ctx2.t index 7782ac9125cd..f5e6c74ecff7 100644 --- a/t/core/ctx2.t +++ b/t/core/ctx2.t @@ -447,3 +447,133 @@ 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 +find ctx.var.a6_count1: 1 +find ctx.var.a6_count2: 1 + +