Skip to content

Instantly share code, notes, and snippets.

@shasharoman
Last active Aug 22, 2021
Embed
What would you like to do?
openresty-rate-limit
local limit = require 'limit'
local ignoreIps = {}
local ignoreKeys = {}
local headers = ngx.req.get_headers()
local key = headers['api-key'] or ''
local ip = string.match(headers['x-forwarded-for'] or ngx.var.remote_addr, '%d+.%d+.%d+.%d+')
for i, v in ipairs(ignoreKeys) do
if v == key then
return
end
end
for i, v in ipairs(ignoreIps) do
if v == ip then
return
end
end
limit.enter()
local limit = require 'limit'
local json = require 'cjson'
local file, err = io.open('limit.json', 'r')
if not file then
ngx.log(ngx.ERR, err)
end
local rules = file:read('*all')
io.close(file)
limit.init({
shm = {
req = 'limit_req_store',
count = 'limit_count_store',
conn = 'limit_conn_store'
},
rules = json.decode(rules),
conn_limit = 128,
conn_burst = 0,
})
{
"POST:/api/v1/order": [10, 5, 120, 60],
"GET:/api/v1/order/items": [10, 5, 300, 60]
}
local _M = {}
local limit_req = require 'resty.limit.req'
local limit_count = require 'resty.limit.count'
local limit_conn = require 'resty.limit.conn'
local shm = nil
local function log(...)
ngx.log(ngx.ERR, 'limit: ', ...)
end
local function split(s, delimiter)
result = {}
for match in (s..delimiter):gmatch('(.-)'..delimiter) do
table.insert(result, match)
end
return result
end
-- https://github.com/openresty/lua-resty-limit-traffic/blob/master/lib/resty/limit/traffic.lua
local function combine(limiters, keys, states)
local n = #limiters
local max_delay = 0
for i = 1, n do
local lim = limiters[i]
local delay, err = lim:incoming(keys[i], i == n)
if states then
states[i] = {delay, err}
end
if not delay then
return nil, err
end
if i == n then
max_delay = delay
end
end
for i = 1, n - 1 do
local lim = limiters[i]
local delay, err = lim:incoming(keys[i], true)
if states then
states[i] = {delay, err}
end
if not delay then
for j = 1, i - 1 do
-- we intentionally ignore any errors returned below.
limiters[j]:uncommit(keys[j])
end
limiters[n]:uncommit(keys[n])
return nil, err
end
max_delay = math.max(max_delay, delay)
end
return max_delay
end
function _M.init(opts)
shm = ngx.shared.limit
shm:set('req_shm_name', opts.shm.req or 'limit_req_store')
shm:set('count_shm_name', opts.shm.count or 'limit_count_store')
shm:set('conn_shm_name', opts.shm.conn or 'limit_conn_store')
if ngx.worker.id() == 0 then
shm:set('conn_limit', opts.conn_limit or 128)
shm:set('conn_burst', opts.conn_burst or 0)
for k, v in pairs(opts.rules) do
shm:set(k, table.concat(v, '-'))
log('rules ' .. k .. ', ' .. table.concat(v, '-'))
end
end
end
function _M.enter()
local key = ngx.var.request_method .. ':' .. ngx.var.uri
local limit = shm:get(key)
if not limit then
return
end
limit = split(limit, '-')
local req_lim, req_err = limit_req.new(shm:get('req_shm_name'), tonumber(limit[1]), tonumber(limit[2]))
local count_lim, count_err = limit_count.new(shm:get('count_shm_name'), tonumber(limit[3]), tonumber(limit[4]))
local conn_lim, conn_err = limit_conn.new(shm:get('conn_shm_name'), shm:get('conn_limit'), shm:get('conn_burst'), 0.5)
if req_err or count_err or conn_err then
log(req_err, count_err, conn_err)
return
end
local headers = ngx.req.get_headers()
local client_ip = headers['x-forwarded-for'] or ngx.var.remote_addr
client_ip = split(client_ip, ',')[1]
local limit_key = client_ip .. key
local limiters = {req_lim, count_lim, conn_lim}
local keys = {limit_key, limit_key, client_ip}
local states = {}
local delay, err = combine(limiters, keys, states)
-- rate limit request per second
ngx.header['x-rl-rps'] = limit[1]
-- rate limit window time
ngx.header['x-rl-wt'] = limit[4]
-- rate limit window limit
ngx.header['x-rl-wl'] = limit[3]
if not delay then
local reason = err
if err == 'rejected' then
if states then
if states[1] and states[1][2] == 'rejected' then
reason = 'The number of requests per second cannot exceed x-rl-rps:' .. limit[1] .. '.'
elseif states[2] and states[2][2] == 'rejected' then
reason = 'The number of requests per ' .. limit[4] .. ' seconds cannot exceed x-rl-wl:' .. limit[3] .. '.'
elseif states[3] and states[3][2] == 'rejected' then
reason = 'The number of concurrent requests cannot exceed ' .. limit[5] .. '.'
end
else
log('rejected but no states')
reason = 'rejected.'
end
end
log('rejected: ', client_ip, ' reason: ', reason)
ngx.status = 429
ngx.say(reason)
return ngx.exit(ngx.OK)
end
-- rate limit window remaining
if states and states[2] then
ngx.header['x-rl-wr'] = states[2][2]
end
-- for conn leaving
if conn_lim:is_committed() then
ngx.ctx.conn_limit = conn_lim
ngx.ctx.conn_limit_key = client_ip
end
if delay > 0.001 then
ngx.header['x-rl-delay'] = delay
log('sleeping: ' .. client_ip, ' delay: ', delay)
ngx.sleep(delay)
end
end
function _M.leave()
local limit = ngx.ctx.conn_limit
if limit then
local conn, err = limit:leaving(ngx.ctx.conn_limit_key, tonumber(ngx.var.request_time))
if not conn then
log('failed to record the connection leaving, request: ', err)
return
end
end
end
return _M
local limit = require 'limit'
limit.leave()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment