def test_check_for_ip_invalid_address():
"""Ensure RateLimitedAction.check_for_ip can't take an invalid IP."""
ip = "123.456.789.123"
action = RateLimitedAction("testaction", timedelta(hours=1), 10)
with raises(ValueError):
action.check_for_ip(ip)
def test_simple_rate_limiting_by_ip(redis):
"""Ensure simple rate-limiting by IP address is working."""
limit = 5
ip = "123.123.123.123"
# define an action with max_burst equal to the full limit
action = RateLimitedAction("testaction",
timedelta(hours=1),
limit,
max_burst=limit,
redis=redis)
# run the action the full number of times, should all be allowed
for _ in range(limit):
result = action.check_for_ip(ip)
assert result.is_allowed
# try one more time, should be rejected
result = action.check_for_ip(ip)
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)
def test_check_by_ip_disabled():
"""Ensure non-by_ip RateLimitedAction can't be checked by ip."""
action = RateLimitedAction("test", timedelta(hours=1), 5, by_ip=False)
with raises(RateLimitError):
action.check_for_ip("123.123.123.123")