Coverage for /home/runner/work/viur-core/viur-core/viur/src/viur/core/contrib/ratelimit.py: 0%

1""" 

2Request-level rate limiter using App Engine Memcache. 

3 

4Registers as a :class:`~viur.core.request.RequestValidator` and therefore 

5runs *before* any session, routing, or handler logic — making it the 

6earliest possible place to shed excess traffic. 

7 

8Guests are identified by their IP address (IPv6 addresses are bucketed into 

9/64 prefixes so that a single host cannot trivially rotate around the limit). 

10Authenticated users are identified by their Datastore user key. 

11 

12Usage:: 

13 

14 from viur.core.request import Router 

15 from viur.core.contrib.ratelimit import RequestRateLimit, TimeWindow 

16 

17 Router.requestValidators.append( 

18 RequestRateLimit( 

19 rate_for_guests=TimeWindow(limit=200, time_window=60), 

20 rate_for_users=TimeWindow(limit=500, time_window=60), 

21 ) 

22 ) 

23""" 

24import dataclasses 

25import ipaddress 

26import logging 

27import time 

28import typing as t 

29 

30from google.appengine.api.memcache import Client 

31 

32from viur.core import current 

33from viur.core.request import RequestValidator, Router 

34 

35logger = logging.getLogger(__name__) 

36 

37_memcache = Client() 

38 

39 

40@dataclasses.dataclass(frozen=True) 

41class TimeWindow: 

42 """Rate-limit budget for a single time window. 

43 

44 :param limit: Maximum number of requests allowed within *time_window*. 

45 :param time_window: Length of the window in seconds. 

46 """ 

47 limit: int 

48 time_window: int 

49 

50 

51class RequestRateLimit(RequestValidator): 

52 """Global HTTP request rate limiter. 

53 

54 Enforces separate budgets for anonymous (guest) and authenticated 

55 requests. When the budget is exceeded the validator returns HTTP 429 

56 and sets the ``Retry-After`` header so clients know when to retry. 

57 

58 :param rate_for_guests: Budget applied to unauthenticated requests. 

59 :param rate_for_users: Budget applied to authenticated requests. 

60 :param namespace: Memcache namespace used for all rate-limit keys. 

61 """ 

62 

63 name = "RequestRateLimit" 

64 

65 def __init__( 

66 self, 

67 rate_for_guests: TimeWindow = TimeWindow(limit=1000, time_window=60), 

68 rate_for_users: TimeWindow = TimeWindow(limit=2000, time_window=60), 

69 namespace: str = "viur_rate_limit", 

70 ): 

71 self.rate_for_guests = rate_for_guests 

72 self.rate_for_users = rate_for_users 

73 self.namespace = namespace 

74 

75 def validate(self, request: "Router") -> t.Optional[tuple[int, str, str]]: 

76 if request.is_deferred: 

77 return None # Task Queue requests are always allowed 

78 

79 if user := current.user.get(): 

80 client_id = str(user["key"]) 

81 rate = self.rate_for_users 

82 else: 

83 client_id = self._get_request_ip() 

84 rate = self.rate_for_guests 

85 

86 current_time = time.time() / rate.time_window 

87 key = f"rate_limit:{client_id}:{int(current_time)}" 

88 

89 count = _memcache.get(key, namespace=self.namespace) 

90 logger.debug(f"rate limit check: {client_id=} {count=} limit={rate.limit}") 

91 

92 if count is None: 

93 _memcache.add(key, 1, time=rate.time_window, namespace=self.namespace) 

94 return None 

95 

96 if count < rate.limit: 

97 _memcache.incr(key, initial_value=1, namespace=self.namespace) 

98 return None 

99 

100 # Budget exhausted — tell the client when the current window expires. 

101 seconds_into_window = (current_time - int(current_time)) * rate.time_window 

102 retry_after = int(rate.time_window - seconds_into_window) 

103 request.response.headers["Retry-After"] = str(retry_after) 

104 return 429, "Too Many Requests", "Too Many Requests. Please try again later." 

105 

106 @staticmethod 

107 def _get_request_ip() -> str: 

108 """Return a stable client identifier derived from the remote address. 

109 

110 IPv4 addresses are returned as-is. For IPv6 the /64 network prefix 

111 is returned so that a single host cannot trivially rotate its 

112 interface identifier to bypass the limit. 

113 """ 

114 raw = current.request.get().request.remote_addr 

115 ip = ipaddress.ip_address(raw) 

116 

117 if isinstance(ip, ipaddress.IPv4Address): 

118 return str(ip) 

119 

120 if isinstance(ip, ipaddress.IPv6Address): 

121 return str(ipaddress.IPv6Network((ip, 64), strict=False)) 

122 

123 raise NotImplementedError(f"Unsupported IP version: {ip!r}")