daniel156161
541b950519
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.
127 lines
4.5 KiB
Python
127 lines
4.5 KiB
Python
"""Per-process pool of authenticated remote connections for reuse.
|
|
|
|
A ``browser-cli serve`` connection stays open after its first (encrypted)
|
|
command, so the client can send further commands over it without re-running the
|
|
TCP/TLS/challenge/auth handshake (~hundreds of ms each). Only encrypted (PQ)
|
|
sessions are pooled — plaintext/legacy sessions stay one-shot, matching the
|
|
server, which only loops for encrypted sessions.
|
|
|
|
Connections are checked out exclusively (never shared between threads at once),
|
|
returned on success, and dropped on any I/O error or once older than an idle
|
|
bound (kept below the server's idle timeout so we don't reuse a connection the
|
|
server has already closed).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import atexit
|
|
import json
|
|
import socket
|
|
import threading
|
|
import time
|
|
|
|
from browser_cli.constants import REMOTE_SESSION_IDLE_TIMEOUT
|
|
from browser_cli.framing import frame
|
|
|
|
# Retire a pooled connection a few seconds before the server would, so we never
|
|
# 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")
|
|
|
|
def __init__(self, sock: socket.socket, secret: bytes) -> None:
|
|
self.sock = sock
|
|
self.secret = secret
|
|
self.last_used = time.monotonic()
|
|
|
|
_POOL: dict[str, list[PooledConnection]] = {}
|
|
_LOCK = threading.Lock()
|
|
|
|
def _close(sock: socket.socket) -> None:
|
|
try:
|
|
sock.close()
|
|
except OSError:
|
|
pass
|
|
|
|
def checkout(endpoint: str) -> PooledConnection | None:
|
|
"""Take an idle authenticated connection for *endpoint*, or None."""
|
|
now = time.monotonic()
|
|
with _LOCK:
|
|
conns = _POOL.get(endpoint)
|
|
while conns:
|
|
conn = conns.pop()
|
|
if now - conn.last_used <= _MAX_IDLE_SECONDS:
|
|
return conn
|
|
_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)
|
|
return
|
|
bucket.append(conn)
|
|
|
|
def discard(conn: PooledConnection) -> None:
|
|
"""Drop a connection that errored or is no longer usable."""
|
|
_close(conn.sock)
|
|
|
|
def close_all() -> None:
|
|
"""Close every pooled connection (process exit / test isolation)."""
|
|
with _LOCK:
|
|
for bucket in _POOL.values():
|
|
for conn in bucket:
|
|
_close(conn.sock)
|
|
_POOL.clear()
|
|
|
|
def session_inner_message(msg: dict) -> dict:
|
|
"""Strip auth/transport fields, leaving the command for an established session."""
|
|
keep = {"id", "command", "args", "user_agent", "accept_encoding", "_route", "_suppress_pq_warning"}
|
|
return {k: v for k, v in msg.items() if k in keep}
|
|
|
|
def send_over(conn: PooledConnection, msg: dict) -> bytes | None:
|
|
"""Send one command over an existing encrypted session. Raises on I/O error."""
|
|
from browser_cli.auth import pq_encrypt
|
|
from browser_cli.remote.socket import recv_all
|
|
from browser_cli.remote.transport import _decode_pq_response
|
|
|
|
inner = json.dumps(session_inner_message(msg)).encode("utf-8")
|
|
envelope = json.dumps({"encrypted": pq_encrypt(conn.secret, "request", inner)}).encode("utf-8")
|
|
conn.sock.sendall(frame(envelope))
|
|
response = recv_all(conn.sock)
|
|
if not response:
|
|
# EOF — an older server (no session loop) closed after one command. Treat as
|
|
# a transport failure so the caller re-handshakes; never as an app error,
|
|
# which could double-execute a non-idempotent command on retry.
|
|
raise EOFError("remote closed the pooled connection")
|
|
return _decode_pq_response(response, conn.secret)
|
|
|
|
atexit.register(close_all)
|