From c967e8326179c86680c0f34d82ee087765aed19a Mon Sep 17 00:00:00 2001 From: Jun Ouyang Date: Mon, 23 Dec 2024 22:19:51 +0800 Subject: [PATCH] feat: control the proxy_upstream in lua side (#98) Set upstream next enablement of current request to the given string of table argument . Global setting set by proxy_next_upstream will be overwritten. AG-186 --- .github/workflows/tests.yml | 2 +- README.md | 39 ++++ lualib/resty/kong/upstream.lua | 98 ++++++++++ src/ngx_http_lua_kong_common.h | 1 + src/ngx_http_lua_kong_module.c | 44 ++++- src/ngx_http_lua_kong_module.h | 4 + t/013-upstream.t | 323 +++++++++++++++++++++++++++++++++ 7 files changed, 509 insertions(+), 2 deletions(-) create mode 100644 lualib/resty/kong/upstream.lua create mode 100644 t/013-upstream.t diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 2df0719b..72a8b928 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -5,7 +5,7 @@ on: push: env: - KONG_VERSION: master + KONG_VERSION: feat-next-upstream BUILD_ROOT: ${{ github.workspace }}/kong/bazel-bin/build concurrency: diff --git a/README.md b/README.md index b8d84f8d..f7343679 100644 --- a/README.md +++ b/README.md @@ -31,6 +31,8 @@ Table of Contents * [resty.kong.log.set\_log\_level](#restykonglogset_log_level) * [resty.kong.log.get\_log\_level](#restykonglogget_log_level) * [resty.kong.peer_conn.get\_last\_peer\_connection\_cached](#restykongpeer_connget_last_peer_connection_cached) + * [resty.kong.upstream.set\_next\_upstream](#restykongupstreamset_next_upstream) + * [License](#license) Description @@ -576,6 +578,43 @@ balancer_by_lua_block { [Back to TOC](#table-of-contents) + +resty.kong.upstream.set\_next\_upstream +---------------------------------- +**syntax:** *res = resty.kong.upstream.set_next_upstream("http_404")* + +**context:** *rewrite_by_lua*, access_by_lua*, balancer_by_lua** + + +**subsystems:** *http* + +Set upstream next enablement of current request to the given string of table +argument . Global setting set by [`proxy_next_upstream`](http://nginx.org/en/docs/http/ngx_http_proxy_module.html#proxy_next_upstream) will be overwritten. + +The `set_next_upstream` function supports variable length of arguments, and each argument must be one of the following strings (also defined in [`proxy_next_upstream`](http://nginx.org/en/docs/http/ngx_http_proxy_module.html#proxy_next_upstream)): +- `error` +- `timeout` +- `invalid_header` +- `http_500` +- `http_502` +- `http_503` +- `http_504` +- `http_403` +- `http_404` +- `http_429` +- `non_idempotent` +- `off` + +On success, this function returns `nil`. Otherwise throw a string +describing the error will be returned. + +This function can be called multiple times in the same request. Later calls override +previous ones. + +[Back to TOC](#table-of-contents) + + + License ======= diff --git a/lualib/resty/kong/upstream.lua b/lualib/resty/kong/upstream.lua new file mode 100644 index 00000000..18159fcc --- /dev/null +++ b/lualib/resty/kong/upstream.lua @@ -0,0 +1,98 @@ +-- Copyright 2019-2022 Kong Inc. + +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at + +-- http://www.apache.org/licenses/LICENSE-2.0 + +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +local _M = {} + +local ffi = require("ffi") +local base = require("resty.core.base") +local select = select +base.allows_subsystem("http") + +ffi.cdef([[ +int +ngx_http_lua_ffi_set_next_upstream(ngx_http_request_t *r, uint32_t next_upstream, char **err); +const uint32_t ngx_http_lua_kong_next_upstream_mask_error; +const uint32_t ngx_http_lua_kong_next_upstream_mask_timeout; +const uint32_t ngx_http_lua_kong_next_upstream_mask_invalid_header; +const uint32_t ngx_http_lua_kong_next_upstream_mask_http_500; +const uint32_t ngx_http_lua_kong_next_upstream_mask_http_502; +const uint32_t ngx_http_lua_kong_next_upstream_mask_http_503; +const uint32_t ngx_http_lua_kong_next_upstream_mask_http_504; +const uint32_t ngx_http_lua_kong_next_upstream_mask_http_403; +const uint32_t ngx_http_lua_kong_next_upstream_mask_http_404; +const uint32_t ngx_http_lua_kong_next_upstream_mask_http_429; +const uint32_t ngx_http_lua_kong_next_upstream_mask_off; +const uint32_t ngx_http_lua_kong_next_upstream_mask_non_idempotent; +]]) + +local type = type +local C = ffi.C +local get_request = base.get_request +local ffi_str = ffi.string + +local NGX_OK = ngx.OK + +local next_upstream_table = { + error = C.ngx_http_lua_kong_next_upstream_mask_error, + timeout = C.ngx_http_lua_kong_next_upstream_mask_timeout, + invalid_header = C.ngx_http_lua_kong_next_upstream_mask_invalid_header, + http_500 = C.ngx_http_lua_kong_next_upstream_mask_http_500, + http_502 = C.ngx_http_lua_kong_next_upstream_mask_http_502, + http_503 = C.ngx_http_lua_kong_next_upstream_mask_http_503, + http_504 = C.ngx_http_lua_kong_next_upstream_mask_http_504, + http_403 = C.ngx_http_lua_kong_next_upstream_mask_http_403, + http_404 = C.ngx_http_lua_kong_next_upstream_mask_http_404, + http_429 = C.ngx_http_lua_kong_next_upstream_mask_http_429, + off = C.ngx_http_lua_kong_next_upstream_mask_off, + non_idempotent = C.ngx_http_lua_kong_next_upstream_mask_non_idempotent, +} + +function _M.set_next_upstream(...) + local nargs = select("#", ...) + if nargs == 0 then + return "no argument" + end + + local r = get_request() + if not r then + return "no request found" + end + + local arg_table = { ... } + local next_upstream = 0 + for i = 1, nargs do + local v = arg_table[i] + if type(v) ~= "string" then + return "argument #" .. i .. " is not a string" + end + + local next_upstream_value = next_upstream_table[v] + if not next_upstream_value then + return "argument #" .. i .. " is not a valid argument" + end + + next_upstream = bit.bor(next_upstream, next_upstream_value) + end + + local err = ffi.new("char *[1]") + local rc = C.ngx_http_lua_ffi_set_next_upstream(r, next_upstream, err) + + if rc ~= NGX_OK then + return "failed to set upstream next: " .. ffi_str(err[0]) + end + + return nil +end + +return _M diff --git a/src/ngx_http_lua_kong_common.h b/src/ngx_http_lua_kong_common.h index 525f1b84..c6f79d01 100644 --- a/src/ngx_http_lua_kong_common.h +++ b/src/ngx_http_lua_kong_common.h @@ -28,6 +28,7 @@ typedef struct { ngx_lua_kong_ssl_ctx_t ssl_ctx; ngx_str_t grpc_authority; ngx_http_log_handler_pt orig_log_handler; + ngx_uint_t next_upstream; } ngx_http_lua_kong_ctx_t; diff --git a/src/ngx_http_lua_kong_module.c b/src/ngx_http_lua_kong_module.c index eae8c723..629d011f 100644 --- a/src/ngx_http_lua_kong_module.c +++ b/src/ngx_http_lua_kong_module.c @@ -16,7 +16,7 @@ #include "ngx_http_lua_kong_directive.h" - +#include "ngx_http_upstream.h" static ngx_int_t ngx_http_lua_kong_init(ngx_conf_t *cf); static void* ngx_http_lua_kong_create_loc_conf(ngx_conf_t* cf); @@ -158,4 +158,46 @@ ngx_http_lua_kong_merge_loc_conf(ngx_conf_t *cf, void *parent, void *child) return NGX_CONF_OK; } +const ngx_uint_t ngx_http_lua_kong_next_upstream_mask_error = NGX_HTTP_UPSTREAM_FT_ERROR; +const ngx_uint_t ngx_http_lua_kong_next_upstream_mask_timeout = NGX_HTTP_UPSTREAM_FT_TIMEOUT; +const ngx_uint_t ngx_http_lua_kong_next_upstream_mask_invalid_header = NGX_HTTP_UPSTREAM_FT_INVALID_HEADER; +const ngx_uint_t ngx_http_lua_kong_next_upstream_mask_http_500 = NGX_HTTP_UPSTREAM_FT_HTTP_500; +const ngx_uint_t ngx_http_lua_kong_next_upstream_mask_http_502 = NGX_HTTP_UPSTREAM_FT_HTTP_502; +const ngx_uint_t ngx_http_lua_kong_next_upstream_mask_http_503 = NGX_HTTP_UPSTREAM_FT_HTTP_503; +const ngx_uint_t ngx_http_lua_kong_next_upstream_mask_http_504 = NGX_HTTP_UPSTREAM_FT_HTTP_504; +const ngx_uint_t ngx_http_lua_kong_next_upstream_mask_http_403 = NGX_HTTP_UPSTREAM_FT_HTTP_403; +const ngx_uint_t ngx_http_lua_kong_next_upstream_mask_http_404 = NGX_HTTP_UPSTREAM_FT_HTTP_404; +const ngx_uint_t ngx_http_lua_kong_next_upstream_mask_http_429 = NGX_HTTP_UPSTREAM_FT_HTTP_429; +const ngx_uint_t ngx_http_lua_kong_next_upstream_mask_off = NGX_HTTP_UPSTREAM_FT_OFF; +const ngx_uint_t ngx_http_lua_kong_next_upstream_mask_non_idempotent = NGX_HTTP_UPSTREAM_FT_NON_IDEMPOTENT; + +ngx_flag_t +ngx_http_lua_kong_get_next_upstream_mask(ngx_http_request_t *r, + ngx_flag_t upstream_next) +{ + ngx_http_lua_kong_ctx_t *ctx; + + ctx = ngx_http_lua_kong_get_module_ctx(r); + if (ctx == NULL) { + return upstream_next; + } + + if(ctx->next_upstream != 0) { + return ctx->next_upstream; + } + return upstream_next; +} + +int +ngx_http_lua_ffi_set_next_upstream(ngx_http_request_t *r, ngx_uint_t next_upstream, char **err) +{ + ngx_http_lua_kong_ctx_t *ctx; + + ctx = ngx_http_lua_kong_get_module_ctx(r); + if (ctx == NULL) { + return NGX_ERROR; + } + ctx->next_upstream = next_upstream; + return NGX_OK; +} \ No newline at end of file diff --git a/src/ngx_http_lua_kong_module.h b/src/ngx_http_lua_kong_module.h index 528b3318..ba68387b 100644 --- a/src/ngx_http_lua_kong_module.h +++ b/src/ngx_http_lua_kong_module.h @@ -40,4 +40,8 @@ ngx_flag_t ngx_http_lua_kong_ssl_get_http2_alpn_enabled(ngx_ssl_connection_t *ssl, ngx_flag_t enable_http2); +ngx_flag_t +ngx_http_lua_kong_get_next_upstream_mask(ngx_http_request_t *r, + ngx_flag_t upstream_next); + #endif /* _NGX_HTTP_LUA_KONG_MODULE_H_INCLUDED_ */ diff --git a/t/013-upstream.t b/t/013-upstream.t new file mode 100644 index 00000000..270d00f3 --- /dev/null +++ b/t/013-upstream.t @@ -0,0 +1,323 @@ +# vim:set ft= ts=4 sw=4 et fdm=marker: + +use Test::Nginx::Socket::Lua; + +#worker_connections(1014); +#master_on(); +#workers(2); +log_level('info'); + +repeat_each(2); +#repeat_each(1); + +plan tests => repeat_each() * (blocks() * 2); + +#no_diff(); +#no_long_string(); +run_tests(); + +__DATA__ +=== TEST 1: default behavior +--- http_config + upstream balancer { + server 127.0.0.1; + balancer_by_lua_block { + local balancer = require "ngx.balancer" + local host = "127.0.0.1" + local port + ngx.ctx.count = (ngx.ctx.count or 0) + 1 + if ngx.ctx.count == 1 then + port = $TEST_NGINX_RAND_PORT_1 + elseif ngx.ctx.count == 2 then + port = $TEST_NGINX_RAND_PORT_2 + elseif ngx.ctx.count == 3 then + port = $TEST_NGINX_RAND_PORT_3 + else + port = $TEST_NGINX_RAND_PORT_4 + end + ngx.log(ngx.ERR, "balancer_by_lua_block: host: ", host, ", port: ", port, ", count: ", ngx.ctx.count) + local ok, err = balancer.set_current_peer(host, port) + if not ok then + ngx.log(ngx.ERR, "failed to set the current peer: ", err) + return ngx.exit(500) + end + balancer.set_more_tries(4) + } + } + server { + # this is the real entry point + listen $TEST_NGINX_RAND_PORT_1; + location / { + content_by_lua_block{ + ngx.exit(404) + } + } + } + server { + # this is the real entry point + listen $TEST_NGINX_RAND_PORT_2; + location / { + content_by_lua_block{ + ngx.exit(404) + } + } + } + server { + # this is the real entry point + listen $TEST_NGINX_RAND_PORT_3; + location / { + content_by_lua_block{ + ngx.exit(404) + } + } + } + server { + # this is the real entry point + listen $TEST_NGINX_RAND_PORT_4; + location / { + content_by_lua_block{ + ngx.print("this is backend peer $TEST_NGINX_RAND_PORT_4") + } + } + } +--- config + location =/balancer { + proxy_pass http://balancer; + } +--- pipelined_requests eval +["GET /balancer", "GET /balancer"] +--- error_code eval +[404, 404] + +=== TEST 2: proxy_next_upstream directive behavior +--- http_config + upstream balancer { + server 127.0.0.1; + balancer_by_lua_block { + local balancer = require "ngx.balancer" + local host = "127.0.0.1" + local port + ngx.ctx.count = (ngx.ctx.count or 0) + 1 + if ngx.ctx.count == 1 then + port = $TEST_NGINX_RAND_PORT_1 + elseif ngx.ctx.count == 2 then + port = $TEST_NGINX_RAND_PORT_2 + elseif ngx.ctx.count == 3 then + port = $TEST_NGINX_RAND_PORT_3 + else + port = $TEST_NGINX_RAND_PORT_4 + end + ngx.log(ngx.ERR, "balancer_by_lua_block: host: ", host, ", port: ", port, ", count: ", ngx.ctx.count) + local ok, err = balancer.set_current_peer(host, port) + if not ok then + ngx.log(ngx.ERR, "failed to set the current peer: ", err) + return ngx.exit(500) + end + balancer.set_more_tries(4) + } + } + server { + # this is the real entry point + listen $TEST_NGINX_RAND_PORT_1; + location / { + content_by_lua_block{ + ngx.exit(404) + } + } + } + server { + # this is the real entry point + listen $TEST_NGINX_RAND_PORT_2; + location / { + content_by_lua_block{ + ngx.exit(404) + } + } + } + server { + # this is the real entry point + listen $TEST_NGINX_RAND_PORT_3; + location / { + content_by_lua_block{ + ngx.exit(404) + } + } + } + server { + # this is the real entry point + listen $TEST_NGINX_RAND_PORT_4; + location / { + content_by_lua_block{ + ngx.print("this is backend peer $TEST_NGINX_RAND_PORT_4") + } + } + } +--- config + proxy_next_upstream error timeout http_404; + location =/balancer { + proxy_pass http://balancer; + } +--- pipelined_requests eval +["GET /balancer", "GET /balancer"] +--- error_code eval +[200, 200] + +=== TEST 3: lua resty.kong.upstream.set_next_upstream() +--- http_config + upstream balancer { + server 127.0.0.1; + balancer_by_lua_block { + local balancer = require "ngx.balancer" + local host = "127.0.0.1" + local port + ngx.ctx.count = (ngx.ctx.count or 0) + 1 + if ngx.ctx.count == 1 then + port = $TEST_NGINX_RAND_PORT_1 + elseif ngx.ctx.count == 2 then + port = $TEST_NGINX_RAND_PORT_2 + elseif ngx.ctx.count == 3 then + port = $TEST_NGINX_RAND_PORT_3 + else + port = $TEST_NGINX_RAND_PORT_4 + end + ngx.log(ngx.ERR, "balancer_by_lua_block: host: ", host, ", port: ", port, ", count: ", ngx.ctx.count) + local ok, err = balancer.set_current_peer(host, port) + if not ok then + ngx.log(ngx.ERR, "failed to set the current peer: ", err) + return ngx.exit(500) + end + balancer.set_more_tries(4) + } + } + server { + # this is the real entry point + listen $TEST_NGINX_RAND_PORT_1; + location / { + content_by_lua_block{ + ngx.exit(404) + } + } + } + server { + # this is the real entry point + listen $TEST_NGINX_RAND_PORT_2; + location / { + content_by_lua_block{ + ngx.exit(403) + } + } + } + server { + # this is the real entry point + listen $TEST_NGINX_RAND_PORT_3; + location / { + content_by_lua_block{ + ngx.exit(500) + } + } + } + server { + # this is the real entry point + listen $TEST_NGINX_RAND_PORT_4; + location / { + content_by_lua_block{ + ngx.print("this is backend peer $TEST_NGINX_RAND_PORT_4") + } + } + } +--- config + access_by_lua_block { + local upstream = require "resty.kong.upstream" + local err = upstream.set_next_upstream("error", "timeout", "http_500", "http_502", "http_503", "http_504", "http_404", "http_403", "http_429", "non_idempotent") + if err then + ngx.log(ngx.ERR, "failed to set upstream next: ", err) + return ngx.exit(503) + end + } + location =/balancer { + proxy_pass http://balancer; + } +--- pipelined_requests eval +["GET /balancer", "GET /balancer"] +--- error_code eval +[200, 200] + +=== TEST 4: lua resty.kong.upstream.set_next_upstream() with error +--- http_config + upstream balancer { + server 127.0.0.1; + balancer_by_lua_block { + local balancer = require "ngx.balancer" + local host = "127.0.0.1" + local port + ngx.ctx.count = (ngx.ctx.count or 0) + 1 + if ngx.ctx.count == 1 then + port = $TEST_NGINX_RAND_PORT_1 + elseif ngx.ctx.count == 2 then + port = $TEST_NGINX_RAND_PORT_2 + elseif ngx.ctx.count == 3 then + port = $TEST_NGINX_RAND_PORT_3 + else + port = $TEST_NGINX_RAND_PORT_4 + end + ngx.log(ngx.ERR, "balancer_by_lua_block: host: ", host, ", port: ", port, ", count: ", ngx.ctx.count) + local ok, err = balancer.set_current_peer(host, port) + if not ok then + ngx.log(ngx.ERR, "failed to set the current peer: ", err) + return ngx.exit(500) + end + balancer.set_more_tries(4) + } + } + server { + # this is the real entry point + listen $TEST_NGINX_RAND_PORT_1; + location / { + content_by_lua_block{ + ngx.exit(404) + } + } + } + server { + # this is the real entry point + listen $TEST_NGINX_RAND_PORT_2; + location / { + content_by_lua_block{ + ngx.exit(403) + } + } + } + server { + # this is the real entry point + listen $TEST_NGINX_RAND_PORT_3; + location / { + content_by_lua_block{ + ngx.exit(500) + } + } + } + server { + # this is the real entry point + listen $TEST_NGINX_RAND_PORT_4; + location / { + content_by_lua_block{ + ngx.print("this is backend peer $TEST_NGINX_RAND_PORT_4") + } + } + } +--- config + access_by_lua_block { + local upstream = require "resty.kong.upstream" + local err = upstream.set_next_upstream("error", "timeout", "http_500", "http_502", "http_503", "http_504", "http_404", "http_403", "http_429", "non_idempotent", "bala") + if err then + ngx.log(ngx.ERR, "failed to set upstream next: ", err) + return ngx.exit(503) + end + } + location =/balancer { + proxy_pass http://balancer; + } +--- pipelined_requests eval +["GET /balancer", "GET /balancer"] +--- error_code eval +[503, 503] \ No newline at end of file