增加限流

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);

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

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"