feat(tools): request-aware-table pr feedback

* refactor
* only use static log level to turn concurrency checks on or off
* use single metatable, keep data and request_id in the instance
* use table.new -style initialization
* reintroduce :clear() method

Co-authored-by: Datong Sun <[email protected]>

changelog/unreleased/kong/11017.yaml → changelog/unreleased/kong/request-aware-table.yml

@@ -1,5 +1,3 @@
 message: Add a request-aware table able to detect accesses from different requests.
 type: feature
 scope: Core
-prs:
-  - 11017

kong/tools/request_aware_table.lua

@@ -1,93 +1,118 @@
-local table_clear = require "table.clear"
+--- NOTE: tool is designed to assist with **detecting** request contamination
+-- issues on CI, during test runs. It does not offer security safeguards.
 
-local LOG_LEVELS = require "kong.constants".LOG_LEVELS
-local get_log_level = require("resty.kong.log").get_log_level
+local table_new = require "table.new"
+local table_clear = require "table.clear"
 
-local get_phase = ngx.get_phase
-local fmt = string.format
-local ngx_var = ngx.var
+local debug_mode  = kong.configuration.log_level == "debug"
+local get_phase   = ngx.get_phase
+local var         = ngx.var
+local log         = ngx.log
 
 local NGX_VAR_PHASES = {
-  set = true,
-  rewrite = true,
-  balancer = true,
-  access = true,
-  content = true,
+  set           = true,
+  rewrite       = true,
+  access        = true,
+  content       = true,
   header_filter = true,
-  body_filter = true,
-  log = true,
+  body_filter   = true,
+  log           = true,
+  balancer      = true,
 }
 
+-- Check if access is allowed for table, based on the request ID
+local function enforce_sequential_access(table)
+  if not NGX_VAR_PHASES[get_phase()] then
+    -- not in a request context: allow access and reset allowed request ID
+    rawset(table, "__allowed_request_id", nil)
+    return
+  end
+
+  local curr_request_id = var.request_id
+  local allowed_request_id = rawget(table, "__allowed_request_id")
+  if not allowed_request_id then
+    -- first access. Set allowed request ID and allow access
+    rawset(table, "__allowed_request_id", curr_request_id)
+    return
+  end
 
-local function is_debug_mode()
-  local log_level = get_log_level(LOG_LEVELS[kong.configuration.log_level])
-  return log_level == LOG_LEVELS.debug
+  if curr_request_id ~= table.__allowed_request_id then
+    error("race condition detected; access to table forbidden", 2)
+  end
 end
 
 
---- Request aware table constructor
--- Wraps an existing table (or creates a new one) with request-aware access
--- logic to protect the underlying data from race conditions.
--- @param data_table (optional) The target table to use as underlying data
--- @return The newly created table with request-aware access
-local function new(data_table)
-  if data_table and type(data_table) ~= "table" then
-    error("data_table must be a table", 2)
+local function clear_table(self)
+  if debug_mode == false then
+    table_clear(self)
+    return
   end
 
-  local allowed_request_id
-  local proxy = {}
-  local data  = data_table or {}
+  table_clear(self.__data)
+  rawset(self, "__allowed_request_id", nil)
+end
 
-  -- Check if access is allowed based on the request ID
-  local function enforce_sequential_access()
-    local curr_phase = get_phase()
-    if not NGX_VAR_PHASES[curr_phase] then
-      error(fmt("cannot enforce sequential access in %s phase", curr_phase), 2)
-    end
 
-    local curr_request_id = ngx_var.request_id
+local __proxy_mt = {
+  __index = function(t, k)
+    if k == "clear" then
+      return clear_table
+    end
 
-    allowed_request_id = allowed_request_id or curr_request_id
+    enforce_sequential_access(t)
+    return t.__data[k]
+  end,
 
-    if curr_request_id ~= allowed_request_id then
-      error("race condition detected; access to table forbidden", 2)
+  __newindex = function(t, k, v)
+    if k == "clear" then
+      log(ngx.WARN, "cannot overwrite 'clear' method")
+      return
     end
-  end
 
-  --- Clear data table
-  -- @tparam function fn (optional) An optional function to use instead
-  -- of `table.clear` to clear the data table
-  function proxy.clear(fn)
-    if fn then
-      fn(data)
+    enforce_sequential_access(t)
+    t.__data[k] = v
+  end,
+}
+
 
-    else
-      table_clear(data)
+local __direct_mt = {
+  __index = { clear = clear_table },
+
+  __newindex = function(t, k, v)
+    if k == "clear" then
+      log(ngx.WARN, "cannot overwrite 'clear' method")
+      return
     end
 
-    allowed_request_id = nil
-  end
+    rawset(t, k, v)
+  end,
+}
 
-  local _proxy_mt = {
-    __index = function(_, k)
-      if is_debug_mode() then
-        enforce_sequential_access()
-      end
 
-      return data[k]
-    end,
+-- Request aware table constructor
+--
+-- Creates a new table with request-aware access logic to protect the
+-- underlying data from race conditions.
+-- The table includes a :clear() method to delete all elements.
+--
+-- The request-aware access logic is turned off when `debug_mode` is disabled.
+--
+-- @param narr (optional) pre allocated array elements
+-- @param nrec (optional) pre allocated hash elements
+-- @return The newly created table with request-aware access
+local function new(narr, nrec)
+  if not narr then narr = 0 end
+  if not nrec then nrec = 0 end
 
-    __newindex = function(_, k, v)
-      if is_debug_mode() then
-        enforce_sequential_access()
-      end
+  local data = table_new(narr, nrec)
 
-      data[k] = v
-    end
-  }
+  -- return table without proxy when debug_mode is disabled
+  if debug_mode == false then
+    return setmetatable(data, __direct_mt)
+  end
 
-  return setmetatable(proxy, _proxy_mt)
+  -- wrap table in proxy (for access checks) when debug_mode is enabled
+  return setmetatable({ __data = data }, __proxy_mt)
 end
 
 return {

spec/02-integration/05-proxy/31-request-aware-table_spec.lua → spec/02-integration/05-proxy/33-request-aware-table_spec.lua

@@ -7,14 +7,17 @@ local LOG_LEVELS = {
 }
 
 
-local function clear_table(client)
+local function new_table()
+  local client = helpers.proxy_client()
   local res = client:get("/", {
     query = {
+      new_tab = true,
       clear = true,
     }
   })
   assert.response(res).has.status(200)
   assert.logfile().has.no.line("[error]", true)
+  client:close()
 end
 
 for _, log_level in ipairs(LOG_LEVELS) do
@@ -60,8 +63,8 @@ for _, log_level in ipairs(LOG_LEVELS) do
 
       before_each(function()
         helpers.clean_logfile()
+        new_table()
         client = helpers.proxy_client()
-        clear_table(client)
       end)
 
       after_each(function()
@@ -75,32 +78,40 @@ for _, log_level in ipairs(LOG_LEVELS) do
         assert.response(res).has.status(200)
         assert.logfile().has.no.line("[error]", true)
       end)
-      it("denies access when there are race conditions and checks are enabled", function()
+
+      it("denies access when there are race conditions and checks are enabled (else allows)", function()
         -- access from request 1 (don't clear)
-        local ok, r = pcall(client.get, client, "/")
-        assert(ok)
+        local r = client:get("/")
         assert.response(r).has.status(200)
 
         -- access from request 2
-        ok, r = pcall(client.get, client, "/")
+        r = client:get("/")
         if concurrency_checks then
-          assert(not ok)
           assert.logfile().has.line("race condition detected", true)
         else
-          assert(ok)
           assert.response(r).has.status(200)
         end
       end)
+
       it("allows access when table is cleared between requests", function()
         -- access from request 1 (clear)
         local r = client:get("/", {
           query = {
             clear = true,
-          },
+          }
         })
         assert.response(r).has.status(200)
 
-        -- access from request 2
+        -- access from request 2 (clear)
+        r = client:get("/", {
+          query = {
+            clear = true,
+          }
+        })
+        assert.response(r).has.status(200)
+        assert.logfile().has.no.line("[error]", true)
+
+        -- access from request 3
         r = client:get("/")
         assert.response(r).has.status(200)
         assert.logfile().has.no.line("[error]", true)

spec/fixtures/custom_plugins/kong/plugins/request-aware-table/handler.lua

@@ -1,46 +1,39 @@
 local RAT = require "kong.tools.request_aware_table"
 
-local get_phase = ngx.get_phase
 local kong = kong
-
+local tab
 
 local _M = {
   PRIORITY = 1001,
   VERSION = "1.0",
 }
 
-local tab = RAT.new({})
-
-local function access_tables()
-  local query = kong.request.get_query()
-
-  if query.clear == "true" and get_phase() == "access" then
-    -- clear before access
-    tab.clear()
-  end
-
+local function access_table()
   -- write access
   tab.foo = "bar"
   tab.bar = "baz"
   -- read/write access
   tab.baz = tab.foo .. tab.bar
-
-  if query.clear == "true" and get_phase() == "body_filter" then
-    -- clear after access
-    tab.clear()
-  end
 end
 
+
 function _M:access(conf)
-  access_tables()
-end
+  local query = kong.request.get_query()
+  if query.new_tab == "true" then
+    -- new table
+    tab = RAT.new()
+  end
 
-function _M:header_filter(conf)
-  access_tables()
-end
+  -- access multiple times during same request
+  access_table()
+  access_table()
+  access_table()
 
-function _M:body_filter(conf)
-  access_tables()
+  if query.clear == "true" then
+    -- clear table
+    tab:clear()
+  end
 end
 
+
 return _M