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.
86 lines
2.6 KiB
Python
86 lines
2.6 KiB
Python
"""Unit tests for the remote connection pool (browser_cli.remote.pool)."""
|
|
import socket
|
|
|
|
from browser_cli.remote import pool
|
|
|
|
def _socketpair():
|
|
a, b = socket.socketpair()
|
|
return a, b
|
|
|
|
def test_checkin_then_checkout_returns_same_connection():
|
|
pool.close_all()
|
|
a, b = _socketpair()
|
|
conn = pool.PooledConnection(a, b"secret")
|
|
pool.checkin("host:1", conn)
|
|
assert pool.checkout("host:1") is conn
|
|
assert pool.checkout("host:1") is None # only one was pooled
|
|
b.close()
|
|
pool.close_all()
|
|
|
|
def test_checkout_drops_stale_connection(monkeypatch):
|
|
pool.close_all()
|
|
a, b = _socketpair()
|
|
conn = pool.PooledConnection(a, b"secret")
|
|
pool.checkin("host:2", conn)
|
|
# Make the pooled connection look older than the idle bound.
|
|
conn.last_used -= (pool._MAX_IDLE_SECONDS + 1)
|
|
assert pool.checkout("host:2") is None # stale → dropped, not returned
|
|
b.close()
|
|
pool.close_all()
|
|
|
|
def test_checkin_caps_pool_size():
|
|
pool.close_all()
|
|
sockets = []
|
|
for i in range(pool._MAX_PER_ENDPOINT + 3):
|
|
a, b = _socketpair()
|
|
sockets.append(b)
|
|
pool.checkin("host:3", pool.PooledConnection(a, b"secret"))
|
|
drained = 0
|
|
while pool.checkout("host:3") is not None:
|
|
drained += 1
|
|
assert drained == pool._MAX_PER_ENDPOINT
|
|
for b in sockets:
|
|
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",
|
|
"pubkey": "x", "sig": "y", "pq_kex": {}, "encrypted": {}, "accept_encoding": {"x": 1},
|
|
}
|
|
inner = pool.session_inner_message(msg)
|
|
assert inner == {"id": "1", "command": "tabs.list", "args": {}, "user_agent": "browser-cli/1", "accept_encoding": {"x": 1}}
|