def test_check_global_disabled():
"""Ensure global check is disabled if action is by_user or by_ip."""
action = RateLimitedAction("test",
timedelta(hours=1),
5,
by_user=True,
by_ip=False)
with raises(RateLimitError):
action.check_global()
action = RateLimitedAction("test",
timedelta(hours=1),
5,
by_user=False,
by_ip=True)
with raises(RateLimitError):
action.check_global()
def test_simple_global_rate_limiting(redis):
"""Ensure simple global rate-limiting is working."""
limit = 5
# define an action with max_burst equal to the full limit
action = RateLimitedAction(
"testaction",
timedelta(hours=1),
limit,
max_burst=limit,
by_user=False,
by_ip=False,
redis=redis,
)
# run the action the full number of times, should all be allowed
for _ in range(limit):
result = action.check_global()
assert result.is_allowed
# try one more time, should be rejected
result = action.check_global()
assert not result.is_allowed
def check_rate_limit(request: Request, action_name: str) -> RateLimitResult:
"""Check the rate limit for a particular action on a request."""
action = None
# check for a custom rate-limit for the user
if request.user:
user_limit = (
request.query(UserRateLimit)
.filter(
UserRateLimit.user == request.user, UserRateLimit.action == action_name
)
.one_or_none()
)
if user_limit:
action = RateLimitedAction(
action_name,
user_limit.period,
user_limit.limit,
by_user=True,
by_ip=False,
)
# if a custom rate-limit wasn't found, use the default, global rate-limit
if not action:
try:
action = RATE_LIMITED_ACTIONS[action_name]
except KeyError as exc:
raise ValueError("Invalid action name: %s" % action_name) from exc
action.redis = request.redis
results = []
if action.is_global:
results.append(action.check_global())
if action.by_user and request.user:
results.append(action.check_for_user_id(request.user.user_id))
if action.by_ip and request.remote_addr:
results.append(action.check_for_ip(request.remote_addr))
# no checks were done, return the "not limited" result
if not results:
return RateLimitResult.unlimited_result()
return RateLimitResult.merged_result(results)