fix(security): bound agent shell token store without deadlock

This commit is contained in:
Shadowbroker
2026-08-17 01:05:12 -06:00
parent 4f009150c2
commit 7afa085d44
+18 -11
View File
@@ -13,25 +13,31 @@ _store: dict[str, float] = {}
_lock = Lock()
def _purge_expired(*, force: bool = False) -> None:
def _purge_expired_locked(now: float) -> None:
"""Drop expired tokens while the caller holds ``_lock``."""
expired = [token for token, expires in _store.items() if expires <= now]
for token in expired:
_store.pop(token, None)
def _purge_expired() -> None:
now = time.time()
with _lock:
expired = [token for token, expires in _store.items() if expires <= now]
for token in expired:
_store.pop(token, None)
if force and len(_store) > _MAX_ACTIVE_TOKENS:
for token in list(_store.keys())[: len(_store) - _MAX_ACTIVE_TOKENS]:
_store.pop(token, None)
_purge_expired_locked(now)
def mint_agent_shell_ws_token() -> tuple[str, int]:
"""Return (token, expires_in_seconds)."""
_purge_expired()
"""Return (token, expires_in_seconds) without exceeding the store bound."""
token = secrets.token_urlsafe(32)
expires_at = time.time() + _TOKEN_TTL_SECONDS
now = time.time()
expires_at = now + _TOKEN_TTL_SECONDS
with _lock:
_purge_expired_locked(now)
if len(_store) >= _MAX_ACTIVE_TOKENS:
_purge_expired(force=True)
# Evict the token that will expire soonest. This keeps the store
# bounded without recursively acquiring the non-reentrant lock.
oldest = min(_store, key=_store.get)
_store.pop(oldest, None)
_store[token] = expires_at
return token, int(_TOKEN_TTL_SECONDS)
@@ -43,6 +49,7 @@ def consume_agent_shell_ws_token(token: str) -> bool:
return False
now = time.time()
with _lock:
_purge_expired_locked(now)
expires_at = _store.pop(cleaned, None)
return expires_at is not None and expires_at > now