bunkerweb/lua/limitreq.lua

local M		= {}
local logger	= require "logger"

function M.decr (key, delay)
	local function callback (premature, key)
		if premature then
			ngx.shared.limit_req:delete(key)
			return
		end
		local value, flags = ngx.shared.limit_req:get(key)
		if value ~= nil then
			if value - 1 == 0 then
				ngx.shared.limit_req:delete(key)
				return
			end
			ngx.shared.limit_req:set(key, value-1, 0)
		end
	end
	local hdl, err = ngx.timer.at(delay, callback, key)
	if not ok then
		logger.log(ngx.ERR, "REQ LIMIT", "can't setup decrement timer : " .. err)
		return false
	end
	return true
end

function M.incr (key)
	local newval, err, forcible = ngx.shared.limit_req:incr(key, 1, 0, 0)
	if not newval then
		logger.log(ngx.ERR, "REQ LIMIT", "can't increment counter : " .. err)
		return false
	end
	return true
end

function M.check (rate, burst, sleep)
	local key = ngx.var.remote_addr .. ngx.var.uri
	local rate_split = rate:gmatch("([^r/]+)")
	local max = rate_split[1]
	local unit = rate_split[2]
	local delay = 0
	if unit == "s" then
		delay = 1
	elseif unit == "m" then
		delay = 60
	elseif unit == "h" then
		delay = 3600
	elseif unit == "d" then
		delay = 86400
	end
	if M.incr(key) then
		local current, flags = ngx.shared.limit_req:get(key)
		if M.decr(key, delay) then
			if current > max + burst then
				logger.log(ngx.WARN, "REQ LIMIT", "ip " .. ngx.var.remote_addr .. " has reached the limit for uri " .. ngx.var.uri .. " : " .. current .. "r/" .. unit .. " (max = " .. rate .. ")")
				return true
			elseif current > max then
				if sleep > 0 then
					ngx.sleep(sleep)
				end
			end
		else
			ngx.shared.limit_req:set(key, current-1, 0)
		end
	end
	return false	
end

return M