-
Notifications
You must be signed in to change notification settings - Fork 30
Broker: thread-safe rate limiters + stale-key eviction #118
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -53,42 +53,72 @@ class BrokerConfig: | |
|
|
||
|
|
||
| class RateLimiter: | ||
| _EVICTION_INTERVAL_SECONDS = 60.0 | ||
|
|
||
| def __init__(self, max_events: int, window_seconds: int = 60) -> None: | ||
| self._max_events = max_events | ||
| self._window_seconds = window_seconds | ||
| self._events: dict[str, Deque[float]] = defaultdict(deque) | ||
| self._lock = threading.Lock() | ||
| self._last_eviction = time.monotonic() | ||
|
|
||
| def allow(self, key: str) -> bool: | ||
| now = time.monotonic() | ||
| events = self._events[key] | ||
| cutoff = now - self._window_seconds | ||
| while events and events[0] < cutoff: | ||
| events.popleft() | ||
| if len(events) >= self._max_events: | ||
| return False | ||
| events.append(now) | ||
| return True | ||
| with self._lock: | ||
| self._maybe_evict_stale_keys(now, cutoff) | ||
| events = self._events[key] | ||
| while events and events[0] < cutoff: | ||
| events.popleft() | ||
| if len(events) >= self._max_events: | ||
| return False | ||
| events.append(now) | ||
| return True | ||
|
|
||
| def _maybe_evict_stale_keys(self, now: float, cutoff: float) -> None: | ||
| """Drop keys whose newest event fell out of the window. Caller holds lock.""" | ||
| if now - self._last_eviction < self._EVICTION_INTERVAL_SECONDS: | ||
| return | ||
| self._last_eviction = now | ||
| stale = [key for key, events in self._events.items() if not events or events[-1] < cutoff] | ||
| for key in stale: | ||
| del self._events[key] | ||
|
|
||
|
|
||
| class ByteRateLimiter: | ||
| _EVICTION_INTERVAL_SECONDS = 60.0 | ||
|
|
||
| def __init__(self, max_bytes: int, window_seconds: int = 60) -> None: | ||
| self._max_bytes = max_bytes | ||
| self._window_seconds = window_seconds | ||
| self._events: dict[str, Deque[tuple[float, int]]] = defaultdict(deque) | ||
| self._lock = threading.Lock() | ||
| self._last_eviction = time.monotonic() | ||
|
|
||
| def allow(self, key: str, size: int) -> bool: | ||
| if self._max_bytes <= 0: | ||
| return True | ||
| now = time.monotonic() | ||
| events = self._events[key] | ||
| cutoff = now - self._window_seconds | ||
| while events and events[0][0] < cutoff: | ||
| events.popleft() | ||
| total = sum(event_size for _, event_size in events) | ||
| if total + size > self._max_bytes: | ||
| return False | ||
| events.append((now, size)) | ||
| return True | ||
| with self._lock: | ||
| self._maybe_evict_stale_keys(now, cutoff) | ||
| events = self._events[key] | ||
| while events and events[0][0] < cutoff: | ||
| events.popleft() | ||
| total = sum(event_size for _, event_size in events) | ||
| if total + size > self._max_bytes: | ||
| return False | ||
| events.append((now, size)) | ||
| return True | ||
|
|
||
| def _maybe_evict_stale_keys(self, now: float, cutoff: float) -> None: | ||
| """Drop keys whose newest event fell out of the window. Caller holds lock.""" | ||
| if now - self._last_eviction < self._EVICTION_INTERVAL_SECONDS: | ||
| return | ||
| self._last_eviction = now | ||
| stale = [key for key, events in self._events.items() if not events or events[-1][0] < cutoff] | ||
| for key in stale: | ||
| del self._events[key] | ||
|
|
||
|
|
||
| class BrokerHandler(http.server.BaseHTTPRequestHandler): | ||
|
|
@@ -470,7 +500,12 @@ def __init__(self, server_address: tuple[str, int], config: BrokerConfig) -> Non | |
| def validate_token(self, token: str | None) -> str | None: | ||
| if token is None: | ||
| return None | ||
| if token in self.config.session_tokens: | ||
| token_matches = False | ||
| for session_token in self.config.session_tokens: | ||
| # Check every configured token so the comparison time does not | ||
| # depend on where (or whether) the match occurs. | ||
| token_matches |= hmac.compare_digest(token, session_token) | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
When any static session token is configured, a malformed Useful? React with 👍 / 👎. |
||
| if token_matches: | ||
| digest = hashlib.sha256(token.encode("utf-8")).hexdigest()[:16] | ||
| return f"static:{digest}" | ||
| if self.config.token_secret is None: | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Because
nowandcutoffare captured before entering this newly added lock, a request that waits here while another thread is evicting/checking a large limiter map can make its decision using an already-stale window. If entries expire during that wait, they remain counted for this request and can incorrectly return 429 or byte-limit a valid request; the same pre-lock timestamp pattern exists inByteRateLimiter.allow().Useful? React with 👍 / 👎.