@castlekit/castle 0.1.5 → 0.3.0

package/src/lib/gateway-connection.ts

@@ -1,7 +1,11 @@
 import WebSocket from "ws";
 import { EventEmitter } from "events";
 import { randomUUID } from "crypto";
+import { execSync } from "child_process";
+import { readFileSync, realpathSync } from "fs";
+import { dirname, join } from "path";
 import { getGatewayUrl, readOpenClawToken, readConfig, configExists } from "./config";
+import { getOrCreateIdentity, signDeviceAuth, saveDeviceToken, getDeviceToken, clearDeviceToken } from "./device-identity";
 
 // ============================================================================
 // Types
@@ -44,7 +48,7 @@ interface PendingRequest {
   timer: ReturnType<typeof setTimeout>;
 }
 
-export type ConnectionState = "disconnected" | "connecting" | "connected" | "error";
+export type ConnectionState = "disconnected" | "connecting" | "connected" | "pairing" | "error";
 
 export interface GatewayEvent {
   event: string;
@@ -52,6 +56,21 @@ export interface GatewayEvent {
   seq?: number;
 }
 
+// ============================================================================
+// Helpers
+// ============================================================================
+
+const MAX_MESSAGE_SIZE = 10 * 1024 * 1024; // 10MB
+const MAX_NONCE_LENGTH = 1024;
+
+/** Strip tokens and key material from strings before logging. */
+function sanitize(str: string): string {
+  return str
+    .replace(/rew_[a-f0-9]+/gi, "rew_***")
+    .replace(/-----BEGIN [A-Z ]+-----[\s\S]*?-----END [A-Z ]+-----/g, "[REDACTED KEY]")
+    .replace(/[a-f0-9]{32,}/gi, (m) => m.slice(0, 8) + "***");
+}
+
 // ============================================================================
 // Singleton Gateway Connection
 // ============================================================================
@@ -69,6 +88,10 @@ class GatewayConnection extends EventEmitter {
   private _serverInfo: { version?: string; connId?: string } = {};
   private _features: { methods?: string[]; events?: string[] } = {};
   private shouldReconnect = true;
+  // Device auth: set to true after "device identity mismatch" to fall back to token-only
+  private _skipDeviceAuth = false;
+  // Avatar URL mapping: hash → original Gateway URL (for proxying)
+  private _avatarUrls = new Map<string, string>();
 
   get state(): ConnectionState {
     return this._state;
@@ -78,6 +101,16 @@ class GatewayConnection extends EventEmitter {
     return this._serverInfo;
   }
 
+  /** Store a mapping from avatar hash to its original Gateway URL */
+  setAvatarUrl(hash: string, originalUrl: string): void {
+    this._avatarUrls.set(hash, originalUrl);
+  }
+
+  /** Get the original Gateway URL for an avatar hash */
+  getAvatarUrl(hash: string): string | null {
+    return this._avatarUrls.get(hash) || null;
+  }
+
   get isConnected(): boolean {
     return this._state === "connected";
   }
@@ -94,6 +127,7 @@ class GatewayConnection extends EventEmitter {
   start(): void {
     if (this._state === "connecting" || this._state === "connected") return;
     this.shouldReconnect = true;
+    this._skipDeviceAuth = false; // Reset — try device auth on fresh start
     this.connect();
   }
 
@@ -119,7 +153,7 @@ class GatewayConnection extends EventEmitter {
     try {
       this.ws = new WebSocket(url);
     } catch (err) {
-      console.error("[Gateway] Failed to create WebSocket:", err);
+      console.error("[Gateway] Failed to create WebSocket:", (err as Error).message);
       this._state = "error";
       this.emit("stateChange", this._state);
       this.scheduleReconnect();
@@ -132,66 +166,224 @@ class GatewayConnection extends EventEmitter {
       this.scheduleReconnect();
     }, this.connectTimeout);
 
+    // Load device identity for auth (skip if previously rejected)
+    let deviceIdentity: { deviceId: string; publicKey: string } | null = null;
+    if (!this._skipDeviceAuth) {
+      try {
+        const identity = getOrCreateIdentity();
+        deviceIdentity = { deviceId: identity.deviceId, publicKey: identity.publicKey };
+      } catch (err) {
+        console.warn("[Gateway] Could not load device identity:", (err as Error).message);
+      }
+    } else {
+      console.log("[Gateway] Device auth disabled — using token-only");
+    }
+
+    // Check for a saved device token from previous pairing
+    const savedDeviceToken = getDeviceToken();
+
     this.ws.on("open", () => {
-      // Build connect handshake
       const connectId = randomUUID();
-      const connectFrame: RequestFrame = {
-        type: "req",
-        id: connectId,
-        method: "connect",
-        params: {
+
+      // Build connect frame per Gateway protocol:
+      // - Initial connect: token-only (no device field)
+      // - Challenge response: include device { id, publicKey, signature, signedAt, nonce }
+      // - Reconnect with deviceToken: use deviceToken as auth.token
+      //
+      // The device field is ONLY sent when responding to connect.challenge,
+      // because the Gateway requires signature + signedAt whenever device is present.
+      let challengeReceived = false;
+
+      // Connection identity constants — must match what's signed
+      const CLIENT_ID = "gateway-client";
+      const CLIENT_MODE = "backend";
+      const ROLE = "operator";
+      const SCOPES = ["operator.admin"];
+      const authToken = savedDeviceToken || token;
+
+      const buildConnectFrame = (challenge?: {
+        nonce: string;
+        signature: string;
+        signedAt: number;
+      }): RequestFrame => {
+        const params: Record<string, unknown> = {
           minProtocol: 3,
           maxProtocol: 3,
           client: {
-            id: "gateway-client",
+            id: CLIENT_ID,
             displayName: "Castle",
             version: "0.0.1",
             platform: process.platform,
-            mode: "backend",
+            mode: CLIENT_MODE,
           },
-          auth: { token },
-          role: "operator",
-          scopes: ["operator.admin"],
+          auth: { token: authToken },
+          role: ROLE,
+          scopes: SCOPES,
           caps: [],
-        },
+        };
+
+        // Only include device when responding to a challenge (with signature)
+        if (challenge && deviceIdentity) {
+          params.device = {
+            id: deviceIdentity.deviceId,
+            publicKey: deviceIdentity.publicKey,
+            signature: challenge.signature,
+            nonce: challenge.nonce,
+            signedAt: challenge.signedAt,
+          };
+        }
+
+        return { type: "req", id: connectId, method: "connect", params };
       };
 
-      // Handle handshake messages (may include connect.challenge events)
+      // Handle handshake messages
       const onHandshakeMessage = (data: WebSocket.RawData) => {
+        const rawSize = Buffer.isBuffer(data) ? data.length : Buffer.byteLength(data.toString());
+        if (rawSize > MAX_MESSAGE_SIZE) {
+          console.error(`[Gateway] Message too large (${rawSize} bytes) — ignoring`);
+          return;
+        }
         try {
           const msg = JSON.parse(data.toString());
 
-          // Handle connect.challenge event -- re-send connect with nonce
+          // Handle connect.challenge — sign nonce and re-send connect with device identity
           if (msg.type === "event" && msg.event === "connect.challenge") {
-            // For now we don't support device-signed nonce challenges
-            // Just proceed with the connect
+            const nonce = msg.payload?.nonce;
+            if (nonce && typeof nonce === "string" && nonce.length > MAX_NONCE_LENGTH) {
+              console.error(`[Gateway] Challenge nonce too large (${nonce.length} bytes) — rejecting`);
+              return;
+            }
+            if (nonce && typeof nonce === "string" && deviceIdentity) {
+              challengeReceived = true;
+              console.log("[Gateway] Challenge received — signing with device key");
+              try {
+                const { signature, signedAt } = signDeviceAuth({
+                  nonce,
+                  clientId: CLIENT_ID,
+                  clientMode: CLIENT_MODE,
+                  role: ROLE,
+                  scopes: SCOPES,
+                  token: authToken!,
+                });
+                const challengeFrame = buildConnectFrame({ nonce, signature, signedAt });
+                this.ws?.send(JSON.stringify(challengeFrame));
+              } catch (err) {
+                console.error("[Gateway] Failed to sign challenge:", (err as Error).message);
+                challengeReceived = false;
+              }
+            } else {
+              console.log("[Gateway] Challenge received but no device identity — skipping");
+            }
+            return;
+          }
+
+          // Handle device.pairing.required — waiting for operator approval
+          if (msg.type === "event" && msg.event === "device.pairing.required") {
+            console.log("[Gateway] Device pairing approval required");
+            console.log("[Gateway] Approve this device in your OpenClaw dashboard to continue...");
+            this._state = "pairing";
+            this.emit("stateChange", this._state);
+            this.emit("pairingRequired", msg.payload);
+            this.emit("gatewayEvent", {
+              event: msg.event,
+              payload: msg.payload,
+              seq: msg.seq,
+            } as GatewayEvent);
             return;
           }
 
-          // Standard response frame to our connect request
+          // Handle device.pairing.approved — save the device token
+          if (msg.type === "event" && msg.event === "device.pairing.approved") {
+            const approvedToken = msg.payload?.deviceToken;
+            if (approvedToken && typeof approvedToken === "string") {
+              console.log("[Gateway] Device pairing approved — saving token");
+              try {
+                saveDeviceToken(approvedToken, url);
+              } catch (err) {
+                console.error("[Gateway] Failed to save device token:", (err as Error).message);
+              }
+            }
+            this.emit("pairingApproved", msg.payload);
+            this.emit("gatewayEvent", {
+              event: msg.event,
+              payload: msg.payload,
+              seq: msg.seq,
+            } as GatewayEvent);
+            return;
+          }
+
+          // Standard response to our connect request
           if (msg.type === "res" && msg.id === connectId) {
+            // If we already sent a challenge response, ignore error from the
+            // initial (token-only) connect — the signed response is in flight.
+            if (!msg.ok && challengeReceived) {
+              console.log("[Gateway] Ignoring error from initial connect — challenge response pending");
+              return;
+            }
+
             clearTimeout(connectTimer);
 
             if (msg.ok) {
-              // hello-ok is embedded in the payload
               const helloOk = msg.payload || {};
+
+              // Save deviceToken from hello-ok (may be at payload.auth.deviceToken
+              // or payload.deviceToken depending on Gateway version)
+              const helloDeviceToken =
+                helloOk.auth?.deviceToken || helloOk.deviceToken;
+              if (helloDeviceToken && typeof helloDeviceToken === "string") {
+                try {
+                  saveDeviceToken(helloDeviceToken, url);
+                  console.log("[Gateway] Device token received and saved");
+                } catch (err) {
+                  console.error("[Gateway] Failed to save device token:", (err as Error).message);
+                }
+              }
+
               this._state = "connected";
               this._serverInfo = helloOk.server || {};
+              // If Gateway reports "dev" or missing version, read from installed package.json
+              if (!this._serverInfo.version || this._serverInfo.version === "dev") {
+                try {
+                  const bin = execSync("which openclaw", { timeout: 2000, encoding: "utf-8" }).trim();
+                  const real = realpathSync(bin);
+                  const pkgPath = join(dirname(real), "package.json");
+                  const pkg = JSON.parse(readFileSync(pkgPath, "utf-8"));
+                  if (pkg.version) this._serverInfo.version = pkg.version;
+                } catch { /* keep whatever the Gateway reported */ }
+              }
               this._features = helloOk.features || {};
               this.reconnectAttempts = 0;
               this.emit("stateChange", this._state);
               this.emit("connected", helloOk);
-              console.log(`[Gateway] Connected to OpenClaw v${helloOk.server?.version || "unknown"}`);
+              console.log(`[Gateway] Connected to OpenClaw ${this._serverInfo.version || "unknown"}`);
               // Switch to normal message handler
               this.ws?.off("message", onHandshakeMessage);
               this.ws?.on("message", this.onMessage.bind(this));
             } else {
+              const errCode = msg.error?.code;
               const errMsg = msg.error?.message || "Connect rejected";
-              console.error(`[Gateway] Connect failed: ${errMsg}`);
+              console.error(`[Gateway] Connect failed: ${sanitize(errMsg)}`);
               this.ws?.off("message", onHandshakeMessage);
               this.cleanup();
-              // Don't reconnect on auth errors
-              if (msg.error?.code === "auth_failed") {
+
+              if (errCode === "auth_failed") {
+                // If we were using a stale device token, clear it and retry
+                // with the gateway token instead
+                if (savedDeviceToken) {
+                  console.log("[Gateway] Device token rejected — clearing and retrying with gateway token");
+                  clearDeviceToken();
+                  this._skipDeviceAuth = false;
+                  this.reconnectAttempts = 0;
+                  this.scheduleReconnect();
+                } else {
+                  // Real auth failure — gateway token is invalid
+                  this._state = "error";
+                  this.emit("stateChange", this._state);
+                  this.emit("authError", msg.error);
+                }
+              } else if (errCode === "protocol_mismatch" || errCode === "protocol_unsupported") {
+                // Permanent failure — don't retry with the same protocol
+                console.error("[Gateway] Protocol version not supported by this Gateway");
                 this._state = "error";
                 this.emit("stateChange", this._state);
                 this.emit("authError", msg.error);
@@ -202,7 +394,7 @@ class GatewayConnection extends EventEmitter {
             return;
           }
 
-          // Forward any events that arrive during handshake
+          // Forward any other events during handshake
           if (msg.type === "event") {
             this.emit("gatewayEvent", {
               event: msg.event,
@@ -211,12 +403,25 @@ class GatewayConnection extends EventEmitter {
             } as GatewayEvent);
           }
         } catch (err) {
-          console.error("[Gateway] Failed to parse handshake message:", err);
+          console.error("[Gateway] Failed to parse handshake message:", (err as Error).message);
         }
       };
 
       this.ws!.on("message", onHandshakeMessage);
-      this.ws!.send(JSON.stringify(connectFrame));
+      this.ws!.send(JSON.stringify(buildConnectFrame()));
+    });
+
+    // Handle WebSocket upgrade failures (e.g. corporate proxies blocking WS)
+    this.ws.on("unexpected-response", (_req, res) => {
+      clearTimeout(connectTimer);
+      console.error(`[Gateway] WebSocket upgrade failed (HTTP ${res.statusCode})`);
+      if (res.statusCode === 401 || res.statusCode === 403) {
+        console.error("[Gateway] Authentication rejected at HTTP level");
+      } else if (res.statusCode === 502 || res.statusCode === 503) {
+        console.error("[Gateway] Gateway may be behind a reverse proxy that doesn't support WebSocket");
+      }
+      this.cleanup();
+      this.scheduleReconnect();
     });
 
     this.ws.on("error", (err) => {
@@ -229,15 +434,32 @@ class GatewayConnection extends EventEmitter {
     this.ws.on("close", (code, reason) => {
       clearTimeout(connectTimer);
       const wasConnected = this._state === "connected";
+      const reasonStr = reason?.toString() || "none";
       this.cleanup();
+
+      // If Gateway rejected device auth (identity mismatch, bad signature, etc),
+      // fall back to token-only auth
+      if (code === 1008 && !this._skipDeviceAuth) {
+        console.log(`[Gateway] Device auth rejected (${reasonStr}) — retrying with token-only auth`);
+        this._skipDeviceAuth = true;
+        this.reconnectAttempts = 0;
+        this.scheduleReconnect();
+        return;
+      }
+
       if (wasConnected) {
-        console.log(`[Gateway] Disconnected (code: ${code}, reason: ${reason?.toString() || "none"})`);
+        console.log(`[Gateway] Disconnected (code: ${code}, reason: ${reasonStr})`);
       }
       this.scheduleReconnect();
     });
   }
 
   private onMessage(data: WebSocket.RawData): void {
+    const rawSize = Buffer.isBuffer(data) ? data.length : Buffer.byteLength(data.toString());
+    if (rawSize > MAX_MESSAGE_SIZE) {
+      console.error(`[Gateway] Message too large (${rawSize} bytes) — ignoring`);
+      return;
+    }
     let msg: GatewayFrame;
     try {
       msg = JSON.parse(data.toString());
@@ -308,14 +530,19 @@ class GatewayConnection extends EventEmitter {
   // --------------------------------------------------------------------------
 
   private resolveToken(): string | null {
-    // 1. Castle config token
+    // 1. Saved device token from previous pairing
+    const deviceToken = getDeviceToken();
+    if (deviceToken) return deviceToken;
+
+    // 2. Castle config token
     if (configExists()) {
       const config = readConfig();
       if (config.openclaw.gateway_token) {
         return config.openclaw.gateway_token;
       }
     }
-    // 2. Auto-detect from OpenClaw config
+
+    // 3. Auto-detect from OpenClaw config
     return readOpenClawToken();
   }
 
@@ -335,7 +562,7 @@ class GatewayConnection extends EventEmitter {
       this.pending.delete(id);
     }
 
-    if (this._state !== "error") {
+    if (this._state !== "error" && this._state !== "pairing") {
       this._state = "disconnected";
       this.emit("stateChange", this._state);
     }
@@ -366,16 +593,26 @@ class GatewayConnection extends EventEmitter {
 }
 
 // ============================================================================
-// Singleton export
+// Singleton export (uses globalThis to survive HMR in dev mode)
 // ============================================================================
 
-let _gateway: GatewayConnection | null = null;
+const GATEWAY_KEY = "__castle_gateway__" as const;
+
+function getGlobalGateway(): GatewayConnection | null {
+  return (globalThis as Record<string, unknown>)[GATEWAY_KEY] as GatewayConnection | null ?? null;
+}
+
+function setGlobalGateway(gw: GatewayConnection): void {
+  (globalThis as Record<string, unknown>)[GATEWAY_KEY] = gw;
+}
 
 export function getGateway(): GatewayConnection {
-  if (!_gateway) {
-    _gateway = new GatewayConnection();
+  let gw = getGlobalGateway();
+  if (!gw) {
+    gw = new GatewayConnection();
+    setGlobalGateway(gw);
   }
-  return _gateway;
+  return gw;
 }
 
 /**