limit_api_flow.lua 2.51 KB

local rateLimit = require "limit_common_flow"
local cjson = require "cjson"



local open_limit_flow=true
local api_default_max_per_sencond=60
local service_default_max_per_sencond=60
local default_err_code=9999991
local default_err_msg="系统正忙,请稍后重试!"

function get_req_param(req_param)        
  local method = ngx.var.request_method
  if "GET" == method then
    return ngx.req.get_uri_args()[req_param]
  else
    ngx.req.read_body()
    return ngx.req.get_post_args()[req_param]
  end
end

function  get_req_uri()
  local limit_config=lua_context.configs["api_rate_limit_conf"]
  local req_service_url = ngx.var.uri
  if (not limit_config.is_open) or (not req_service_url) then
    return nil
  end
  local beginIndex, endIndex = string.find(req_service_url, "apigateway")
  if endIndex == nil then
    return nil
  end
  local real_url=string.sub(req_service_url,endIndex+1,string.len(req_service_url))
  return real_url
end


function extract_limit_method()
   
    local uri_method=get_req_uri()
	
	if uri_method and uri_method ~= "/" and uri_method ~= ""  then
	   return uri_method
	end

    local  method= get_req_param("method") 
    return method
end

function rate_limit()
  local common_config=lua_context.configs["common_conf"]
  if not common_config.lua_golbal_switch then
     return 
  end

  local limit_config=lua_context.configs["api_rate_limit_conf"]
  local api_rate_limit=limit_config.api_rate_limit
  local req_uri_method = extract_limit_method() 
  --ngx.log(ngx.INFO,"=================>>" .. cjson.encode(api_rate_limit[req_uri_method]))
  if (not limit_config.is_open) or (not req_uri_method) then
    return
  end
  
  local max_per_sencond=api_default_max_per_sencond
  if api_rate_limit[req_uri_method] and api_rate_limit[req_uri_method][1] then
    max_per_sencond=api_rate_limit[req_uri_method][1]
  end

  local err_code=default_err_code
  if api_rate_limit[req_uri_method] and api_rate_limit[req_uri_method][2] then
    err_code=api_rate_limit[req_uri_method][2]
  end
  
  local err_msg=default_err_msg
  if api_rate_limit[req_uri_method] and api_rate_limit[req_uri_method][3] then
    err_msg=api_rate_limit[req_uri_method][3]
  end
  
  local flag=rateLimit.limit_flow("yh:nginx:limitflow:" .. req_uri_method,max_per_sencond)
  if not flag then
    ngx.log(ngx.ERR,"The request is limited :" .. req_uri_method)
    ngx.header["Content-Type"]="application/json;charset=utf-8"
    local msg='{"code":' .. err_code .. ',"message":"'.. err_msg .. '"}'
    ngx.say(msg)
    ngx.exit(ngx.HTTP_OK)
  end
end


rate_limit()