From f78588bf019d757e1281326c70fd788ff2b982c4 Mon Sep 17 00:00:00 2001 From: Qi Date: Fri, 7 Jun 2024 07:19:52 +0000 Subject: [PATCH] fixup! feat(pdk): implement private PDK functions to process rate_limting response headers --- kong/pdk/private/rate_limiting.lua | 397 ++++++++++++++--------------- 1 file changed, 194 insertions(+), 203 deletions(-) diff --git a/kong/pdk/private/rate_limiting.lua b/kong/pdk/private/rate_limiting.lua index a93d11841840..bfeb19774fe0 100644 --- a/kong/pdk/private/rate_limiting.lua +++ b/kong/pdk/private/rate_limiting.lua @@ -1,16 +1,16 @@ -local table_new = require("table.new") -local buffer = require("string.buffer") +local table_new = require("table.new") +local buffer = require("string.buffer") -local type = type -local pairs = pairs -local assert = assert -local tostring = tostring -local resp_header = ngx.header +local type = type +local pairs = pairs +local assert = assert +local tostring = tostring +local resp_header = ngx.header -local RL_LIMIT = "Ratelimit-Limit" -local RL_REMAINING = "Ratelimit-Remaining" -local RL_RESET = "Ratelimit-Reset" -local RETRY_AFTER = "Retry-After" +local RL_LIMIT = "Ratelimit-Limit" +local RL_REMAINING = "Ratelimit-Remaining" +local RL_RESET = "Ratelimit-Reset" +local RETRY_AFTER = "Retry-After" -- determine the number of pre-allocated fields at runtime @@ -18,255 +18,246 @@ local max_fields_n = 4 local buf = buffer.new(64) local LIMIT_BY = { - second = { - limit = "X-Ratelimit-Limit-Second", - remain = "X-Ratelimit-Remaining-Second", - limit_segment_0 = "X-", - limit_segment_1 = "Ratelimit-Limit-", - limit_segment_3 = "-Second", - remain_segment_0 = "X-", - remain_segment_1 = "Ratelimit-Remaining-", - remain_segment_3 = "-Second", - }, - minute = { - limit = "X-Ratelimit-Limit-Minute", - remain = "X-Ratelimit-Remaining-Minute", - limit_segment_0 = "X-", - limit_segment_1 = "Ratelimit-Limit-", - limit_segment_3 = "-Minute", - remain_segment_0 = "X-", - remain_segment_1 = "Ratelimit-Remaining-", - remain_segment_3 = "-Minute", - }, - hour = { - limit = "X-Ratelimit-Limit-Hour", - remain = "X-Ratelimit-Remaining-Hour", - limit_segment_0 = "X-", - limit_segment_1 = "Ratelimit-Limit-", - limit_segment_3 = "-Hour", - remain_segment_0 = "X-", - remain_segment_1 = "Ratelimit-Remaining-", - remain_segment_3 = "-Hour", - }, - day = { - limit = "X-Ratelimit-Limit-Day", - remain = "X-Ratelimit-Remaining-Day", - limit_segment_0 = "X-", - limit_segment_1 = "Ratelimit-Limit-", - limit_segment_3 = "-Day", - remain_segment_0 = "X-", - remain_segment_1 = "Ratelimit-Remaining-", - remain_segment_3 = "-Day", - }, - month = { - limit = "X-Ratelimit-Limit-Month", - remain = "X-Ratelimit-Remaining-Month", - limit_segment_0 = "X-", - limit_segment_1 = "Ratelimit-Limit-", - limit_segment_3 = "-Month", - remain_segment_0 = "X-", - remain_segment_1 = "Ratelimit-Remaining-", - remain_segment_3 = "-Month", - }, - year = { - limit = "X-Ratelimit-Limit-Year", - remain = "X-Ratelimit-Remaining-Year", - limit_segment_0 = "X-", - limit_segment_1 = "Ratelimit-Limit-", - limit_segment_3 = "-Year", - remain_segment_0 = "X-", - remain_segment_1 = "Ratelimit-Remaining-", - remain_segment_3 = "-Year", - }, + second = { + limit = "X-Ratelimit-Limit-Second", + remain = "X-Ratelimit-Remaining-Second", + limit_segment_0 = "X-", + limit_segment_1 = "Ratelimit-Limit-", + limit_segment_3 = "-Second", + remain_segment_0 = "X-", + remain_segment_1 = "Ratelimit-Remaining-", + remain_segment_3 = "-Second", + }, + minute = { + limit = "X-Ratelimit-Limit-Minute", + remain = "X-Ratelimit-Remaining-Minute", + limit_segment_0 = "X-", + limit_segment_1 = "Ratelimit-Limit-", + limit_segment_3 = "-Minute", + remain_segment_0 = "X-", + remain_segment_1 = "Ratelimit-Remaining-", + remain_segment_3 = "-Minute", + }, + hour = { + limit = "X-Ratelimit-Limit-Hour", + remain = "X-Ratelimit-Remaining-Hour", + limit_segment_0 = "X-", + limit_segment_1 = "Ratelimit-Limit-", + limit_segment_3 = "-Hour", + remain_segment_0 = "X-", + remain_segment_1 = "Ratelimit-Remaining-", + remain_segment_3 = "-Hour", + }, + day = { + limit = "X-Ratelimit-Limit-Day", + remain = "X-Ratelimit-Remaining-Day", + limit_segment_0 = "X-", + limit_segment_1 = "Ratelimit-Limit-", + limit_segment_3 = "-Day", + remain_segment_0 = "X-", + remain_segment_1 = "Ratelimit-Remaining-", + remain_segment_3 = "-Day", + }, + month = { + limit = "X-Ratelimit-Limit-Month", + remain = "X-Ratelimit-Remaining-Month", + limit_segment_0 = "X-", + limit_segment_1 = "Ratelimit-Limit-", + limit_segment_3 = "-Month", + remain_segment_0 = "X-", + remain_segment_1 = "Ratelimit-Remaining-", + remain_segment_3 = "-Month", + }, + year = { + limit = "X-Ratelimit-Limit-Year", + remain = "X-Ratelimit-Remaining-Year", + limit_segment_0 = "X-", + limit_segment_1 = "Ratelimit-Limit-", + limit_segment_3 = "-Year", + remain_segment_0 = "X-", + remain_segment_1 = "Ratelimit-Remaining-", + remain_segment_3 = "-Year", + }, } local _M = {} local function _has_rl_ctx(ngx_ctx) - return ngx_ctx.__rate_limiting_context__ ~= nil + return ngx_ctx.__rate_limiting_context__ ~= nil end local function _create_rl_ctx(ngx_ctx) - assert(not _has_rl_ctx(ngx_ctx), "rate limiting context already exists") - local ctx = table_new(0, max_fields_n) - ngx_ctx.__rate_limiting_context__ = ctx - return ctx + assert(not _has_rl_ctx(ngx_ctx), "rate limiting context already exists") + local ctx = table_new(0, max_fields_n) + ngx_ctx.__rate_limiting_context__ = ctx + return ctx end local function _get_rl_ctx(ngx_ctx) - assert(_has_rl_ctx(ngx_ctx), "rate limiting context does not exist") - return ngx_ctx.__rate_limiting_context__ + assert(_has_rl_ctx(ngx_ctx), "rate limiting context does not exist") + return ngx_ctx.__rate_limiting_context__ end local function _get_or_create_rl_ctx(ngx_ctx) - if not _has_rl_ctx(ngx_ctx) then - _create_rl_ctx(ngx_ctx) - end + if not _has_rl_ctx(ngx_ctx) then + _create_rl_ctx(ngx_ctx) + end - local rl_ctx = _get_rl_ctx(ngx_ctx) - return rl_ctx + local rl_ctx = _get_rl_ctx(ngx_ctx) + return rl_ctx end function _M.set_basic_limit(ngx_ctx, limit, remaining, reset) - local rl_ctx = _get_or_create_rl_ctx(ngx_ctx or ngx.ctx) + local rl_ctx = _get_or_create_rl_ctx(ngx_ctx or ngx.ctx) - assert(type(limit) == "number", "limit must be a number") - assert(type(remaining) == "number", "remaining must be a number") - assert(type(reset) == "number", "reset must be a number") + assert(type(limit) == "number", "limit must be a number") + assert(type(remaining) == "number", "remaining must be a number") + assert(type(reset) == "number", "reset must be a number") - rl_ctx[RL_LIMIT] = limit - rl_ctx[RL_REMAINING] = remaining - rl_ctx[RL_RESET] = reset + rl_ctx[RL_LIMIT] = limit + rl_ctx[RL_REMAINING] = remaining + rl_ctx[RL_RESET] = reset end - function _M.set_retry_after(ngx_ctx, reset) - local rl_ctx = _get_or_create_rl_ctx(ngx_ctx or ngx.ctx) + local rl_ctx = _get_or_create_rl_ctx(ngx_ctx or ngx.ctx) - assert(type(reset) == "number", "reset must be a number") + assert(type(reset) == "number", "reset must be a number") - rl_ctx[RETRY_AFTER] = reset + rl_ctx[RETRY_AFTER] = reset end - function _M.set_limit_by(ngx_ctx, limit_by, limit, remaining) - local rl_ctx = _get_or_create_rl_ctx(ngx_ctx or ngx.ctx) + local rl_ctx = _get_or_create_rl_ctx(ngx_ctx or ngx.ctx) - assert(type(limit_by) == "string", "limit_by must be a string") - assert(type(limit) == "number", "limit must be a number") - assert(type(remaining) == "number", "remaining must be a number") + assert(type(limit_by) == "string", "limit_by must be a string") + assert(type(limit) == "number", "limit must be a number") + assert(type(remaining) == "number", "remaining must be a number") - limit_by = LIMIT_BY[limit_by] - assert(limit_by, "invalid limit_by") + limit_by = LIMIT_BY[limit_by] + assert(limit_by, "invalid limit_by") - rl_ctx[limit_by.limit] = limit - rl_ctx[limit_by.remain] = remaining + rl_ctx[limit_by.limit] = limit + rl_ctx[limit_by.remain] = remaining end - function _M.set_limit_by_with_identifier(ngx_ctx, limit_by, limit, remaining, id_seg_1, id_seg_2) - local rl_ctx = _get_or_create_rl_ctx(ngx_ctx or ngx.ctx) - - assert(type(id_seg_1) == "string" or type(id_seg_2) == "string", "id_seg_1 or id_seg_2 must be a string") - assert(type(limit_by) == "string", "limit_by must be a string") - assert(type(limit) == "number", "limit must be a number") - assert(type(remaining) == "number", "remaining must be a number") - - limit_by = LIMIT_BY[limit_by] - assert(limit_by, "invalid limit_by") - - id_seg_1 = id_seg_1 or "" - id_seg_2 = id_seg_2 or "" - - -- construct the key like X--Ratelimit-Limit-- - local limit_key = buf:reset():put( - limit_by.limit_segment_0, - id_seg_1, - limit_by.limit_segment_1, - id_seg_2, - limit_by.limit_segment_3 - ):get() - - -- construct the key like X--Ratelimit-Remaining-- - local remain_key = buf:reset():put( - limit_by.remain_segment_0, - id_seg_1, - limit_by.remain_segment_1, - id_seg_2, - limit_by.remain_segment_3 - ):get() - - rl_ctx[limit_key] = limit - rl_ctx[remain_key] = remaining + local rl_ctx = _get_or_create_rl_ctx(ngx_ctx or ngx.ctx) + + assert(type(id_seg_1) == "string" or type(id_seg_2) == "string", "id_seg_1 or id_seg_2 must be a string") + assert(type(limit_by) == "string", "limit_by must be a string") + assert(type(limit) == "number", "limit must be a number") + assert(type(remaining) == "number", "remaining must be a number") + + limit_by = LIMIT_BY[limit_by] + assert(limit_by, "invalid limit_by") + + id_seg_1 = id_seg_1 or "" + id_seg_2 = id_seg_2 or "" + + -- construct the key like X--Ratelimit-Limit-- + local limit_key = buf:reset():put( + limit_by.limit_segment_0, + id_seg_1, + limit_by.limit_segment_1, + id_seg_2, + limit_by.limit_segment_3 + ):get() + + -- construct the key like X--Ratelimit-Remaining-- + local remain_key = buf:reset():put( + limit_by.remain_segment_0, + id_seg_1, + limit_by.remain_segment_1, + id_seg_2, + limit_by.remain_segment_3 + ):get() + + rl_ctx[limit_key] = limit + rl_ctx[remain_key] = remaining end - function _M.get_basic_limit(ngx_ctx) - local rl_ctx = _get_rl_ctx(ngx_ctx or ngx.ctx) - return rl_ctx[RL_LIMIT], rl_ctx[RL_REMAINING], rl_ctx[RL_RESET] + local rl_ctx = _get_rl_ctx(ngx_ctx or ngx.ctx) + return rl_ctx[RL_LIMIT], rl_ctx[RL_REMAINING], rl_ctx[RL_RESET] end - function _M.get_retry_after(ngx_ctx) - local rl_ctx = _get_rl_ctx(ngx_ctx or ngx.ctx) - return rl_ctx[RETRY_AFTER] + local rl_ctx = _get_rl_ctx(ngx_ctx or ngx.ctx) + return rl_ctx[RETRY_AFTER] end - function _M.get_limit_by(ngx_ctx, limit_by) - local rl_ctx = _get_rl_ctx(ngx_ctx or ngx.ctx) + local rl_ctx = _get_rl_ctx(ngx_ctx or ngx.ctx) - assert(type(limit_by) == "string", "limit_by must be a string") + assert(type(limit_by) == "string", "limit_by must be a string") - limit_by = LIMIT_BY[limit_by] - assert(limit_by, "invalid limit_by") + limit_by = LIMIT_BY[limit_by] + assert(limit_by, "invalid limit_by") - return rl_ctx[limit_by.limit], rl_ctx[limit_by.remain] + return rl_ctx[limit_by.limit], rl_ctx[limit_by.remain] end - function _M.get_limit_by_with_identifier(ngx_ctx, limit_by, id_seg_1, id_seg_2) - local rl_ctx = _get_rl_ctx(ngx_ctx or ngx.ctx) - - assert(type(id_seg_1) == "string" or type(id_seg_2) == "string", "id_seg_1 or id_seg_2 must be a string") - assert(type(limit_by) == "string", "limit_by must be a string") - - limit_by = LIMIT_BY[limit_by] - assert(limit_by, "invalid limit_by") - - id_seg_1 = id_seg_1 or "" - id_seg_2 = id_seg_2 or "" - - -- construct the key like X--Ratelimit-Limit-- - local limit_key = buf:reset():put( - limit_by.limit_segment_0, - id_seg_1, - limit_by.limit_segment_1, - id_seg_2, - limit_by.limit_segment_3 - ):get() - - -- construct the key like X--Ratelimit-Remaining-- - local remain_key = buf:reset():put( - limit_by.remain_segment_0, - id_seg_1, - limit_by.remain_segment_1, - id_seg_2, - limit_by.remain_segment_3 - ):get() - - return rl_ctx[limit_key], rl_ctx[remain_key] + local rl_ctx = _get_rl_ctx(ngx_ctx or ngx.ctx) + + assert(type(id_seg_1) == "string" or type(id_seg_2) == "string", "id_seg_1 or id_seg_2 must be a string") + assert(type(limit_by) == "string", "limit_by must be a string") + + limit_by = LIMIT_BY[limit_by] + assert(limit_by, "invalid limit_by") + + id_seg_1 = id_seg_1 or "" + id_seg_2 = id_seg_2 or "" + + -- construct the key like X--Ratelimit-Limit-- + local limit_key = buf:reset():put( + limit_by.limit_segment_0, + id_seg_1, + limit_by.limit_segment_1, + id_seg_2, + limit_by.limit_segment_3 + ):get() + + -- construct the key like X--Ratelimit-Remaining-- + local remain_key = buf:reset():put( + limit_by.remain_segment_0, + id_seg_1, + limit_by.remain_segment_1, + id_seg_2, + limit_by.remain_segment_3 + ):get() + + return rl_ctx[limit_key], rl_ctx[remain_key] end - function _M.set_response_headers(ngx_ctx) - if not _has_rl_ctx(ngx_ctx) then - return - end - - local rl_ctx = _get_rl_ctx(ngx_ctx) - local actual_fields_n = 0 - - for k, v in pairs(rl_ctx) do - resp_header[k] = tostring(v) - actual_fields_n = actual_fields_n + 1 - end - - if actual_fields_n > max_fields_n then - local msg = string.format( - "[private-rl-pdk] bumpping pre-allocated fields from %d to %d for performance reasons", - max_fields_n, - actual_fields_n - ) - ngx.log(ngx.INFO, msg) - max_fields_n = actual_fields_n - end + if not _has_rl_ctx(ngx_ctx) then + return + end + + local rl_ctx = _get_rl_ctx(ngx_ctx) + local actual_fields_n = 0 + + for k, v in pairs(rl_ctx) do + resp_header[k] = tostring(v) + actual_fields_n = actual_fields_n + 1 + end + + if actual_fields_n > max_fields_n then + local msg = string.format( + "[private-rl-pdk] bumpping pre-allocated fields from %d to %d for performance reasons", + max_fields_n, + actual_fields_n + ) + ngx.log(ngx.INFO, msg) + max_fields_n = actual_fields_n + end end - return _M