diff --git a/src/access.lua b/src/access.lua index 5fe9098..f80a76b 100755 --- a/src/access.lua +++ b/src/access.lua @@ -77,7 +77,10 @@ function _M.dispatch() ngx.exit(404) end - -- TODO:限流,IP黑名单... + -- TODO:IP黑名单... + -- 限流 + local al = require "access_limit" + al.checkAccessLimit(api_info["request_uri"], api_info["uri_limit_seconds"], api_info["uri_limit_times"], api_info["ip_uri_limit_seconds"], api_info["ip_uri_limit_times"]) -- 负载均衡,选择服务器开始 local server = select_server(api_info); diff --git a/src/access_limit.lua b/src/access_limit.lua new file mode 100644 index 0000000..2ba4385 --- /dev/null +++ b/src/access_limit.lua @@ -0,0 +1,69 @@ +local config = require "config" +local _M = {} +local log = ngx.log +local ERR = ngx.ERR + +function open_redis() + local redis = require "resty.redis" + local red = redis:new() + red:set_timeout(1000) -- 1 sec + local ok, err = red:connect(config["redis_host"], config["redis_port"]) + if not ok then + log(ERR, "failed to connect: ", err) + end + return red, err +end + +function return_redis(red) + -- put it into the connection pool of size 100, + -- with 10 seconds max idle time + local ok, err = red:set_keepalive(10000, 100) + if not ok then + log(ERR, "failed to set keepalive: ", err) + return + end +end + +function _M.checkAccessLimit(uri, uriSeconds, uriTimes, ipUriSeconds, ipUriTimes) + if not uriSeconds or not ipUriSeconds then + return + end + + local red, err = open_redis() + if err then + return + end + + -- 针对整个接口限流 + if uriSeconds then + local value, _ = red:get(uri) + if value then + if value >= uriTimes then + return_redis() + ngx.exit(403) + end + red:incr(uri) + else + red:setex(uri, uriSeconds, 1) + end + end + + -- 针对IP接口限流 + if ipUriSeconds then + local headers = ngx.req.get_headers() + local ip = headers["X-REAL-IP"] or headers["X_FORWARDED_FOR"] or ngx.var.remote_addr or "0.0.0.0" + local value, _ = red:get(ip .. uri) + if value then + if value >= ipUriTimes then + return_redis() + ngx.exit(403) + end + red:incr(ip .. uri) + else + red:setex(ip .. uri, ipUriSeconds, 1) + end + end + return_redis() +end + +return _M \ No newline at end of file diff --git a/src/config.lua b/src/config.lua index a421fd6..bc33f2f 100755 --- a/src/config.lua +++ b/src/config.lua @@ -1,10 +1,14 @@ -local config={}; +local config = {}; config['mysql_host'] = "127.0.0.1" config['mysql_port'] = "3306" config['mysql_user'] = "root" config['mysql_pass'] = "" config['mysql_db'] = "agw" +-- 如果开启限流需要正确配置redis +config["redis_host"] = "127.0.0.1" +config["redis_port"] = 6379 + config['admin_name'] = "admin" config['admin_pass'] = "admin"