Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(websocket): harden security and robustness

Browse Source 
- Use hmac.compare_digest for timing-safe static token comparison
- Add issued token capacity limit (_MAX_ISSUED_TOKENS=10000) with 429 response
- Use atomic pop in _take_issued_token_if_valid to eliminate TOCTOU window
- Enforce TLSv1.2 minimum version for SSL connections
- Extract _safe_send helper for consistent ConnectionClosed handling
- Move connection registration after ready send to prevent out-of-order delivery
- Add HTTP-level allow_from check and client_id truncation in process_request
- Make stop() idempotent with graceful shutdown error handling
- Normalize path via validator instead of leaving raw value
- Default websocket_requires_token to True for secure-by-default behavior
- Add integration tests and ws_test_client helper
- Refactor tests to use shared _ch factory and bus fixture
This commit is contained in:
chengyongru
2026-04-09 18:22:35 +08:00
committed by Xubin Ren
parent e0ccc401c0
commit 56a5906db5
4 changed files with 1102 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/channels/websocket.py
+77 -47
View File
@@ -65,7 +65,7 @@ class WebSocketConfig(Base):
token_issue_path: str = ""
token_issue_secret: str = ""
token_ttl_s: int = Field(default=300, ge=30, le=86_400)
websocket_requires_token: bool = False
websocket_requires_token: bool = True
allow_from: list[str] = Field(default_factory=lambda: ["*"])
streaming: bool = True
max_message_bytes: int = Field(default=1_048_576, ge=1024, le=16_777_216)
@@ -79,7 +79,7 @@ class WebSocketConfig(Base):
def path_must_start_with_slash(cls, value: str) -> str:
if not value.startswith("/"):
raise ValueError('path must start with "/"')
return value
return _normalize_config_path(value)
@field_validator("token_issue_path")
@classmethod
@@ -130,6 +130,12 @@ def _parse_query(path_with_query: str) -> dict[str, list[str]]:
return _parse_request_path(path_with_query)[1]
def _query_first(query: dict[str, list[str]], key: str) -> str | None:
"""Return the first value for *key*, or None."""
values = query.get(key)
return values[0] if values else None
def _parse_inbound_payload(raw: str) -> str | None:
"""Parse a client frame into text; return None for empty or unrecognized content."""
text = raw.strip()
@@ -197,9 +203,12 @@ class WebSocketChannel(BaseChannel):
"websocket: ssl_certfile and ssl_keyfile must both be set for WSS, or both left empty"
)
ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
ctx.minimum_version = ssl.TLSVersion.TLSv1_2
ctx.load_cert_chain(certfile=cert, keyfile=key)
return ctx
_MAX_ISSUED_TOKENS = 10_000
def _purge_expired_issued_tokens(self) -> None:
now = time.monotonic()
for token_key, expiry in list(self._issued_tokens.items()):
@@ -207,17 +216,19 @@ class WebSocketChannel(BaseChannel):
self._issued_tokens.pop(token_key, None)
def _take_issued_token_if_valid(self, token_value: str | None) -> bool:
"""Validate and consume one issued token (single use per connection attempt)."""
"""Validate and consume one issued token (single use per connection attempt).
Uses single-step pop to minimize the window between lookup and removal;
safe under asyncio's single-threaded cooperative model.
"""
if not token_value:
return False
self._purge_expired_issued_tokens()
expiry = self._issued_tokens.get(token_value)
expiry = self._issued_tokens.pop(token_value, None)
if expiry is None:
return False
if time.monotonic() > expiry:
self._issued_tokens.pop(token_value, None)
return False
self._issued_tokens.pop(token_value, None)
return True
def _handle_token_issue_http(self, connection: Any, request: Any) -> Any:
@@ -231,6 +242,12 @@ class WebSocketChannel(BaseChannel):
"any client can obtain connection tokens — set token_issue_secret for production."
)
self._purge_expired_issued_tokens()
if len(self._issued_tokens) >= self._MAX_ISSUED_TOKENS:
logger.error(
"websocket: too many outstanding issued tokens ({}), rejecting issuance",
len(self._issued_tokens),
)
return _http_json_response({"error": "too many outstanding tokens"}, status=429)
token_value = f"nbwt_{secrets.token_urlsafe(32)}"
self._issued_tokens[token_value] = time.monotonic() + float(self.config.token_ttl_s)
@@ -238,13 +255,12 @@ class WebSocketChannel(BaseChannel):
{"token": token_value, "expires_in": self.config.token_ttl_s}
)
def _authorize_websocket_handshake(self, connection: Any, request_path: str) -> Any:
query = _parse_query(request_path)
supplied = (query.get("token") or [None])[0]
def _authorize_websocket_handshake(self, connection: Any, query: dict[str, list[str]]) -> Any:
supplied = _query_first(query, "token")
static_token = self.config.token.strip()
if static_token:
if supplied == static_token:
if supplied and hmac.compare_digest(supplied, static_token):
return None
if supplied and self._take_issued_token_if_valid(supplied):
return None
@@ -279,7 +295,15 @@ class WebSocketChannel(BaseChannel):
expected_ws = self._expected_path()
if got != expected_ws:
return connection.respond(404, "Not Found")
return self._authorize_websocket_handshake(connection, request.path)
# Early reject before WebSocket upgrade to avoid unnecessary overhead;
# _handle_message() performs a second check as defense-in-depth.
query = _parse_query(request.path)
client_id = _query_first(query, "client_id") or ""
if len(client_id) > 128:
client_id = client_id[:128]
if not self.is_allowed(client_id):
return connection.respond(403, "Forbidden")
return self._authorize_websocket_handshake(connection, query)
async def handler(connection: ServerConnection) -> None:
await self._connection_loop(connection)
@@ -321,26 +345,30 @@ class WebSocketChannel(BaseChannel):
request = connection.request
path_part = request.path if request else "/"
_, query = _parse_request_path(path_part)
client_id_raw = (query.get("client_id") or [None])[0]
client_id_raw = _query_first(query, "client_id")
client_id = client_id_raw.strip() if client_id_raw else ""
if not client_id:
client_id = f"anon-{uuid.uuid4().hex[:12]}"
elif len(client_id) > 128:
logger.warning("websocket: client_id too long ({} chars), truncating", len(client_id))
client_id = client_id[:128]
chat_id = str(uuid.uuid4())
self._connections[chat_id] = connection
await connection.send(
json.dumps(
{
"event": "ready",
"chat_id": chat_id,
"client_id": client_id,
},
ensure_ascii=False,
)
)
try:
await connection.send(
json.dumps(
{
"event": "ready",
"chat_id": chat_id,
"client_id": client_id,
},
ensure_ascii=False,
)
)
# Register only after ready is successfully sent to avoid out-of-order sends
self._connections[chat_id] = connection
async for raw in connection:
if isinstance(raw, bytes):
try:
@@ -363,15 +391,34 @@ class WebSocketChannel(BaseChannel):
self._connections.pop(chat_id, None)
async def stop(self) -> None:
if not self._running:
return
self._running = False
if self._stop_event:
self._stop_event.set()
if self._server_task:
await self._server_task
try:
await self._server_task
except Exception as e:
logger.warning("websocket: server task error during shutdown: {}", e)
self._server_task = None
self._connections.clear()
self._issued_tokens.clear()
async def _safe_send(self, chat_id: str, raw: str, *, label: str = "") -> None:
"""Send a raw frame, cleaning up dead connections on ConnectionClosed."""
connection = self._connections.get(chat_id)
if connection is None:
return
try:
await connection.send(raw)
except ConnectionClosed:
self._connections.pop(chat_id, None)
logger.warning("websocket{}connection gone for chat_id={}", label, chat_id)
except Exception as e:
logger.error("websocket{}send failed: {}", label, e)
raise
async def send(self, msg: OutboundMessage) -> None:
connection = self._connections.get(msg.chat_id)
if connection is None:
@@ -386,14 +433,7 @@ class WebSocketChannel(BaseChannel):
if msg.reply_to:
payload["reply_to"] = msg.reply_to
raw = json.dumps(payload, ensure_ascii=False)
try:
await connection.send(raw)
except ConnectionClosed:
self._connections.pop(msg.chat_id, None)
logger.warning("websocket: connection gone for chat_id={}", msg.chat_id)
except Exception as e:
logger.error("websocket send failed: {}", e)
raise
await self._safe_send(msg.chat_id, raw, label=" ")
async def send_delta(
self,
@@ -401,27 +441,17 @@ class WebSocketChannel(BaseChannel):
delta: str,
metadata: dict[str, Any] | None = None,
) -> None:
connection = self._connections.get(chat_id)
if connection is None:
if self._connections.get(chat_id) is None:
return
meta = metadata or {}
if meta.get("_stream_end"):
body: dict[str, Any] = {"event": "stream_end"}
if meta.get("_stream_id") is not None:
body["stream_id"] = meta["_stream_id"]
else:
body = {
"event": "delta",
"text": delta,
}
if meta.get("_stream_id") is not None:
body["stream_id"] = meta["_stream_id"]
if meta.get("_stream_id") is not None:
body["stream_id"] = meta["_stream_id"]
raw = json.dumps(body, ensure_ascii=False)
try:
await connection.send(raw)
except ConnectionClosed:
self._connections.pop(chat_id, None)
logger.warning("websocket: stream connection gone for chat_id={}", chat_id)
except Exception as e:
logger.error("websocket stream send failed: {}", e)
raise
await self._safe_send(chat_id, raw, label=" stream ")