Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor(pdk): simplify private RL pdk #13213

Merged
merged 1 commit into from
Jun 17, 2024
Merged

refactor(pdk): simplify private RL pdk #13213

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
267 changes: 16 additions & 251 deletions kong/pdk/private/rate_limiting.lua
View file
Open in desktop
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
50 changes: 38 additions & 12 deletions kong/plugins/rate-limiting/handler.lua
View file
Open in desktop
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)
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)
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
Loading