--[[
執行過載限流策略
--]]

-- 當執行限流時，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

resty-limit-multiple-strategy.lua