From b5821bca6d4243045929d49aee9483a3a9703f7e Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Fri, 24 Apr 2026 09:40:15 +0100 Subject: [PATCH] Fix socket token handling --- .../trustgraph/api/async_socket_client.py | 56 +++++++++++++++++-- .../trustgraph/api/socket_client.py | 53 ++++++++++++++++-- 2 files changed, 99 insertions(+), 10 deletions(-) diff --git a/trustgraph-base/trustgraph/api/async_socket_client.py b/trustgraph-base/trustgraph/api/async_socket_client.py index e5d553ea..ca9146b9 100644 --- a/trustgraph-base/trustgraph/api/async_socket_client.py +++ b/trustgraph-base/trustgraph/api/async_socket_client.py @@ -49,21 +49,67 @@ class AsyncSocketClient: return f"ws://{url}" def _build_ws_url(self): - ws_url = f"{self.url.rstrip('/')}/api/v1/socket" - if self.token: - ws_url = f"{ws_url}?token={self.token}" - return ws_url + # /api/v1/socket uses the first-frame auth protocol — the + # token is sent as the first frame after connecting rather + # than in the URL. This avoids browser issues with 401 on + # the WebSocket handshake and lets long-lived sockets + # refresh credentials mid-session. + return f"{self.url.rstrip('/')}/api/v1/socket" async def connect(self): - """Establish the persistent websocket connection.""" + """Establish the persistent websocket connection and run the + first-frame auth handshake.""" if self._connected: return + if not self.token: + raise ProtocolException( + "AsyncSocketClient requires a token for first-frame " + "auth against /api/v1/socket" + ) + ws_url = self._build_ws_url() self._connect_cm = websockets.connect( ws_url, ping_interval=20, ping_timeout=self.timeout ) self._socket = await self._connect_cm.__aenter__() + + # First-frame auth: send {"type":"auth","token":"..."} and + # wait for auth-ok / auth-failed. Run before starting the + # reader task so the response isn't consumed by the reader's + # id-based routing. + await self._socket.send(json.dumps({ + "type": "auth", "token": self.token, + })) + try: + raw = await asyncio.wait_for( + self._socket.recv(), timeout=self.timeout, + ) + except asyncio.TimeoutError: + await self._socket.close() + raise ProtocolException("Timeout waiting for auth response") + + try: + resp = json.loads(raw) + except Exception: + await self._socket.close() + raise ProtocolException( + f"Unexpected non-JSON auth response: {raw!r}" + ) + + if resp.get("type") == "auth-ok": + self.workspace = resp.get("workspace", self.workspace) + elif resp.get("type") == "auth-failed": + await self._socket.close() + raise ProtocolException( + f"auth failure: {resp.get('error', 'unknown')}" + ) + else: + await self._socket.close() + raise ProtocolException( + f"Unexpected auth response: {resp!r}" + ) + self._connected = True self._reader_task = asyncio.create_task(self._reader()) diff --git a/trustgraph-base/trustgraph/api/socket_client.py b/trustgraph-base/trustgraph/api/socket_client.py index 4eade3e8..aeb15f85 100644 --- a/trustgraph-base/trustgraph/api/socket_client.py +++ b/trustgraph-base/trustgraph/api/socket_client.py @@ -112,10 +112,10 @@ class SocketClient: return f"ws://{url}" def _build_ws_url(self): - ws_url = f"{self.url.rstrip('/')}/api/v1/socket" - if self.token: - ws_url = f"{ws_url}?token={self.token}" - return ws_url + # /api/v1/socket uses the first-frame auth protocol — the + # token is sent as the first frame after connecting rather + # than in the URL. + return f"{self.url.rstrip('/')}/api/v1/socket" def _get_loop(self): """Get or create the event loop, reusing across calls.""" @@ -132,15 +132,58 @@ class SocketClient: return self._loop async def _ensure_connected(self): - """Lazily establish the persistent websocket connection.""" + """Lazily establish the persistent websocket connection and + run the first-frame auth handshake.""" if self._connected: return + if not self.token: + raise ProtocolException( + "SocketClient requires a token for first-frame auth " + "against /api/v1/socket" + ) + ws_url = self._build_ws_url() self._connect_cm = websockets.connect( ws_url, ping_interval=20, ping_timeout=self.timeout ) self._socket = await self._connect_cm.__aenter__() + + # First-frame auth — run before starting the reader so the + # auth-ok / auth-failed response isn't consumed by the reader + # loop's id-based routing. + await self._socket.send(json.dumps({ + "type": "auth", "token": self.token, + })) + try: + raw = await asyncio.wait_for( + self._socket.recv(), timeout=self.timeout, + ) + except asyncio.TimeoutError: + await self._socket.close() + raise ProtocolException("Timeout waiting for auth response") + + try: + resp = json.loads(raw) + except Exception: + await self._socket.close() + raise ProtocolException( + f"Unexpected non-JSON auth response: {raw!r}" + ) + + if resp.get("type") == "auth-ok": + self.workspace = resp.get("workspace", self.workspace) + elif resp.get("type") == "auth-failed": + await self._socket.close() + raise ProtocolException( + f"auth failure: {resp.get('error', 'unknown')}" + ) + else: + await self._socket.close() + raise ProtocolException( + f"Unexpected auth response: {resp!r}" + ) + self._connected = True self._reader_task = asyncio.create_task(self._reader())