From f9289af67f6a7603e8e5519386fdfe7a27c5db77 Mon Sep 17 00:00:00 2001 From: Blasius Patrick Date: Tue, 4 Aug 2026 13:20:57 +0700 Subject: [PATCH] fix: add resolve_token to TokenStore for E2E handshake The E2E handshake needs the raw pairing token to derive the HKDF handshake key, but TokenStore only stored SHA-256 hashes. This adds: - token field to _StoredRecord (Fernet-encrypted at rest) - resolve_token(name) method to retrieve raw token - Graceful fallback to legacy auth when token not available (nodes paired before this fix get plaintext auth until re-paired) Also fixes auth logic: E2E proof verification now correctly skips the legacy token_store.validate() call instead of overwriting is_valid=True with the (empty) auth.token. Signed-off-by: Blasius Patrick --- tokens.py | 26 +++++++++++++++++++- wsserver/server.py | 61 ++++++++++++++++++++++++++-------------------- 2 files changed, 60 insertions(+), 27 deletions(-) diff --git a/tokens.py b/tokens.py index 293077f..541c8ef 100644 --- a/tokens.py +++ b/tokens.py @@ -238,15 +238,19 @@ class _StoredRecord: created_at: str revoked: bool last_used_at: str | None = None + token: str | None = None # raw token (encrypted at rest via Fernet file-level) def to_dict(self) -> dict[str, Any]: - return { + d: dict[str, Any] = { "name": self.name, "token_hash": self.token_hash, "created_at": self.created_at, "revoked": self.revoked, "last_used_at": self.last_used_at, } + if self.token is not None: + d["token"] = self.token + return d @classmethod def from_dict(cls, data: Mapping[str, Any]) -> _StoredRecord: @@ -259,6 +263,7 @@ def from_dict(cls, data: Mapping[str, Any]) -> _StoredRecord: created_at=data["created_at"], revoked=bool(data.get("revoked", False)), last_used_at=data.get("last_used_at"), + token=data.get("token"), ) def to_public(self) -> TokenRecord: @@ -377,6 +382,7 @@ def _write(records: list[_StoredRecord]) -> list[_StoredRecord]: name=name, token_hash=token_hash, created_at=now, + token=token, revoked=False, ) ) @@ -522,6 +528,24 @@ def validate(self, name: str, presented_token: str) -> bool: pass return True + def resolve_token(self, name: str) -> str: + """Return the raw pairing token for ``name``, or raise TokenStoreError. + + Only needed for E2E key derivation — the raw token is mixed into + the HKDF to produce the handshake key. The token is stored encrypted + at rest (Fernet file-level). + """ + records = self._read() + for rec in records: + if rec.name == name and not rec.revoked: + if rec.token is None: + raise TokenStoreError( + f"node {name!r} has no stored token " + f"(paired before v0.1.2 — re-pair to enable E2E)" + ) + return rec.token + raise TokenStoreError(f"node {name!r} not found or revoked") + # -- internals --------------------------------------------------------- def _read(self) -> list[_StoredRecord]: diff --git a/wsserver/server.py b/wsserver/server.py index 3488d43..3d4747c 100644 --- a/wsserver/server.py +++ b/wsserver/server.py @@ -511,6 +511,12 @@ async def ws_nodes(websocket: WebSocket) -> None: e2e_handshake_key = _e2e.derive_handshake_key( shared, salt_bytes, token ) + except TokenStoreError as exc: + logger.warning( + "E2E token resolution failed for %r: %s — falling back to legacy auth", + hello.node_name, exc, + ) + e2e_handshake_key = None except (ValueError, Exception) as exc: logger.warning("E2E key exchange failed for %r: %s", hello.node_name, exc) await _send_json_safe( @@ -524,15 +530,21 @@ async def ws_nodes(websocket: WebSocket) -> None: await _safe_close(websocket, CLOSE_PROTOCOL_VERSION) return - await _send_json_safe( - websocket, - _build_hello_ack( - hello.protocol_version, - session_id, - ecdh_pub=_e2e.encode_public_key(server_kp.public), - salt=_e2e.encode_b64(salt_bytes), - ), - ) + if e2e_handshake_key is not None: + await _send_json_safe( + websocket, + _build_hello_ack( + hello.protocol_version, + session_id, + ecdh_pub=_e2e.encode_public_key(server_kp.public), + salt=_e2e.encode_b64(salt_bytes), + ), + ) + else: + await _send_json_safe( + websocket, + _build_hello_ack(hello.protocol_version, session_id), + ) else: e2e_handshake_key = None await _send_json_safe( @@ -624,23 +636,20 @@ async def ws_nodes(websocket: WebSocket) -> None: else: # Legacy auth — validate token is_valid = False - try: - # tokens.validate() does a full read-decrypt-write cycle with - # os.fsync on the token store. That blocks the event loop if - # called directly from async code. Push it to a thread. - is_valid = await asyncio.to_thread( - token_store.validate, auth.node_name, auth.token - ) - except TokenStoreError as exc: - logging.getLogger(__name__).error( - "token store error during auth: %s", exc - ) - await _send_json_safe( - websocket, - _build_auth_err(reason="invalid_token", code=CLOSE_AUTH_FAILED), - ) - await _safe_close(websocket, CLOSE_AUTH_FAILED) - return + try: + is_valid = await asyncio.to_thread( + token_store.validate, auth.node_name, auth.token + ) + except TokenStoreError as exc: + logging.getLogger(__name__).error( + "token store error during auth: %s", exc + ) + await _send_json_safe( + websocket, + _build_auth_err(reason="invalid_token", code=CLOSE_AUTH_FAILED), + ) + await _safe_close(websocket, CLOSE_AUTH_FAILED) + return if not is_valid: await _send_json_safe(