diff --git a/browser_cli/remote/pool.py b/browser_cli/remote/pool.py index e00bbec..f60bc38 100644 --- a/browser_cli/remote/pool.py +++ b/browser_cli/remote/pool.py @@ -26,6 +26,7 @@ from browser_cli.framing import frame # hand back one the server has just timed out and closed. _MAX_IDLE_SECONDS = max(5, REMOTE_SESSION_IDLE_TIMEOUT - 5) _MAX_PER_ENDPOINT = 8 +_MAX_ENDPOINTS = 64 class PooledConnection: __slots__ = ("sock", "secret", "last_used") @@ -56,10 +57,32 @@ def checkout(endpoint: str) -> PooledConnection | None: _close(conn.sock) # too old — assume the server has dropped it return None +def _prune_endpoints_locked(now: float) -> None: + """Keep the number of endpoint buckets bounded for long-running SDK users.""" + for endpoint, bucket in list(_POOL.items()): + fresh = [conn for conn in bucket if now - conn.last_used <= _MAX_IDLE_SECONDS] + if fresh: + _POOL[endpoint] = fresh + else: + for conn in bucket: + _close(conn.sock) + _POOL.pop(endpoint, None) + + while len(_POOL) >= _MAX_ENDPOINTS: + oldest_endpoint, bucket = min( + _POOL.items(), + key=lambda item: min(conn.last_used for conn in item[1]) if item[1] else 0.0, + ) + for conn in bucket: + _close(conn.sock) + _POOL.pop(oldest_endpoint, None) + def checkin(endpoint: str, conn: PooledConnection) -> None: """Return a still-healthy connection to the pool for reuse.""" conn.last_used = time.monotonic() with _LOCK: + if endpoint not in _POOL and len(_POOL) >= _MAX_ENDPOINTS: + _prune_endpoints_locked(conn.last_used) bucket = _POOL.setdefault(endpoint, []) if len(bucket) >= _MAX_PER_ENDPOINT: _close(conn.sock) diff --git a/browser_cli/serve/security.py b/browser_cli/serve/security.py index 65c181d..5ab67ad 100644 --- a/browser_cli/serve/security.py +++ b/browser_cli/serve/security.py @@ -70,19 +70,48 @@ class RateLimiter: ``rate`` is the sustained refill in tokens/second; ``burst`` is the bucket capacity (defaults to ``rate``). ``rate <= 0`` disables limiting entirely. Thread-safe so it can be shared across all connections of one serve process. + + The bucket table is capped. Without that bound, a long-running public server + could retain one entry per ever-seen identity/IP forever; GC cannot reclaim + those entries because the limiter still references them. """ - def __init__(self, rate: float, burst: float | None = None) -> None: + def __init__(self, rate: float, burst: float | None = None, max_buckets: int = 4096) -> None: self.rate = float(rate) self.capacity = float(burst) if burst is not None else max(float(rate), 1.0) + self.max_buckets = max(1, int(max_buckets)) self._buckets: dict[str, tuple[float, float]] = {} self._lock = threading.Lock() + def _prune_locked(self, now: float) -> None: + """Drop idle/full buckets, then oldest buckets, until the table is bounded.""" + if len(self._buckets) < self.max_buckets or self.rate <= 0: + return + + # Once a bucket has fully refilled, keeping it around carries no useful + # throttling state. Use at least 60s so normal active identities are not + # churned out aggressively on high-rate configs. + idle_seconds = max(60.0, (self.capacity / self.rate) * 2) + full_epsilon = 1e-9 + for bucket_key, (tokens, last) in list(self._buckets.items()): + refilled = min(self.capacity, tokens + (now - last) * self.rate) + if refilled >= self.capacity - full_epsilon and now - last >= idle_seconds: + self._buckets.pop(bucket_key, None) + + # If an attacker keeps creating fresh identities faster than they go idle, + # still keep memory bounded. Evict the oldest identity state; that may reset + # throttling for that identity, but bounded memory is more important here. + while len(self._buckets) >= self.max_buckets: + oldest_key = min(self._buckets, key=lambda k: self._buckets[k][1]) + self._buckets.pop(oldest_key, None) + def allow(self, key: str) -> bool: if self.rate <= 0: return True now = time.monotonic() with self._lock: + if key not in self._buckets and len(self._buckets) >= self.max_buckets: + self._prune_locked(now) tokens, last = self._buckets.get(key, (self.capacity, now)) tokens = min(self.capacity, tokens + (now - last) * self.rate) if tokens < 1.0: diff --git a/tests/test_remote_pool.py b/tests/test_remote_pool.py index 5c9e776..220e226 100644 --- a/tests/test_remote_pool.py +++ b/tests/test_remote_pool.py @@ -43,6 +43,39 @@ def test_checkin_caps_pool_size(): b.close() pool.close_all() +def test_checkin_caps_endpoint_buckets(): + pool.close_all() + peers = [] + try: + for i in range(pool._MAX_ENDPOINTS + 5): + a, b = _socketpair() + peers.append(b) + pool.checkin(f"host-{i}:443", pool.PooledConnection(a, b"secret")) + assert len(pool._POOL) <= pool._MAX_ENDPOINTS + finally: + for peer in peers: + peer.close() + pool.close_all() + +def test_checkin_prunes_stale_endpoint_buckets(): + pool.close_all() + old_a, old_b = _socketpair() + old = pool.PooledConnection(old_a, b"secret") + pool.checkin("old:443", old) + old.last_used -= pool._MAX_IDLE_SECONDS + 1 + peers = [old_b] + try: + for i in range(pool._MAX_ENDPOINTS): + a, b = _socketpair() + peers.append(b) + pool.checkin(f"new-{i}:443", pool.PooledConnection(a, b"secret")) + assert "old:443" not in pool._POOL + assert len(pool._POOL) <= pool._MAX_ENDPOINTS + finally: + for peer in peers: + peer.close() + pool.close_all() + def test_session_inner_message_strips_auth_fields(): msg = { "id": "1", "command": "tabs.list", "args": {}, "user_agent": "browser-cli/1", diff --git a/tests/test_serve_security.py b/tests/test_serve_security.py index 89384bd..b3d2330 100644 --- a/tests/test_serve_security.py +++ b/tests/test_serve_security.py @@ -181,6 +181,23 @@ def test_rate_limiter_is_per_key(): assert limiter.allow("a") is False assert limiter.allow("b") is False +def test_rate_limiter_caps_identity_buckets(): + limiter = RateLimiter(rate=0.0001, burst=1, max_buckets=3) + for i in range(10): + assert limiter.allow(f"key-{i}") is True + assert len(limiter._buckets) <= 3 + +def test_rate_limiter_prunes_refilled_idle_buckets(monkeypatch): + current = 1000.0 + monkeypatch.setattr("browser_cli.serve.security.time.monotonic", lambda: current) + limiter = RateLimiter(rate=1, burst=2, max_buckets=2) + assert limiter.allow("old") is True + current += 120.0 + assert limiter.allow("a") is True + assert limiter.allow("b") is True + assert "old" not in limiter._buckets + assert len(limiter._buckets) <= 2 + # ── ServeSecurity ──────────────────────────────────────────────────────────────── def test_effective_policy_prefers_per_key_override():