Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: Add suspicious_ips_maxsize, suspicious_ips_threshold_ratio
Browse files Browse the repository at this point in the history
jopemachine committed Apr 30, 2024
1 parent d075b8e commit f8caa97
Showing 4 changed files with 27 additions and 10 deletions.
2 changes: 1 addition & 1 deletion changes/2083.feature.md
Original file line number Diff line number Diff line change
@@ -1 +1 @@
truncate `suspicious_ips` when it exceed the max_size
truncate `hot_clients_ips` when it exceed the max_size
1 change: 1 addition & 0 deletions changes/2087.feature.md
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Add `hot_clients_ips_maxsize`, `hot_clients_ips_threshold_ratio`
27 changes: 19 additions & 8 deletions src/ai/backend/manager/api/ratelimit.py
Original file line number Diff line number Diff line change
@@ -50,15 +50,15 @@
if id_type == "ip" then
local rate_limit = tonumber(ARGV[3])
local score_threshold = rate_limit * 0.8
local hot_clients_ips_maxsize = tonumber(ARGV[4])
local hot_clients_ips_threshold_ratio = tonumber(ARGV[5])
-- Add IP to hot_clients_ips only if count is greater than score_threshold
if rolling_count >= score_threshold then
-- Add the IP address to "hot_clients_ips" only if rolling_count is greater than the threshold
if rolling_count >= rate_limit * hot_clients_ips_threshold_ratio then
redis.call('ZADD', 'hot_clients_ips', rolling_count, id_value)
local max_size = 1000
local current_size = redis.call('ZCARD', 'hot_clients_ips')
if current_size > max_size then
if current_size > hot_clients_ips_maxsize then
redis.call('ZREMRANGEBYRANK', 'hot_clients_ips', 0, 0)
end
end
@@ -103,23 +103,34 @@ async def rlim_middleware(
return response
else:
root_ctx: RootContext = app["_root.context"]
rate_limit = root_ctx.shared_config["anonymous_ratelimit"]
anonymous_ratelimiter = root_ctx.shared_config["anonymous_ratelimiter"]

ip_address = get_client_ip(request)

if not ip_address or rate_limit is None:
if not ip_address or anonymous_ratelimiter is None:
# No checks for rate limiting.
response = await handler(request)
# Arbitrary number for indicating no rate limiting.
response.headers["X-RateLimit-Limit"] = "1000"
response.headers["X-RateLimit-Remaining"] = "1000"
else:
rate_limit, hot_clients_ips_maxsize, hot_clients_ips_threshold_ratio = (
anonymous_ratelimiter["rlimit"],
anonymous_ratelimiter["hot_clients_ips_maxsize"],
anonymous_ratelimiter["hot_clients_ips_threshold_ratio"],
)
ret = await redis_helper.execute_script(
rr,
"ratelimit",
_rlim_script,
["ip", ip_address],
[str(now), str(_rlim_window), str(rate_limit)],
[
str(now),
str(_rlim_window),
str(rate_limit),
str(hot_clients_ips_maxsize),
str(hot_clients_ips_threshold_ratio),
],
)
if ret is None:
remaining = rate_limit
7 changes: 6 additions & 1 deletion src/ai/backend/manager/config.py
Original file line number Diff line number Diff line change
@@ -456,7 +456,12 @@ def container_registry_serialize(v: dict[str, Any]) -> dict[str, str]:
},
).allow_extra("*"),
t.Key("roundrobin_states", default=None): t.Null | tx.RoundRobinStatesJSONString,
t.Key("anonymous_ratelimit", default=None): t.Null | t.ToInt,
t.Key("anonymous_ratelimiter", default=None): t.Null
| t.Dict({
t.Key("rlimit"): t.ToInt(),
t.Key("hot_clients_ips_maxsize", default=1000): t.Null | t.ToInt(),
t.Key("hot_clients_ips_threshold_ratio", default=0.8): t.Null | t.ToFloat(),
}),
}).allow_extra("*")

_volume_defaults: dict[str, Any] = {

0 comments on commit f8caa97

Please sign in to comment.