Bound remote pool and rate-limiter memory
Both tables grew one entry per distinct key and never shrank, so a long-running SDK client or a public serve process accumulated state for every endpoint or identity it had ever seen. GC cannot reclaim them while the pool and limiter still reference them. Cap the pool at 64 endpoint buckets and the limiter at 4096 buckets. Both evict useless state first: connections past the idle timeout the server has likely dropped anyway, and buckets that have fully refilled, which carry no throttling information. Only then fall back to evicting the oldest entry. Evicting a limiter bucket resets throttling for that identity, which is the deliberate trade: an attacker cycling identities faster than they go idle can regain tokens, but unbounded growth would take the process down instead.
This commit is contained in:
@@ -26,6 +26,7 @@ from browser_cli.framing import frame
|
|||||||
# hand back one the server has just timed out and closed.
|
# hand back one the server has just timed out and closed.
|
||||||
_MAX_IDLE_SECONDS = max(5, REMOTE_SESSION_IDLE_TIMEOUT - 5)
|
_MAX_IDLE_SECONDS = max(5, REMOTE_SESSION_IDLE_TIMEOUT - 5)
|
||||||
_MAX_PER_ENDPOINT = 8
|
_MAX_PER_ENDPOINT = 8
|
||||||
|
_MAX_ENDPOINTS = 64
|
||||||
|
|
||||||
class PooledConnection:
|
class PooledConnection:
|
||||||
__slots__ = ("sock", "secret", "last_used")
|
__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
|
_close(conn.sock) # too old — assume the server has dropped it
|
||||||
return None
|
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:
|
def checkin(endpoint: str, conn: PooledConnection) -> None:
|
||||||
"""Return a still-healthy connection to the pool for reuse."""
|
"""Return a still-healthy connection to the pool for reuse."""
|
||||||
conn.last_used = time.monotonic()
|
conn.last_used = time.monotonic()
|
||||||
with _LOCK:
|
with _LOCK:
|
||||||
|
if endpoint not in _POOL and len(_POOL) >= _MAX_ENDPOINTS:
|
||||||
|
_prune_endpoints_locked(conn.last_used)
|
||||||
bucket = _POOL.setdefault(endpoint, [])
|
bucket = _POOL.setdefault(endpoint, [])
|
||||||
if len(bucket) >= _MAX_PER_ENDPOINT:
|
if len(bucket) >= _MAX_PER_ENDPOINT:
|
||||||
_close(conn.sock)
|
_close(conn.sock)
|
||||||
|
|||||||
@@ -70,19 +70,48 @@ class RateLimiter:
|
|||||||
``rate`` is the sustained refill in tokens/second; ``burst`` is the bucket
|
``rate`` is the sustained refill in tokens/second; ``burst`` is the bucket
|
||||||
capacity (defaults to ``rate``). ``rate <= 0`` disables limiting entirely.
|
capacity (defaults to ``rate``). ``rate <= 0`` disables limiting entirely.
|
||||||
Thread-safe so it can be shared across all connections of one serve process.
|
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.rate = float(rate)
|
||||||
self.capacity = float(burst) if burst is not None else max(float(rate), 1.0)
|
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._buckets: dict[str, tuple[float, float]] = {}
|
||||||
self._lock = threading.Lock()
|
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:
|
def allow(self, key: str) -> bool:
|
||||||
if self.rate <= 0:
|
if self.rate <= 0:
|
||||||
return True
|
return True
|
||||||
now = time.monotonic()
|
now = time.monotonic()
|
||||||
with self._lock:
|
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, last = self._buckets.get(key, (self.capacity, now))
|
||||||
tokens = min(self.capacity, tokens + (now - last) * self.rate)
|
tokens = min(self.capacity, tokens + (now - last) * self.rate)
|
||||||
if tokens < 1.0:
|
if tokens < 1.0:
|
||||||
|
|||||||
@@ -43,6 +43,39 @@ def test_checkin_caps_pool_size():
|
|||||||
b.close()
|
b.close()
|
||||||
pool.close_all()
|
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():
|
def test_session_inner_message_strips_auth_fields():
|
||||||
msg = {
|
msg = {
|
||||||
"id": "1", "command": "tabs.list", "args": {}, "user_agent": "browser-cli/1",
|
"id": "1", "command": "tabs.list", "args": {}, "user_agent": "browser-cli/1",
|
||||||
|
|||||||
@@ -181,6 +181,23 @@ def test_rate_limiter_is_per_key():
|
|||||||
assert limiter.allow("a") is False
|
assert limiter.allow("a") is False
|
||||||
assert limiter.allow("b") 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 ────────────────────────────────────────────────────────────────
|
# ── ServeSecurity ────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
def test_effective_policy_prefers_per_key_override():
|
def test_effective_policy_prefers_per_key_override():
|
||||||
|
|||||||
Reference in New Issue
Block a user