Skip to content

Commit 46d57b4

Browse files
travisbelltravisbell
committed
Add basic support to skip whitelisted API keys.
1 parent 5b3788b commit 46d57b4

File tree

2 files changed

+64
-48
lines changed

2 files changed

+64
-48
lines changed

‎README.md

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,12 +43,15 @@ server {
4343
rate = 40,
4444
interval = 10,
4545
log_level = ngx.NOTICE,
46-
redis_config = { host = "127.0.0.1", port = 6379, timeout = 1, pool_size = 100 } }
46+
redis_config = { host = "127.0.0.1", port = 6379, timeout = 1, pool_size = 100 },
47+
whitelisted_api_keys = { 'XXX', 'ZZZ' } }
4748
';
4849
4950
proxy_set_header Host $host;
51+
proxy_set_header X-Server-Scheme $scheme;
5052
proxy_set_header X-Real-IP $remote_addr;
5153
proxy_set_header X-Forwarded-For $remote_addr;
54+
proxy_set_header X-Forwarded-Proto $x_forwarded_proto;
5255
5356
proxy_connect_timeout 1s;
5457
proxy_read_timeout 30s;

‎lib/resty/rate/limit.lua

Lines changed: 60 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ end
1212

1313
local function bump_request(redis_connection, redis_pool_size, ip_key, rate, interval, current_time, log_level)
1414
local key = "RL:" .. ip_key
15-
15+
1616
local count, error = redis_connection:incr(key)
1717
if not count then
1818
ngx.log(log_level, "failed to incr count: ", error)
@@ -46,62 +46,75 @@ local function bump_request(redis_connection, redis_pool_size, ip_key, rate, int
4646
end
4747

4848
function _M.limit(config)
49-
local log_level = config.log_level or ngx.ERR
49+
local uri_parameters = ngx.req.get_uri_args()
50+
local enforce_limit = true
5051

51-
if not config.connection then
52-
local ok, redis = pcall(require, "resty.redis")
53-
if not ok then
54-
ngx.log(log_level, "failed to require redis")
55-
return
52+
for i, api_key in ipairs(config.whitelisted_api_keys) do
53+
if api_key == uri_parameters.api_key then
54+
enforce_limit = false
5655
end
56+
end
5757

58-
local redis_config = config.redis_config or {}
59-
redis_config.timeout = redis_config.timeout or 1
60-
redis_config.host = redis_config.host or "127.0.0.1"
61-
redis_config.port = redis_config.port or 6379
62-
redis_config.pool_size = redis_config.pool_size or 100
58+
if enforce_limit then
59+
local log_level = config.log_level or ngx.ERR
60+
61+
if not config.connection then
62+
local ok, redis = pcall(require, "resty.redis")
63+
if not ok then
64+
ngx.log(log_level, "failed to require redis")
65+
return
66+
end
67+
68+
local redis_config = config.redis_config or {}
69+
redis_config.timeout = redis_config.timeout or 1
70+
redis_config.host = redis_config.host or "127.0.0.1"
71+
redis_config.port = redis_config.port or 6379
72+
redis_config.pool_size = redis_config.pool_size or 100
73+
74+
local redis_connection = redis:new()
75+
redis_connection:set_timeout(redis_config.timeout * 1000)
76+
77+
local ok, error = redis_connection:connect(redis_config.host, redis_config.port)
78+
if not ok then
79+
ngx.log(log_level, "failed to connect to redis: ", error)
80+
return
81+
end
82+
83+
config.redis_config = redis_config
84+
config.connection = redis_connection
85+
end
6386

64-
local redis_connection = redis:new()
65-
redis_connection:set_timeout(redis_config.timeout * 1000)
87+
local current_time = ngx.now()
88+
local connection = config.connection
89+
local redis_pool_size = config.redis_config.pool_size
90+
local key = config.key or ngx.var.remote_addr
91+
local rate = config.rate or 10
92+
local interval = config.interval or 1
6693

67-
local ok, error = redis_connection:connect(redis_config.host, redis_config.port)
68-
if not ok then
69-
ngx.log(log_level, "failed to connect to redis: ", error)
94+
local response, error = bump_request(connection, redis_pool_size, key, rate, interval, current_time, log_level)
95+
if not response then
7096
return
7197
end
7298

73-
config.redis_config = redis_config
74-
config.connection = redis_connection
75-
end
76-
77-
local current_time = ngx.now()
78-
local connection = config.connection
79-
local redis_pool_size = config.redis_config.pool_size
80-
local key = config.key or ngx.var.remote_addr
81-
local rate = config.rate or 10
82-
local interval = config.interval or 1
83-
84-
local response, error = bump_request(connection, redis_pool_size, key, rate, interval, current_time, log_level)
85-
if not response then
86-
return
87-
end
88-
89-
if response.count > rate then
90-
local retry_after = math.floor(response.reset - current_time)
91-
if retry_after < 0 then
92-
retry_after = 0
99+
if response.count > rate then
100+
local retry_after = math.floor(response.reset - current_time)
101+
if retry_after < 0 then
102+
retry_after = 0
103+
end
104+
105+
ngx.header["Access-Control-Allow-Origin"] = "*"
106+
ngx.header["Content-Type"] = "application/json; charset=utf-8"
107+
ngx.header["Retry-After"] = retry_after
108+
ngx.status = 429
109+
ngx.say('{"status_code":25,"status_message":"Your request count (' .. response.count .. ') is over the allowed limit of ' .. rate .. '."}')
110+
ngx.exit(ngx.HTTP_OK)
111+
else
112+
ngx.header["X-RateLimit-Limit"] = rate
113+
ngx.header["X-RateLimit-Remaining"] = math.floor(response.remaining)
114+
ngx.header["X-RateLimit-Reset"] = math.floor(response.reset)
93115
end
94-
95-
ngx.header["Access-Control-Allow-Origin"] = "*"
96-
ngx.header["Content-Type"] = "application/json; charset=utf-8"
97-
ngx.header["Retry-After"] = retry_after
98-
ngx.status = 429
99-
ngx.say('{"status_code":25,"status_message":"Your request count (' .. response.count .. ') is over the allowed limit of ' .. rate .. '."}')
100-
ngx.exit(ngx.HTTP_OK)
101116
else
102-
ngx.header["X-RateLimit-Limit"] = rate
103-
ngx.header["X-RateLimit-Remaining"] = math.floor(response.remaining)
104-
ngx.header["X-RateLimit-Reset"] = math.floor(response.reset)
117+
return
105118
end
106119
end
107120

0 commit comments

Comments
 (0)