Coverage for polar/rate_limit.py: 63%

1from collections.abc import Sequence 

2 

3from ratelimit import RateLimitMiddleware, Rule 

4from ratelimit.auths import EmptyInformation 

5from ratelimit.auths.ip import client_ip 

6from ratelimit.backends.redis import RedisBackend 

7from ratelimit.types import ASGIApp, Scope 

8 

9from polar.auth.models import AuthSubject, Subject, is_anonymous 

10from polar.config import Environment, settings 

11from polar.enums import RateLimitGroup 

12from polar.redis import create_redis 

13 

14 

15async def _authenticate(scope: Scope) -> tuple[str, RateLimitGroup]: 

16 auth_subject: AuthSubject[Subject] = scope["state"]["auth_subject"] 

17 

18 if is_anonymous(auth_subject): 

19 try: 

20 ip, _ = await client_ip(scope) 

21 return ip, RateLimitGroup.default 

22 except EmptyInformation: 

23 return auth_subject.rate_limit_key 

24 

25 return auth_subject.rate_limit_key 

26 

27 

28_BASE_RULES: dict[str, Sequence[Rule]] = { 

29 "^/v1/login-code": [Rule(minute=6, hour=12, block_time=900, zone="login-code")], 

30 "^/v1/customer-portal/customer-session/(request|authenticate)": [ 

31 Rule(minute=6, hour=12, block_time=900, zone="customer-session-login") 

32 ], 

33 "^/v1/customer-portal/license-keys/(validate|activate|deactivate)": [ 

34 Rule(second=3, block_time=60, zone="customer-license-key") 

35 ], 

36 "^/v1/customer-seats/claim/.+/stream": [ 

37 Rule(minute=10, block_time=300, zone="seat-claim-stream") 

38 ], 

39} 

40 

41_SANDBOX_RULES: dict[str, Sequence[Rule]] = { 

42 **_BASE_RULES, 

43 "^/v1": [ 

44 Rule(group=RateLimitGroup.default, minute=100, zone="api"), 

45 Rule(group=RateLimitGroup.web, second=50, zone="api"), 

46 Rule(group=RateLimitGroup.elevated, second=50, zone="api"), 

47 ], 

48} 

49 

50_PRODUCTION_RULES: dict[str, Sequence[Rule]] = { 

51 **_BASE_RULES, 

52 "^/v1": [ 

53 Rule(group=RateLimitGroup.default, minute=500, zone="api"), 

54 Rule(group=RateLimitGroup.web, second=100, zone="api"), 

55 Rule(group=RateLimitGroup.elevated, second=100, zone="api"), 

56 ], 

57} 

58 

59 

60def get_middleware(app: ASGIApp) -> RateLimitMiddleware: 

61 match settings.ENV: 

62 case Environment.production: 62 ↛ 63line 62 didn't jump to line 63 because the pattern on line 62 never matched

63 rules = _PRODUCTION_RULES 

64 case Environment.sandbox: 64 ↛ 65line 64 didn't jump to line 65 because the pattern on line 64 never matched

65 rules = _SANDBOX_RULES 

66 case _: 

67 rules = {} 

68 return RateLimitMiddleware( 

69 app, _authenticate, RedisBackend(create_redis("rate-limit")), rules 

70 ) 

71 

72 

73__all__ = ["get_middleware"]