bugfix: better caller context checks in the Lua APIs of ngx.balancer.

lua/ngx/balancer.lua

@@ -20,7 +20,7 @@ int ngx_http_lua_ffi_balancer_set_current_peer(ngx_http_request_t *r,
 int ngx_http_lua_ffi_balancer_set_more_tries(ngx_http_request_t *r,
     int count, char **err);
 
-unsigned ngx_http_lua_ffi_balancer_get_last_failure(ngx_http_request_t *r,
+int ngx_http_lua_ffi_balancer_get_last_failure(ngx_http_request_t *r,
     int *status, char **err);
 ]]
 
@@ -90,7 +90,7 @@ function _M.get_last_failure()
     end
 
     if state == FFI_ERROR then
-        return nil, ffi_str(errmsg[0])
+        return nil, nil, ffi_str(errmsg[0])
     end
 
     return peer_state_names[state] or "unknown", int_out[0]

src/ngx_http_lua_balancer.c

@@ -335,7 +335,6 @@ ngx_http_lua_balancer_by_chunk(lua_State *L, ngx_http_request_t *r)
     }
 
     ctx->context = NGX_HTTP_LUA_CONTEXT_BALANCER;
-    ctx->entered_balancer_phase = 1;
 
     /* init nginx context in Lua VM */
     ngx_http_lua_set_req(L, r);
@@ -431,8 +430,13 @@ ngx_http_lua_ffi_balancer_set_current_peer(ngx_http_request_t *r,
     }
 
     ctx = ngx_http_get_module_ctx(r, ngx_http_lua_module);
-    if (ctx == NULL || !ctx->entered_balancer_phase) {
-        *err = "bad lua ctx";
+    if (ctx == NULL) {
+        *err = "no ctx found";
+        return NGX_ERROR;
+    }
+
+    if ((ctx->context & NGX_HTTP_LUA_CONTEXT_BALANCER) == 0) {
+        *err = "API disabled in the current context";
         return NGX_ERROR;
     }
 
@@ -510,8 +514,13 @@ ngx_http_lua_ffi_balancer_set_more_tries(ngx_http_request_t *r,
     }
 
     ctx = ngx_http_get_module_ctx(r, ngx_http_lua_module);
-    if (ctx == NULL || !ctx->entered_balancer_phase) {
-        *err = "bad lua ctx";
+    if (ctx == NULL) {
+        *err = "no ctx found";
+        return NGX_ERROR;
+    }
+
+    if ((ctx->context & NGX_HTTP_LUA_CONTEXT_BALANCER) == 0) {
+        *err = "API disabled in the current context";
         return NGX_ERROR;
     }
 
@@ -562,8 +571,13 @@ ngx_http_lua_ffi_balancer_get_last_failure(ngx_http_request_t *r,
     }
 
     ctx = ngx_http_get_module_ctx(r, ngx_http_lua_module);
-    if (ctx == NULL || !ctx->entered_balancer_phase) {
-        *err = "bad lua ctx";
+    if (ctx == NULL) {
+        *err = "no ctx found";
+        return NGX_ERROR;
+    }
+
+    if ((ctx->context & NGX_HTTP_LUA_CONTEXT_BALANCER) == 0) {
+        *err = "API disabled in the current context";
         return NGX_ERROR;
     }
 

src/ngx_http_lua_common.h

@@ -453,7 +453,6 @@ typedef struct ngx_http_lua_ctx_s {
     unsigned         entered_rewrite_phase:1;
     unsigned         entered_access_phase:1;
     unsigned         entered_content_phase:1;
-    unsigned         entered_balancer_phase:1;
 
     unsigned         buffering:1; /* HTTP 1.0 response body buffering flag */
 

t/133-balancer.t

@@ -667,3 +667,132 @@ free keepalive peer: saving connection
 ]
 --- no_error_log
 [warn]
+
+
+
+=== TEST 19: set_current_peer called in a wrong context
+--- wait: 0.2
+--- http_config
+    lua_package_path "../lua-resty-core/lib/?.lua;lua/?.lua;;";
+
+    upstream backend {
+        server 127.0.0.1:$TEST_NGINX_SERVER_PORT;
+        balancer_by_lua_block {
+            print("hello from balancer by lua!")
+        }
+    }
+
+--- config
+
+    location = /fake {
+        echo ok;
+    }
+
+    location = /t {
+        proxy_pass http://backend/fake;
+
+        log_by_lua_block {
+            local balancer = require "ngx.balancer"
+            local ok, err = balancer.set_current_peer("127.0.0.1", 1234)
+            if not ok then
+                ngx.log(ngx.ERR, "failed to call: ", err)
+                return
+            end
+            ngx.log(ngx.ALERT, "unexpected success")
+        }
+    }
+
+--- request
+GET /t
+--- response_body
+ok
+--- error_log eval
+qr/\[error\] .*? log_by_lua.*? failed to call: API disabled in the current context/
+--- no_error_log
+[alert]
+
+
+
+=== TEST 20: get_last_failure called in a wrong context
+--- wait: 0.2
+--- http_config
+    lua_package_path "../lua-resty-core/lib/?.lua;lua/?.lua;;";
+
+    upstream backend {
+        server 127.0.0.1:$TEST_NGINX_SERVER_PORT;
+        balancer_by_lua_block {
+            print("hello from balancer by lua!")
+        }
+    }
+
+--- config
+
+    location = /fake {
+        echo ok;
+    }
+
+    location = /t {
+        proxy_pass http://backend/fake;
+
+        log_by_lua_block {
+            local balancer = require "ngx.balancer"
+            local state, status, err = balancer.get_last_failure()
+            if not state and err then
+                ngx.log(ngx.ERR, "failed to call: ", err)
+                return
+            end
+            ngx.log(ngx.ALERT, "unexpected success")
+        }
+    }
+
+--- request
+GET /t
+--- response_body
+ok
+--- error_log eval
+qr/\[error\] .*? log_by_lua.*? failed to call: API disabled in the current context/
+--- no_error_log
+[alert]
+
+
+
+=== TEST 21: set_more_tries called in a wrong context
+--- wait: 0.2
+--- http_config
+    lua_package_path "../lua-resty-core/lib/?.lua;lua/?.lua;;";
+
+    upstream backend {
+        server 127.0.0.1:$TEST_NGINX_SERVER_PORT;
+        balancer_by_lua_block {
+            print("hello from balancer by lua!")
+        }
+    }
+
+--- config
+
+    location = /fake {
+        echo ok;
+    }
+
+    location = /t {
+        proxy_pass http://backend/fake;
+
+        log_by_lua_block {
+            local balancer = require "ngx.balancer"
+            local ok, err = balancer.set_more_tries(1)
+            if not ok then
+                ngx.log(ngx.ERR, "failed to call: ", err)
+                return
+            end
+            ngx.log(ngx.ALERT, "unexpected success")
+        }
+    }
+
+--- request
+GET /t
+--- response_body
+ok
+--- error_log eval
+qr/\[error\] .*? log_by_lua.*? failed to call: API disabled in the current context/
+--- no_error_log
+[alert]