limit_api_flow.lua
3.38 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
require "resty.core"
local cjson = require "cjson"
local http_request = require "http_request"
local iptool=require "iptool"
local default_err_code=9999991
local default_err_msg="系统正忙,请稍后重试!"
-- read config
local common_config=lua_context.configs["common_conf"]
local limit_config=lua_context.configs["api_rate_limit_conf"]
local limit_ip_config=lua_context.configs["limit_ip_access"]
--- do limit using local cache
function limit(limit_key,max_limit,seconds)
if not limit_key or not max_limit then
return true
end
if not seconds then
seconds=2
end
lua_context.lua_conf_cache:safe_add(limit_key,0,seconds)
local limit,err=lua_context.lua_conf_cache:incr(limit_key,1)
if limit then
if tonumber(limit)>max_limit then
return false
else
return true
end
end
return true
end
-------- function: doing rate limit -----
function rate_limit()
if (not common_config) or (not limit_config) then
return
end
if not common_config.lua_golbal_switch then
return
end
-- get method
local req_uri_method = http_request.get_method()
if (not limit_config.is_open) or (not req_uri_method) then
return
end
-- get max QPS from config by method
local api_rate_limit=limit_config.api_rate_limit
local max_per_sencond=limit_config.default_rate_limit
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
if not max_per_sencond then
max_per_sencond=60
end
-- get error code & error message from config by method
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
-- do limit
local flag= limit("yh:nginx:limitflow:" .. req_uri_method, max_per_sencond)
if not flag then
ngx.log(ngx.ERR,"The request is in 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
--check is in white ip list
function is_white()
if limit_ip_config then
local white_ips_length=#limit_ip_config.white_ips
if white_ips_length >0 then
for i=1,white_ips_length do
local is_in_white_ips=iptool:pcall_check_ip_in_ipblock(ngx.var.real_ip,limit_ip_config.white_ips[i],false)
if is_in_white_ips then
return true
end
end
end
end
return false
end
----- check weather ip is in black list
function check_malIp()
local cache=lua_context.mal_ip_cache
local ip=ngx.var.real_ip
local exist = cache:get("yh:mip:" .. ip)
if exist then
ngx.log(ngx.ERR, "The request ip is in blacklist:" .. ip)
ngx.exit(ngx.HTTP_FORBIDDEN)
end
end
---------end check_malIp()-----------
---- function read request body
function init_read_body()
if ngx.var.request_method=="POST" and ngx.var.content_type and string.match(ngx.var.content_type,"application/x%-www%-form%-urlencoded.*") then
ngx.req.read_body()
end
end
function main()
if is_white() then
return
else
init_read_body()
check_malIp()
rate_limit()
end
end
----- running ----
main()