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

Author: oowl

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

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

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

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