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