From 9c0aad44f3109a457e23883b7e9121be37c218fb Mon Sep 17 00:00:00 2001
From: Qi <qiqi.zhang@konghq.com>
Date: Fri, 7 Jun 2024 11:42:41 +0800
Subject: [PATCH] refactor(plugin/response-ratelimiting): unify the response
 header processing using private PDK functions

---
 .../response-ratelimiting/header_filter.lua   | 23 +++++++++++--------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/kong/plugins/response-ratelimiting/header_filter.lua b/kong/plugins/response-ratelimiting/header_filter.lua
index 6323d1fb802d..e45c0ee5b480 100644
--- a/kong/plugins/response-ratelimiting/header_filter.lua
+++ b/kong/plugins/response-ratelimiting/header_filter.lua
@@ -1,3 +1,7 @@
+local kong_string = require "kong.tools.string"
+local pdk_private_rl = require "kong.pdk.private.rate_limiting"
+
+
 local kong = kong
 local next = next
 local type = type
@@ -5,12 +9,12 @@ local pairs = pairs
 local ipairs = ipairs
 local tonumber = tonumber
 local math_max = math.max
-local strip = require("kong.tools.string").strip
-local split = require("kong.tools.string").split
 
 
-local RATELIMIT_LIMIT = "X-RateLimit-Limit"
-local RATELIMIT_REMAINING = "X-RateLimit-Remaining"
+local strip = kong_string.strip
+local split = kong_string.split
+local pdk_rl_set_response_headers = pdk_private_rl.set_response_headers
+local pdk_rl_set_limit_by_with_identifier = pdk_private_rl.set_limit_by_with_identifier
 
 
 local function parse_header(header_value, limits)
@@ -60,16 +64,12 @@ function _M.execute(conf)
   end
 
   local stop
+  local ngx_ctx = ngx.ctx
   for limit_name in pairs(usage) do
     for period_name, lv in pairs(usage[limit_name]) do
       if not conf.hide_client_headers then
-        -- increment_value for this current request
-        local limit_hdr  = RATELIMIT_LIMIT .. "-" .. limit_name .. "-" .. period_name
-        local remain_hdr = RATELIMIT_REMAINING .. "-" .. limit_name .. "-" .. period_name
-        kong.response.set_header(limit_hdr, lv.limit)
-
         local remain = math_max(0, lv.remaining - (increments[limit_name] and increments[limit_name] or 0))
-        kong.response.set_header(remain_hdr, remain)
+        pdk_rl_set_limit_by_with_identifier(ngx_ctx, period_name, lv.limit, remain, nil, limit_name)
       end
 
       if increments[limit_name] and increments[limit_name] > 0 and lv.remaining <= 0 then
@@ -78,6 +78,9 @@ function _M.execute(conf)
     end
   end
 
+  -- Set rate-limiting response headers
+  pdk_rl_set_response_headers(ngx_ctx)
+
   kong.response.clear_header(conf.header_name)
 
   -- If limit is exceeded, terminate the request