Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
67 changes: 51 additions & 16 deletions services/model-broker/openphone_model_broker.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Capture limiter time after acquiring the lock

Because now and cutoff are 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 in ByteRateLimiter.allow().

Useful? React with 👍 / 👎.

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):
Expand Down Expand Up @@ -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)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Reject non-ASCII bearer tokens without crashing

When any static session token is configured, a malformed Authorization: Bearer ... value containing non-ASCII bytes now reaches hmac.compare_digest(token, session_token), which raises TypeError for non-ASCII str inputs. That turns an unauthenticated request to /v1/responses or /v1/audio/transcriptions into an unhandled handler exception/connection drop with no audit event instead of the previous 401 path; validate or encode the bearer token before comparing so bad tokens are rejected normally.

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:
Expand Down
Loading