Handshake under the pod-IP authority for WebSocket upgrades (#11)

DSH binds its auth cookie to the request authority (payload.authority ===
authority) and looks the cookie up under a name derived from the Host. The
upgrade presents Host = pod-IP:3080 (websockets derives it from the URI), so a
cookie minted for the browser authority is rejected with 401. The previous
pod-IP change fixed the Host fence (403) but not this, so the client kept
looping on "[connection] connection lost, retry #N".

Perform the WS handshake under the pod-IP authority too, so the cookie matches
the Host the upgrade carries. The HTTP path keeps the browser authority.

Also serialise the lazy Kubernetes client init: anyio.to_thread can enter it
concurrently, and a thread that observes _crd set while _core is still None
returns a None client. The live log showed the result:

  AttributeError: 'NoneType' object has no attribute 'read_namespaced_pod'

Co-authored-by: SREDevOpsAgents <agents@sredevops.org>

examples/dsh-agents-as-a-service/README.md

@@ -313,11 +313,13 @@ the DSH web GUI is fully interactive. Requirements:
   `60-ingress.yaml` is a starting point (Traefik handles upgrades by default),
 - idle suspension is disabled for users with in-flight proxied requests or open WebSockets, so
   a long-running chat stream is not cut off mid-turn,
-- the upstream upgrade dials the sandbox **pod IP** and forwards no `Host`/`Origin`: DSH's Host
-  fence only trusts its own bind address or a loopback/allow-listed host, and `websockets`
-  derives `Host` from the URI (it cannot be overridden). Forwarding the browser's origin here
-  makes DSH answer `403`, which shows up in the browser as an endless
-  `[connection] connection lost, retry #N` loop,
+- the upstream upgrade dials the sandbox **pod IP**, forwards no `Host`/`Origin`, and performs
+  the launch-token handshake under that same `pod-IP:3080` authority. DSH fences on `Host`
+  (only its bind address, loopback, or an allow-listed host passes) and binds its auth cookie to
+  the request authority (`payload.authority === authority`), while `websockets` derives `Host`
+  from the URI and cannot be overridden. A browser-host cookie or forwarded origin makes the
+  upgrade fail (`401`/`403`, surfaced as close `1011`), which the client retries forever as
+  `[connection] connection lost, retry #N`,
 - missing/invalid gateway auth is rejected right after the upgrade with close code `4401`
   (WebSocket auth has no HTTP `401` path); cross-origin or Origin-less upgrades with the cookie
   are rejected the same way (CSWSH protection), so clients must send a matching `Origin`;
@@ -341,7 +343,7 @@ redirects `/u/<user>` to the trailing-slash form with `308` so relative URLs res
 | `409` on signup | User already has a claim; use `POST /login` with the existing token or delete the user first |
 | Cookie set but browser keeps getting `401` | `COOKIE_SECURE=true` while serving over plain HTTP (secure cookies are not sent); use HTTPS or set `false` for local demos |
 | WebSocket closes immediately (`4401`) | Missing/expired gateway cookie, or the upgrade's `Origin` does not match `Host` (CSWSH check); DSH rejecting `Host` instead surfaces as close `1011` — add the host to `DSH_TRUSTED_HOSTS` |
-| Console spams `[connection] connection lost, retry #N` | The upstream WS upgrade is being rejected (DSH answers `403` to the forward, surfaced as close `1011`). The gateway dials the sandbox **pod IP** and forwards no `Host`/`Origin` for upgrades; a gateway image from before that fix, or a sandbox with an empty `status.podIPs`, hits the loop |
+| Console spams `[connection] connection lost, retry #N` | The upstream WS upgrade is being rejected, surfaced as close `1011`. The gateway dials the sandbox **pod IP**, forwards no `Host`/`Origin`, and handshakes under the same `pod-IP:3080` authority DSH binds its cookie to; an older gateway image, or a sandbox with an empty `status.podIPs`, hits this loop |
 | `manifest.webmanifest` returns `401` in the console | Browsers fetch the Web App Manifest without credentials, so the gateway's session cookie is absent and the request fails closed. The UI is unaffected; only the PWA install prompt is lost |
 | Signup is slow / `202` then `Waking` | Warm pool exhausted (`kubectl get sandboxes -l agents.x-k8s.io/warm-pool-sandbox`); scale the pool |
 | `403` from the Kubernetes API in gateway logs | The gateway ServiceAccount is missing claim/sandbox/pod-log verbs (see `50-gateway.yaml`) |

examples/dsh-agents-as-a-service/gateway/gateway.py

@@ -41,6 +41,7 @@ import json
 import os
 import re
 import secrets
+import threading
 import time
 from dataclasses import dataclass, field
 from urllib.parse import (
@@ -365,15 +366,21 @@ class Kubernetes:
         self.namespace = namespace
         self._crd: k8s_client.CustomObjectsApi | None = None
         self._core: k8s_client.CoreV1Api | None = None
+        # anyio.to_thread runs _call() on a worker pool, so the lazy init can be
+        # entered concurrently. Serialise it: without the lock one thread can
+        # observe _crd set while _core is still None and return a None client.
+        self._clients_lock = threading.Lock()
 
     def _clients(self):
-        if self._crd is None:
-            try:
-                k8s_config.load_incluster_config()
-            except k8s_config.ConfigException:
-                k8s_config.load_kube_config()
-            self._crd = k8s_client.CustomObjectsApi()
-            self._core = k8s_client.CoreV1Api()
+        if self._crd is None or self._core is None:
+            with self._clients_lock:
+                if self._crd is None or self._core is None:
+                    try:
+                        k8s_config.load_incluster_config()
+                    except k8s_config.ConfigException:
+                        k8s_config.load_kube_config()
+                    self._crd = k8s_client.CustomObjectsApi()
+                    self._core = k8s_client.CoreV1Api()
         return self._crd, self._core
 
     async def _call(self, which: str, method: str, *args, **kwargs):
@@ -1148,17 +1155,22 @@ against the JSON API. The browser session is an HttpOnly cookie.</p>
             try:
                 sandbox, fqdn = await self._resolve_sandbox(user, claim)
                 host = websocket.headers.get("host", "")
-                cookie, _, _ = await self.dsh_cookie(user, sandbox, fqdn, host)
                 path = websocket.path_params.get("path", "")
-                # DSH's Host fence accepts its own bind address and requires
-                # any Origin to have the same authority. websockets derives the
-                # Host header from the URI and cannot be overridden, so dial
-                # the pod IP instead of the service FQDN and forward no Origin;
-                # CSWSH is already rejected by same_origin() above.
+                # DSH's Host fence accepts its own bind address and binds the
+                # auth cookie to the request authority. websockets derives the
+                # Host header from the URI and cannot be overridden, so dial the
+                # pod IP (the bind address) and handshake under that same
+                # authority - a cookie minted for the browser authority would be
+                # rejected (401). Forward no Origin either: CSWSH is already
+                # rejected by same_origin() above.
                 upstream_host = self._pod_ip(sandbox) or fqdn
                 if ":" in upstream_host and not upstream_host.startswith("["):
                     upstream_host = f"[{upstream_host}]"  # IPv6 literal
-                url = f"ws://{upstream_host}:{self.settings.dsh_port}/{path}"
+                authority = f"{upstream_host}:{self.settings.dsh_port}"
+                cookie, _, _ = await self.dsh_cookie(
+                    user, sandbox, fqdn, authority
+                )
+                url = f"ws://{authority}/{path}"
                 query = filtered_query(websocket.url.query)
                 if query:
                     url = f"{url}?{query}"

examples/dsh-agents-as-a-service/gateway/test_gateway.py

@@ -10,6 +10,7 @@ use Starlette's ``TestClient`` (which runs the lifespan) against a fake
 import asyncio
 import hashlib
 import json
+import threading
 import time
 from types import SimpleNamespace
 
@@ -529,6 +530,42 @@ async def test_kubernetes_get_pod_normalises_the_v1pod_model(monkeypatch):
     assert pod["metadata"]["uid"] == "pod-uid-9"
 
 
+def test_kubernetes_clients_initialise_once_under_concurrency(monkeypatch):
+    """_call runs on a worker pool: the lazy init must not hand out a None client."""
+    import gateway as gw
+    from gateway import Kubernetes
+
+    created = []
+
+    class Crd:
+        def __init__(self):
+            created.append("crd")
+
+    class SlowCore:
+        def __init__(self):
+            time.sleep(0.05)  # widen the window between _crd and _core
+            created.append("core")
+
+    monkeypatch.setattr(gw.k8s_client, "CustomObjectsApi", Crd)
+    monkeypatch.setattr(gw.k8s_client, "CoreV1Api", SlowCore)
+    monkeypatch.setattr(gw.k8s_config, "load_incluster_config", lambda: None)
+
+    k8s = Kubernetes("test-ns")
+    results = []
+
+    def worker():
+        results.append(k8s._clients())
+
+    threads = [threading.Thread(target=worker) for _ in range(8)]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join()
+
+    assert all(core is not None for _, core in results)
+    assert created.count("core") == 1
+
+
 # --- health / landing / redirect ------------------------------------------