feat: control the proxy_upstream in lua side

.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:

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
 =======
 

lualib/resty/kong/upstream.lua

@@ -0,0 +1,85 @@
+-- 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_next_upstream(ngx_http_request_t *r, uint32_t next_upstream, char **err);
+]])
+
+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 = 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,
+	off = 0x00001000,
+	non_idempotent = 0x00004000,
+}
+
+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

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;
 
 

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_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;
+}

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_ */