resty-limit-multiple-strategy.lua

Posted 2020-11-12 koushr

--[[
执行过载限流策略
--]]

-- 当执行限流时，nginx 返回的状态码
err_code = 403

local limit_config = {
  user_limit = {rate = 15, brust = 5},     --用户限流维度桶设置
  priority_uri_limit = {      --指定 URI 限流维度桶设置
    [‘/elis_smp_portal_dmz/do/cms/getLittleRedDotInfo‘] = {rate = 400, brust = 50},
    [‘/elis_smp_portal_dmz/do/esa/portal.msgcenter.unreadv2‘] = {rate = 400, brust = 50},
    [‘/elis_smp_portal_dmz/do/home/getUpgradeReminder‘] = {rate = 400, brust = 50},
    [‘/elis_smp_portal_dmz/do/userLogin/checkPluginVersion‘] = {rate = 400, brust = 50},
    [‘/elis_smp_portal_dmz/do/createRYMSession‘] = {rate = 400, brust = 50},
    [‘/elis_smp_portal_dmz/do/zone/queryZoneConfig‘] = {rate = 400, brust = 50}
  },
  error_status = err_code,     --请求被限流时返回状态码
  nodelay = no    --是否需要不延迟处理
}

--定义限流策略方法名称
local limit_user_method = "limit_by_user"
local limit_uri_method = "limit_by_priority_uri"

local limit_req = require "resty.limit.req"

ngx.req.read_body()

--过载保护策略总开关，若开关关闭，则全部策略失效
local limit_req_store = ngx.shared.limit_req_store
local overload_protection_switch = limit_req_store:get("overload_protection_switch")
if (overload_protection_switch ~= nil and overload_protection_switch == "N") then
        ngx.log(ngx.INFO, "nginx limit strategy has been shutdown.")
    return
end

-- 获取当前请求 URI
local cur_uri = ngx.var.uri

-- 执行过载操作
local is_in_black_list = limit_req_store:get("black_list:"..cur_uri)
local is_rejected_uri = reject_uri_tab["rejected"..cur_uri]
if (is_in_black_list ~= nil or is_rejected_uri ~= nil) then
  -- 拦截比例，默认全部拦截
  local request_limit_percent = limit_req_store:get("reject_request_percent")
  if request_limit_percent == nil then
    request_limit_percent = 100
  end
  local random_num = math.random(100)
  if random_num <= tonumber(request_limit_percent)  then
      ngx.log(ngx.ERR,"nginx limit strategy current uri=",cur_uri," has been rejected.")
      ngx.exit(ngx.HTTP_FORBIDDEN)
  end
end

-- 执行限流策略，入参：key 和 commit,commit 默认为 true
local function execute_limit_strategy(lim, key, commit, method)
  --  请求流入，如果你的请求需要被延迟则返回delay>0
  local delay, err = lim:incoming(key, commit)
  if (not delay and err == "rejected") then
    ngx.log(ngx.ERR, "nginx limit strategy: ",method,", key: ", key, " request rate was exceeded, current request was rejected.")
    ngx.exit(limit_config.error_status)
  end

  -- 根据需要决定延迟或者不延迟处理
  if delay > 0 then
    if limit_config.nodelay then
      -- 直接突发处理
      -- ngx.log(ngx.ERR, "nginx limit strategy: ", method, ", key: ", key, " request rate was exceeded, current request was rejected. delay")
      -- ngx.exit(error_status)
    else
      --延迟处理
      ngx.sleep(delay)
    end
  end
end

-- 获取限流策略的共享内存
local function get_shared_dict(shared_dict_name, rate, burst)
  local limit_shared_dict, err = limit_req.new(shared_dict_name, rate, burst)
  if not limit_shared_dict then --没定义共享字典
    ngx.log(ngx.ERR, "Nginx limit shared dict:", shared_dict_name, " has not been set.")
    ngx.exit(error_status)
  end
  return limit_shared_dict
end


-- 初始化 URL 限流策略 table
local uri_limit_tab = {}

local user_lim

if limit_config == nil then
  ngx.log(ngx.ERR, "nginx request limit has no config info.")
  return
else
  local user_limit_config = limit_config.user_limit
  ngx.log(ngx.DEBUG,"limit config user rate:", user_limit_config.rate, ", brust:", user_limit_config.brust)
  user_lim = get_shared_dict("limit_req_store",user_limit_config.rate, user_limit_config.brust)
  local priority_uri_limit_config = limit_config.priority_uri_limit
  for key, value in pairs(priority_uri_limit_config) do
    local rate = value.rate
    local brust = value.brust
    ngx.log(ngx.DEBUG,"limit config uri:", key, ", rate:", rate,", brust:",brust)
    local lim = get_shared_dict("limit_req_store",rate,brust)
    uri_limit_tab[key]=lim
  end
end

-- 执行指定 URI 限流策略
if uri_limit_tab[cur_uri] ~= nil then
  execute_limit_strategy(uri_limit_tab[cur_uri], cur_uri , true , limit_uri_method)
end

-- 获取客户端请求 IP
local remote_ip = ngx.var.remote_addr
local user_id
local device_id
local key_by_user

--  获取用户ID 或 设备号
local request_method = ngx.var.request_method
if "GET" == request_method then
  user_id = ngx.req.get_uri_args()["uid"]
  device_id = ngx.req.get_uri_args()["deviceId"]
  ngx.log(ngx.DEBUG, "nginx request limit module GET parameter user_id: ",user_id, " header:",device_id)
elseif "POST" == request_method then
  -- 获取请求头信息，如果缺少请求头信息判断，遇到上传文件操作时，会引起ngx.req.get_post_args()操作异常。
  local receive_headers = ngx.req.get_headers()
  local body_data = ngx.req.get_body_data()
  if not body_data or (receive_headers["content-type"] ~= nil and string.sub(receive_headers["content-type"],1,20) == "multipart/form-data;") then
    ngx.log(ngx.WARN,"nginx request limit module uri: ",cur_uri, " header:",receive_headers["content-type"],", body data:",body_data)
    return
  end
  user_id = ngx.req.get_post_args()[‘uid‘];
  device_id = ngx.req.get_post_args()[‘deviceId‘];
  ngx.log(ngx.DEBUG, "nginx request limit module POST parameter user_id: ",user_id, " device_id:",device_id)
end

-- 调试日志
ngx.log(ngx.ERR, "nginx request limit module cur_uri:",cur_uri, ", user_id: ",user_id)

if (user_id ~= nil or device_id ~= nil) then
  if user_id ~= nil then
    key_by_user = user_id
  else
    key_by_user = device_id
  end
  execute_limit_strategy(user_lim,key_by_user..":"..remote_ip, true, limit_user_method)
end