151 lines
4.1 KiB
Lua
151 lines
4.1 KiB
Lua
-- IP频率限制核心模块
|
||
-- 功能：统计指定IP在时间段内的请求次数，根据配置执行相应处置策略
|
||
|
||
local _M = {}
|
||
|
||
local common_config = require "config.common_config"
|
||
local logger = require "lib.logger"
|
||
|
||
-- 获取配置（直接从公共配置读取）
|
||
local function get_config()
|
||
return common_config.rate_limit
|
||
end
|
||
|
||
-- 检查域名是否应该应用限流
|
||
local function should_apply_rate_limit(host)
|
||
local cfg = get_config()
|
||
|
||
-- 如果未启用，直接返回false
|
||
if not cfg.enabled then
|
||
return false
|
||
end
|
||
|
||
-- 如果没有host，默认不应用
|
||
if not host or host == "" then
|
||
return false
|
||
end
|
||
|
||
-- 从公共配置获取域名列表
|
||
local domains = common_config.domains
|
||
|
||
-- 如果domains列表为空，对所有域名生效
|
||
if #domains == 0 then
|
||
return true
|
||
end
|
||
|
||
-- 检查域名是否在配置的列表中
|
||
for _, domain in ipairs(domains) do
|
||
if host == domain then
|
||
return true
|
||
end
|
||
end
|
||
|
||
return false
|
||
end
|
||
|
||
-- 主函数：检查并限制IP请求频率
|
||
-- 返回值：
|
||
-- - true: 允许访问
|
||
-- - false: 拒绝访问（已执行相应处置）
|
||
function _M.check_rate_limit()
|
||
local cfg = get_config()
|
||
|
||
-- 如果未启用，直接放行
|
||
if not cfg.enabled then
|
||
return true
|
||
end
|
||
|
||
-- 获取当前域名
|
||
local host = ngx.var.host
|
||
|
||
-- 检查是否应该对此域名应用限流
|
||
if not should_apply_rate_limit(host) then
|
||
logger.log_info("Rate limit skipped for domain", {
|
||
domain = host,
|
||
reason = "not_in_domains_list"
|
||
})
|
||
return true
|
||
end
|
||
|
||
-- 获取共享内存字典
|
||
local dict = ngx.shared[cfg.shared_dict]
|
||
if not dict then
|
||
logger.log_error("Shared dictionary not found", {
|
||
dict_name = cfg.shared_dict
|
||
})
|
||
return true
|
||
end
|
||
|
||
-- 获取客户端IP
|
||
local client_ip = ngx.var.remote_addr
|
||
if not client_ip then
|
||
logger.log_warn("Failed to get client IP")
|
||
return true
|
||
end
|
||
|
||
-- 记录请求详情（可选）
|
||
logger.log_request_details(
|
||
client_ip,
|
||
ngx.var.request_uri,
|
||
ngx.var.request_method,
|
||
ngx.var.http_user_agent
|
||
)
|
||
|
||
local current_time = ngx.time()
|
||
local time_window = cfg.time_window
|
||
local max_requests = cfg.max_requests
|
||
|
||
-- 构建key：使用时间窗口的起始时间戳
|
||
local window_start = math.floor(current_time / time_window) * time_window
|
||
local key = client_ip .. ":" .. window_start
|
||
|
||
-- 使用原子操作增加计数
|
||
local new_count, err = dict:incr(key, 1, 0)
|
||
if not new_count then
|
||
logger.log_error("Failed to increment counter", {
|
||
ip = client_ip,
|
||
error = err or "unknown"
|
||
})
|
||
return true
|
||
end
|
||
|
||
-- 设置过期时间
|
||
if new_count == 1 then
|
||
dict:expire(key, time_window + 10)
|
||
end
|
||
|
||
-- 检查是否超过限制
|
||
if new_count > max_requests then
|
||
-- 根据配置的处置策略执行相应操作
|
||
if cfg.action == "allow" then
|
||
-- 放行：仅记录日志，不阻止
|
||
logger.log_allowed(client_ip, new_count, time_window)
|
||
return true
|
||
|
||
elseif cfg.action == "rate_limit" then
|
||
-- 限制速率：返回429
|
||
logger.log_rate_limited(client_ip, new_count, max_requests, time_window)
|
||
|
||
ngx.status = cfg.status_code or 429
|
||
ngx.header["Content-Type"] = "text/plain; charset=utf-8"
|
||
ngx.say(cfg.message or "Rate limit exceeded")
|
||
return false
|
||
|
||
elseif cfg.action == "block" then
|
||
-- 封禁：返回403
|
||
logger.log_blocked(client_ip, new_count, max_requests, time_window)
|
||
|
||
ngx.status = cfg.status_code or 403
|
||
ngx.header["Content-Type"] = "text/plain; charset=utf-8"
|
||
ngx.say(cfg.message or "Access denied")
|
||
return false
|
||
end
|
||
end
|
||
|
||
-- 正常放行
|
||
logger.log_allowed(client_ip, new_count, time_window)
|
||
return true
|
||
end
|
||
|
||
return _M
|