apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-04-25 00:16:23 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Fix socket token handling

Browse source
This commit is contained in:
Cyber MacGeddon 2026-04-24 09:40:15 +01:00
parent 2a071b12d0
commit b5821bca6d
2 changed files with 99 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

56
trustgraph-base/trustgraph/api/async_socket_client.py
View file

@ -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())

53
trustgraph-base/trustgraph/api/socket_client.py
View file

@ -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())