From 36e374bc73c23d51d2af841c5d6c69c23956f313 Mon Sep 17 00:00:00 2001 From: Jun Ouyang Date: Tue, 17 Dec 2024 15:57:38 +0800 Subject: [PATCH] feat: control the proxy_upstream in lua side --- lualib/resty/kong/upstream.lua | 82 ++++++++++++++++++++++++++++++++++ src/ngx_http_lua_kong_common.h | 1 + src/ngx_http_lua_kong_module.c | 71 +++++++++++++++++++++++++++++ 3 files changed, 154 insertions(+) create mode 100644 lualib/resty/kong/upstream.lua diff --git a/lualib/resty/kong/upstream.lua b/lualib/resty/kong/upstream.lua new file mode 100644 index 00000000..91fab174 --- /dev/null +++ b/lualib/resty/kong/upstream.lua @@ -0,0 +1,82 @@ +-- 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") +base.allows_subsystem("http") + +ffi.cdef([[ +int +ngx_http_lua_ffi_set_upstream_next(ngx_http_request_t *r, uint next_upstream, char **err); +]]) + +local type = type +local error = error +local tostring = tostring +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 = 0x00000002, + timeout = 0x00000004, + invalid_header = 0x00000008, + http_500 = 0x00000010, + http_502 = 0x00000020, + http_503 = 0x00000040, + http_504 = 0x00000080, + http_403 = 0x00000100, + http_404 = 0x00000200, + http_429 = 0x00000400, + updating = 0x00000800, + off = 0x00001000, +} + +function _M.set_upstream_next(...) + local nargs = select("#", ...) + if nargs == 0 then + error("no argument") + end + + local r = get_request() + if not r then + error("no request found") + end + + local next_upstream_table = { ... } + local next_upstream = 0 + for i = 1, nargs do + local v = next_upstream_table[i] + if type(v) ~= "number" then + error("argument #" .. i .. " is not a valid argument") + end + + next_upstream = bit.bor(next_upstream, v) + end + + local err = ffi.new("char *[1]") + local rc = C.ngx_http_lua_ffi_set_upstream_next(r, next_upstream, err) + + if rc ~= NGX_OK then + error("failed to set upstream next: " .. tostring(ffi_str(err[0]))) + end + + return true +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..42951878 100644 --- a/src/ngx_http_lua_kong_module.c +++ b/src/ngx_http_lua_kong_module.c @@ -158,4 +158,75 @@ ngx_http_lua_kong_merge_loc_conf(ngx_conf_t *cf, void *parent, void *child) return NGX_CONF_OK; } +#define NGX_HTTP_UPSTREAM_FT_ERROR 0x00000002 +#define NGX_HTTP_UPSTREAM_FT_TIMEOUT 0x00000004 +#define NGX_HTTP_UPSTREAM_FT_INVALID_HEADER 0x00000008 +#define NGX_HTTP_UPSTREAM_FT_HTTP_500 0x00000010 +#define NGX_HTTP_UPSTREAM_FT_HTTP_502 0x00000020 +#define NGX_HTTP_UPSTREAM_FT_HTTP_503 0x00000040 +#define NGX_HTTP_UPSTREAM_FT_HTTP_504 0x00000080 +#define NGX_HTTP_UPSTREAM_FT_HTTP_403 0x00000100 +#define NGX_HTTP_UPSTREAM_FT_HTTP_404 0x00000200 +#define NGX_HTTP_UPSTREAM_FT_HTTP_429 0x00000400 +#define NGX_HTTP_UPSTREAM_FT_UPDATING 0x00000800 +#define NGX_HTTP_UPSTREAM_FT_BUSY_LOCK 0x00001000 +#define NGX_HTTP_UPSTREAM_FT_MAX_WAITING 0x00002000 +#define NGX_HTTP_UPSTREAM_FT_NON_IDEMPOTENT 0x00004000 +#define NGX_HTTP_UPSTREAM_FT_NOLIVE 0x40000000 +#define NGX_HTTP_UPSTREAM_FT_OFF 0x80000000 + +#define NGX_HTTP_UPSTREAM_FT_STATUS (NGX_HTTP_UPSTREAM_FT_HTTP_500 \ + |NGX_HTTP_UPSTREAM_FT_HTTP_502 \ + |NGX_HTTP_UPSTREAM_FT_HTTP_503 \ + |NGX_HTTP_UPSTREAM_FT_HTTP_504 \ + |NGX_HTTP_UPSTREAM_FT_HTTP_403 \ + |NGX_HTTP_UPSTREAM_FT_HTTP_404 \ + |NGX_HTTP_UPSTREAM_FT_HTTP_429) + +static ngx_conf_bitmask_t ngx_http_proxy_next_upstream_masks[] = { + { ngx_string("error"), NGX_HTTP_UPSTREAM_FT_ERROR }, + { ngx_string("timeout"), NGX_HTTP_UPSTREAM_FT_TIMEOUT }, + { ngx_string("invalid_header"), NGX_HTTP_UPSTREAM_FT_INVALID_HEADER }, + { ngx_string("non_idempotent"), NGX_HTTP_UPSTREAM_FT_NON_IDEMPOTENT }, + { ngx_string("http_500"), NGX_HTTP_UPSTREAM_FT_HTTP_500 }, + { ngx_string("http_502"), NGX_HTTP_UPSTREAM_FT_HTTP_502 }, + { ngx_string("http_503"), NGX_HTTP_UPSTREAM_FT_HTTP_503 }, + { ngx_string("http_504"), NGX_HTTP_UPSTREAM_FT_HTTP_504 }, + { ngx_string("http_403"), NGX_HTTP_UPSTREAM_FT_HTTP_403 }, + { ngx_string("http_404"), NGX_HTTP_UPSTREAM_FT_HTTP_404 }, + { ngx_string("http_429"), NGX_HTTP_UPSTREAM_FT_HTTP_429 }, + { ngx_string("updating"), NGX_HTTP_UPSTREAM_FT_UPDATING }, + { ngx_string("off"), NGX_HTTP_UPSTREAM_FT_OFF }, + { ngx_null_string, 0 } +}; + +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_upstream_next(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