Skip to content

Commit

Permalink
refactor(pdk): simplify private RL pdk
Browse files Browse the repository at this point in the history
  • Loading branch information
@ADD-SP
ADD-SP committed Jun 17, 2024
1 parent 9c0aad4 commit fd74647
Show file tree
Hide file tree
Showing 5 changed files with 276 additions and 270 deletions.
267 changes: 16 additions & 251 deletions kong/pdk/private/rate_limiting.lua
View file
Original file line number Diff line number Diff line change
@@ -1,86 +1,13 @@
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 tablex_keys = require("pl.tablex").keys

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
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",
},
}

local _M = {}

Expand Down Expand Up @@ -114,201 +41,39 @@ local function _get_or_create_rl_ctx(ngx_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)

assert(
type(limit) == "number",
"arg #2 `limit` for `set_basic_limit` must be a number"
)
assert(
type(remaining) == "number",
"arg #3 `remaining` for `set_basic_limit` must be a number"
)
function _M.store_response_header(ngx_ctx, key, value)
assert(
type(reset) == "number",
"arg #4 `reset` for `set_basic_limit` must be a number"
type(key) == "string",
"arg #2 `key` for function `store_response_header` must be a string"
)

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)

assert(
type(reset) == "number",
"arg #2 `reset` for `set_retry_after` must be a number"
)

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)

assert(
type(limit_by) == "string",
"arg #2 `limit_by` for `set_limit_by` must be a string"
)
assert(
type(limit) == "number",
"arg #3 `limit` for `set_limit_by` must be a number"
)
local value_type = type(value)
assert(
type(remaining) == "number",
"arg #4 `remaining` for `set_limit_by` must be a number"
value_type == "string" or value_type == "number",
"arg #3 `value` for function `store_response_header` must be a string or a number"
)

limit_by = LIMIT_BY[limit_by]
assert(limit_by, "invalid limit_by")

rl_ctx[limit_by.limit] = limit
rl_ctx[limit_by.remain] = remaining
local rl_ctx = _get_or_create_rl_ctx(ngx_ctx)
rl_ctx[key] = value
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)

function _M.get_stored_response_header(ngx_ctx, key)
assert(
type(limit_by) == "string",
"arg #2 `limit_by` for `set_limit_by_with_identifier` must be a string"
)
assert(
type(limit) == "number",
"arg #3 `limit` for `set_limit_by_with_identifier` must be a number"
)
assert(
type(remaining) == "number",
"arg #4 `remaining` for `set_limit_by_with_identifier` must be a number"
)

local id_seg_1_typ = type(id_seg_1)
local id_seg_2_typ = type(id_seg_2)
assert(
id_seg_1_typ == "nil" or id_seg_1_typ == "string",
"arg #5 `id_seg_1` for `set_limit_by_with_identifier` must be a string or nil"
)
assert(
id_seg_2_typ == "nil" or id_seg_2_typ == "string",
"arg #6 `id_seg_2` for `set_limit_by_with_identifier` must be a string or nil"
type(key) == "string",
"arg #2 `key` for function `get_stored_response_header` must be a string"
)

limit_by = LIMIT_BY[limit_by]
if not limit_by then
local valid_limit_bys = tablex_keys(LIMIT_BY)
local msg = string.format(
"arg #2 `limit_by` for `set_limit_by_with_identifier` must be one of: %s",
table.concat(valid_limit_bys, ", ")
)
error(msg)
if not _has_rl_ctx(ngx_ctx) then
return nil
end

id_seg_1 = id_seg_1 or ""
id_seg_2 = id_seg_2 or ""

-- construct the key like X-<id_seg_1>-RateLimit-Limit-<id_seg_2>-<limit_by>
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-<id_seg_1>-RateLimit-Remaining-<id_seg_2>-<limit_by>
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]
end

function _M.get_retry_after(ngx_ctx)
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)

assert(
type(limit_by) == "string",
"arg #2 `limit_by` for `get_limit_by` must be a string"
)

limit_by = LIMIT_BY[limit_by]
assert(limit_by, "invalid limit_by")

return rl_ctx[limit_by.limit], rl_ctx[limit_by.remain]
local rl_ctx = _get_rl_ctx(ngx_ctx)
return rl_ctx[key]
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(limit_by) == "string",
"arg #2 `limit_by` for `get_limit_by_with_identifier` must be a string"
)

local id_seg_1_typ = type(id_seg_1)
local id_seg_2_typ = type(id_seg_2)
assert(
id_seg_1_typ == "nil" or id_seg_1_typ == "string",
"arg #3 `id_seg_1` for `get_limit_by_with_identifier` must be a string or nil"
)
assert(
id_seg_2_typ == "nil" or id_seg_2_typ == "string",
"arg #4 `id_seg_2` for `get_limit_by_with_identifier` must be a string or nil"
)

limit_by = LIMIT_BY[limit_by]
if not limit_by then
local valid_limit_bys = tablex_keys(LIMIT_BY)
local msg = string.format(
"arg #2 `limit_by` for `get_limit_by_with_identifier` must be one of: %s",
table.concat(valid_limit_bys, ", ")
)
error(msg)
end

id_seg_1 = id_seg_1 or ""
id_seg_2 = id_seg_2 or ""

-- construct the key like X-<id_seg_1>-RateLimit-Limit-<id_seg_2>-<limit_by>
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-<id_seg_1>-RateLimit-Remaining-<id_seg_2>-<limit_by>
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)
function _M.apply_response_headers(ngx_ctx)
if not _has_rl_ctx(ngx_ctx) then
return
end
Expand Down
52 changes: 39 additions & 13 deletions kong/plugins/rate-limiting/handler.lua
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,17 +14,42 @@ local pairs = pairs
local error = error
local tostring = tostring
local timer_at = ngx.timer.at
local pdk_rl_set_basic_limit = pdk_private_rl.set_basic_limit
local pdk_rl_set_retry_after = pdk_private_rl.set_retry_after
local pdk_rl_set_limit_by = pdk_private_rl.set_limit_by
local pdk_rl_set_response_headers = pdk_private_rl.set_response_headers
local SYNC_RATE_REALTIME = -1


local pdk_rl_store_response_header = pdk_private_rl.store_response_header
local pdk_rl_apply_response_headers = pdk_private_rl.apply_response_headers


local EMPTY = {}
local EXPIRATION = require "kong.plugins.rate-limiting.expiration"


local RATELIMIT_LIMIT = "RateLimit-Limit"
local RATELIMIT_REMAINING = "RateLimit-Remaining"
local RATELIMIT_RESET = "RateLimit-Reset"
local RETRY_AFTER = "Retry-After"


local X_RATELIMIT_LIMIT = {
second = "X-RateLimit-Limit-Second",
minute = "X-RateLimit-Limit-Minute",
hour = "X-RateLimit-Limit-Hour",
day = "X-RateLimit-Limit-Day",
month = "X-RateLimit-Limit-Month",
year = "X-RateLimit-Limit-Year",
}

local X_RATELIMIT_REMAINING = {
second = "X-RateLimit-Remaining-Second",
minute = "X-RateLimit-Remaining-Minute",
hour = "X-RateLimit-Remaining-Hour",
day = "X-RateLimit-Remaining-Day",
month = "X-RateLimit-Remaining-Month",
year = "X-RateLimit-Remaining-Year",
}


local RateLimitingHandler = {}


Expand Down Expand Up @@ -126,7 +151,6 @@ function RateLimitingHandler:access(conf)

if usage then
local ngx_ctx = ngx.ctx
-- Adding headers
local reset
if not conf.hide_client_headers then
local timestamps
Expand Down Expand Up @@ -157,21 +181,23 @@ function RateLimitingHandler:access(conf)
reset = max(1, window - floor((current_timestamp - timestamps[k]) / 1000))
end

pdk_rl_set_limit_by(ngx_ctx, k, limit, current_remaining)
pdk_rl_store_response_header(ngx_ctx, X_RATELIMIT_LIMIT[k], current_limit)
pdk_rl_store_response_header(ngx_ctx, X_RATELIMIT_REMAINING[k], current_remaining)
end

pdk_rl_set_basic_limit(ngx_ctx, limit, remaining, reset)
pdk_rl_store_response_header(ngx_ctx, RATELIMIT_LIMIT, limit)
pdk_rl_store_response_header(ngx_ctx, RATELIMIT_REMAINING, remaining)
pdk_rl_store_response_header(ngx_ctx, RATELIMIT_RESET, reset)
end

-- If limit is exceeded, terminate the request
if stop then
pdk_rl_set_retry_after(ngx_ctx, reset)
pdk_rl_set_response_headers(ngx_ctx)
return kong.response.error(conf.error_code, conf.error_message)
pdk_rl_store_response_header(ngx_ctx, RETRY_AFTER, reset)
pdk_rl_apply_response_headers(ngx_ctx)
return kong.response.error(conf.error_code, conf.error_message, headers)
end

-- Set rate-limiting response headers
pdk_rl_set_response_headers(ngx_ctx)
pdk_rl_apply_response_headers(ngx_ctx)
end

if conf.sync_rate ~= SYNC_RATE_REALTIME and conf.policy == "redis" then
Expand All @@ -186,4 +212,4 @@ function RateLimitingHandler:access(conf)
end


return RateLimitingHandler
return RateLimitingHandler
Loading

0 comments on commit fd74647

Please sign in to comment.