fix(pdk): # and ipairs on response.get_headers()

Fix #11546, KAG-2602

changelog/unreleased/kong/pdk-response-header-mt.yaml

@@ -0,0 +1,8 @@
+message: "Fix behavior of kong.service.response.get_headers() when using # and pairs()"
+type: bugfix
+scope: Core
+prs:
+  - 11708
+jiras:
+  - "#11546"
+  - "KAG-2602"

kong/pdk/service/response.lua

@@ -10,15 +10,13 @@ local phase_checker = require "kong.pdk.private.phases"
 
 local ngx = ngx
 local sub = string.sub
-local fmt = string.format
 local gsub = string.gsub
 local find = string.find
 local type = type
 local error = error
 local lower = string.lower
 local pairs = pairs
 local tonumber = tonumber
-local getmetatable = getmetatable
 local setmetatable = setmetatable
 local check_phase = phase_checker.check
 
@@ -35,85 +33,104 @@ local header_body_log = phase_checker.new(PHASES.response,
                                           PHASES.log)
 
 
-local attach_resp_headers_mt
-
+local make_resp_headers_aliases
 
 do
-  local resp_headers_orig_mt_index
-
+  local EMPTY = {}
+  local resp_header_alias_mt = {}
+  local index_caches = setmetatable({}, { __mode = "k" })
 
-  local resp_headers_mt = {
-    __index = function(t, name)
-      if type(name) == "string" then
-        local var = fmt("upstream_http_%s", gsub(lower(name), "-", "_"))
-        if not ngx.var[var] then
-          return nil
-        end
-      end
+  local function normalize_header_name(name)
+    return gsub(lower(name), "-", "_")
+  end
 
-      return resp_headers_orig_mt_index(t, name)
-    end,
-  }
+  function resp_header_alias_mt.__index(t, name)
+    if type(name) ~= "string" then
+      return rawget(t, name)
+    end
 
+    name = normalize_header_name(name)
 
-  attach_resp_headers_mt = function(response_headers, err)
-    if not resp_headers_orig_mt_index then
-      local mt = getmetatable(response_headers)
-      resp_headers_orig_mt_index = mt.__index
+    local value = rawget(t, name)
+    if value ~= nil then
+      return value
     end
 
-    setmetatable(response_headers, resp_headers_mt)
+    local index_cache = index_caches[t]
 
-    return response_headers, err
-  end
-end
+    if index_cache == nil then
+      index_cache = {}
+      for n, v in pairs(t) do
+        local normalized_n = normalize_header_name(n)
+        index_cache[normalized_n] = v
+      end
+      index_caches[t] = index_cache
+    end
 
+    return index_cache[name]
+  end
 
-local attach_buffered_headers_mt
+  function resp_header_alias_mt.__newindex(t, name, value)
+    local normalized_header_name = normalize_header_name(name)
 
-do
-  local EMPTY = {}
+    -- update index cache
+    local index_cache = index_caches[t]
+    if index_cache then
+      index_cache[normalized_header_name] = value
+    end
 
-  attach_buffered_headers_mt = function(response_headers, max_headers)
-    if not response_headers then
-      return EMPTY
+    -- quick path
+    if type(name) ~= "string" or t[name] == nil then
+      rawset(t, name, value)
+      return
     end
 
-    return setmetatable({}, { __index = function(_, name)
-      if type(name) ~= "string" then
-        return nil
+    for n, _ in pairs(t) do
+      local normalized_n = normalize_header_name(n)
+      if normalized_n == normalized_header_name then
+        rawset(t, n, value)
+        return
       end
+    end
 
-      if response_headers[name] then
-        return response_headers[name]
-      end
+    error("unreachable")
+  end
 
-      name = lower(name)
+  function make_resp_headers_aliases(t, err)
+    if not t then
+      return EMPTY, err
+    end
 
-      if response_headers[name] then
-        return response_headers[name]
-      end
+    local ret = setmetatable(t, resp_header_alias_mt)
 
-      name = gsub(name, "-", "_")
+    return ret, err
+  end
+end
 
-      if response_headers[name] then
-        return response_headers[name]
-      end
 
-      local i = 1
-      for n, v in pairs(response_headers) do
-        if i > max_headers then
-          return nil
-        end
+local add_fast_path_headers
 
-        n = gsub(lower(n), "-", "_")
-        if n == name then
-          return v
-        end
+do
+  local FAST_PATH_HEADERS = {
+    ["Location"] = "upstream_http_location",
+    ["Last_Modified"] = "upstream_http_last_modified",
+    ["Keep_Alive"] = "upstream_http_keep_alive",
+    ["Cache_Control"] = "upstream_http_cache_control",
+    ["Link"] = "upstream_http_link",
+    ["Date"] = "upstream_http_date",
+    ["Server"] = "upstream_http_server",
+  }
 
-        i = i + 1
+  -- we do not drop headers exceeding the limit as they are already there
+  function add_fast_path_headers(t, _)
+    for header_name, var_name in pairs(FAST_PATH_HEADERS) do
+      local value = ngx.var[var_name]
+      if value and t[header_name] == nil then
+        t[header_name] = value
       end
-    end })
+    end
+
+    return t
   end
 end
 
@@ -199,17 +216,12 @@ local function new(pdk, major_version)
 
     local buffered_headers = ngx.ctx.buffered_headers
 
-    if max_headers == nil then
-      if buffered_headers then
-        if not MAX_HEADERS_CONFIGURED then
-          MAX_HEADERS_CONFIGURED = pdk and pdk.configuration and pdk.configuration.lua_max_resp_headers
-        end
-        return attach_buffered_headers_mt(buffered_headers, MAX_HEADERS_CONFIGURED or MAX_HEADERS_DEFAULT)
-      end
-
-      return attach_resp_headers_mt(ngx.resp.get_headers())
+    if not MAX_HEADERS_CONFIGURED then
+      MAX_HEADERS_CONFIGURED = pdk and pdk.configuration and pdk.configuration.lua_max_resp_headers
     end
 
+    max_headers = max_headers or MAX_HEADERS_CONFIGURED or MAX_HEADERS_DEFAULT
+
     if type(max_headers) ~= "number" then
       error("max_headers must be a number", 2)
 
@@ -220,11 +232,12 @@ local function new(pdk, major_version)
       error("max_headers must be <= " .. MAX_HEADERS, 2)
     end
 
-    if buffered_headers then
-      return attach_buffered_headers_mt(buffered_headers, max_headers)
+    if not buffered_headers then
+      -- we let get_headers to validate the max_headers
+      return make_resp_headers_aliases(ngx.resp.get_headers(max_headers))
     end
 
-    return attach_resp_headers_mt(ngx.resp.get_headers(max_headers))
+    return make_resp_headers_aliases(add_fast_path_headers(buffered_headers, max_headers))
   end
 
   ---

t/01-pdk/07-service-response/02-get_headers.t

@@ -470,3 +470,65 @@ X-Non-Service-Header: nil
 X-Non-Service-Header: test
 --- no_error_log
 [error]
+
+
+
+=== TEST 11: service.response.get_headers() # and pairs works with returned table
+--- http_config eval
+qq{
+    $t::Util::HttpConfig
+
+    server {
+        listen unix:$ENV{TEST_NGINX_NXSOCK}/nginx.sock;
+
+        location / {
+            content_by_lua_block {
+                ngx.header["test"] = "test"
+            }
+        }
+    }
+}
+--- config
+    location = /t {
+        access_by_lua_block {
+            ngx.header["X-Non-Service-Header"] = "test"
+        }
+
+        proxy_pass http://unix:$TEST_NGINX_NXSOCK/nginx.sock;
+
+        header_filter_by_lua_block {
+            ngx.header.content_length = nil
+        }
+
+        body_filter_by_lua_block {
+            local PDK = require "kong.pdk"
+            local pdk = PDK.new()
+
+            local headers = pdk.service.response.get_headers()
+
+            local ordered_headers = {}
+
+            for k, v in pairs(headers) do
+                table.insert(ordered_headers, k)
+            end
+
+            table.sort(ordered_headers)
+
+            for _, k in ipairs(ordered_headers) do
+                ngx.arg[1] = ngx.arg[1] .. k .. ": " .. headers[k] .. "\n"
+            end
+
+            ngx.arg[1] = ngx.arg[1] .. "#headers: " .. #headers
+
+            ngx.arg[2] = true
+        }
+    }
+--- request
+GET /t
+--- response_body chop
+connection: close
+content-type: text/plain
+test: test
+#headers: 3
+--- no_error_log
+[error]