From 494d0f08dc3cde08c7ec03082731ce50e63d5764 Mon Sep 17 00:00:00 2001 From: Wangchong Zhou Date: Fri, 5 Jul 2024 18:21:42 +0800 Subject: [PATCH] feat(pdk): allow certain API to be called in balancer phase --- kong/pdk/private/phases.lua | 1 + kong/pdk/response.lua | 1 + kong/pdk/service.lua | 2 +- kong/pdk/service/request.lua | 22 ++++++++++++---------- t/01-pdk/09-service/00-phase_checks.t | 4 ++-- 5 files changed, 17 insertions(+), 13 deletions(-) diff --git a/kong/pdk/private/phases.lua b/kong/pdk/private/phases.lua index e5787a0f2ef..d3a2bca5717 100644 --- a/kong/pdk/private/phases.lua +++ b/kong/pdk/private/phases.lua @@ -122,6 +122,7 @@ end local public_phases = setmetatable({ request = new_phase(PHASES.rewrite, PHASES.access, + PHASES.balancer, PHASES.response, PHASES.header_filter, PHASES.body_filter, diff --git a/kong/pdk/response.lua b/kong/pdk/response.lua index 8877cb49431..844d5c7d139 100644 --- a/kong/pdk/response.lua +++ b/kong/pdk/response.lua @@ -56,6 +56,7 @@ local header_body_log = phase_checker.new(PHASES.response, local rewrite_access_header = phase_checker.new(PHASES.rewrite, PHASES.access, PHASES.response, + PHASES.balancer, PHASES.header_filter, PHASES.error, PHASES.admin_api) diff --git a/kong/pdk/service.lua b/kong/pdk/service.lua index 9a8bc068ddf..59b0f9b6203 100644 --- a/kong/pdk/service.lua +++ b/kong/pdk/service.lua @@ -88,7 +88,7 @@ local function new() -- kong.service.set_target("service.local", 443) -- kong.service.set_target("192.168.130.1", 80) function service.set_target(host, port) - check_phase(PHASES.access) + check_phase(access_and_rewrite_and_balancer_preread) if type(host) ~= "string" then error("host must be a string", 2) diff --git a/kong/pdk/service/request.lua b/kong/pdk/service/request.lua index 495dbf0febc..28fab489e6a 100644 --- a/kong/pdk/service/request.lua +++ b/kong/pdk/service/request.lua @@ -83,12 +83,12 @@ local function new(self) -- Enables buffered proxying, which allows plugins to access Service body and -- response headers at the same time. -- @function kong.service.request.enable_buffering - -- @phases `rewrite`, `access` + -- @phases `rewrite`, `access`, `balancer` -- @return Nothing. -- @usage -- kong.service.request.enable_buffering() request.enable_buffering = function() - check_phase(access_and_rewrite) + check_phase(access_rewrite_balancer) if ngx.req.http_version() >= 2 then error("buffered proxying cannot currently be enabled with http/" .. @@ -102,13 +102,13 @@ local function new(self) --- -- Sets the protocol to use when proxying the request to the Service. -- @function kong.service.request.set_scheme - -- @phases `access` + -- @phases `access`, `rewrite`, `balancer` -- @tparam string scheme The scheme to be used. Supported values are `"http"` or `"https"`. -- @return Nothing; throws an error on invalid inputs. -- @usage -- kong.service.request.set_scheme("https") request.set_scheme = function(scheme) - check_phase(PHASES.access) + check_phase(access_rewrite_balancer) if type(scheme) ~= "string" then error("scheme must be a string", 2) @@ -131,14 +131,14 @@ local function new(self) -- -- Input should **not** include the query string. -- @function kong.service.request.set_path - -- @phases `access` + -- @phases `access`, `rewrite`, `balancer` -- @tparam string path The path string. Special characters and UTF-8 -- characters are allowed, for example: `"/v2/movies"` or `"/foo/😀"`. -- @return Nothing; throws an error on invalid inputs. -- @usage -- kong.service.request.set_path("/v2/movies") request.set_path = function(path) - check_phase(PHASES.access) + check_phase(access_rewrite_balancer) if type(path) ~= "string" then error("path must be a string", 2) @@ -440,13 +440,13 @@ local function new(self) -- For a higher-level function to set the body based on the request content type, -- see `kong.service.request.set_body()`. -- @function kong.service.request.set_raw_body - -- @phases `rewrite`, `access` + -- @phases `rewrite`, `access`, `balancer` -- @tparam string body The raw body. -- @return Nothing; throws an error on invalid inputs. -- @usage -- kong.service.request.set_raw_body("Hello, world!") request.set_raw_body = function(body) - check_phase(access_and_rewrite) + check_phase(access_rewrite_balancer) if type(body) ~= "string" then error("body must be a string", 2) @@ -459,7 +459,9 @@ local function new(self) -- Ensure client request body has been read. -- This function is a nop if body has already been read, -- and necessary to write the request to the service if it has not. - ngx.req.read_body() + if ngx.get_phase() ~= "balancer" then + ngx.req.read_body() + end ngx.req.set_body_data(body) end @@ -594,7 +596,7 @@ local function new(self) -- a string with `kong.service.request.set_raw_body()`. -- -- @function kong.service.request.set_body - -- @phases `rewrite`, `access` + -- @phases `rewrite`, `access`, `balancer` -- @tparam table args A table with data to be converted to the appropriate format -- and stored in the body. -- @tparam[opt] string mimetype can be one of: diff --git a/t/01-pdk/09-service/00-phase_checks.t b/t/01-pdk/09-service/00-phase_checks.t index 0262c56c450..87869a82631 100644 --- a/t/01-pdk/09-service/00-phase_checks.t +++ b/t/01-pdk/09-service/00-phase_checks.t @@ -69,7 +69,7 @@ qq{ args = { "example.com", 8000 }, init_worker = false, certificate = "pending", - rewrite = "forced false", + rewrite = true, access = true, response = "forced false", header_filter = "forced false", @@ -243,7 +243,7 @@ qq{ args = { "example.com", 8000 }, init_worker = false, certificate = "pending", - rewrite = "forced false", + rewrite = true, access = true, response = "forced false", header_filter = "forced false",