fix(gateway): reject RPCs from invalidated device-token clients during rotation/revoke race

davidangularme · davidangularme · commit b304ccd07987 · 2026-04-23T20:35:16.000+03:00
device.token.rotate, device.token.revoke and device.pair.remove all
respond 200 OK to the admin, then schedule disconnectClientsForDevice
via queueMicrotask so the response can flush before the socket close.
That microtask window plus the absence of a per-RPC re-check for
device-token auth (unlike shared-auth, which gets checked at
message-handler.ts:1444-1458) created a race: an attacker with RPCs
already pipelined in the WS socket buffer could land a few more
authenticated operations with the rotated/revoked token before the
socket actually closed.

Fix: add a cheap in-memory 'invalidated' flag on GatewayWsClient and
mark it synchronously *before* responding in the three handlers. Add
a mirror check at the start of the per-RPC dispatch that force-closes
the client if the flag is set, regardless of whether socket.close()
has taken effect yet. Disconnect still happens via queueMicrotask so
the admin's rotate/revoke response flushes normally.

Introduces context.invalidateClientsForDevice(deviceId, opts) as a
sync companion to the existing disconnectClientsForDevice. Also
defense-in-depth: disconnectClientsForDevice now sets the flag too,
so any other caller of the hard-disconnect path gets the per-RPC
gate for free.
diff --git a/src/gateway/server-methods/devices.test.ts b/src/gateway/server-methods/devices.test.ts
@@ -69,6 +69,7 @@ function createOptions(
     context: {
       broadcast: vi.fn(),
       disconnectClientsForDevice: vi.fn(),
+      invalidateClientsForDevice: vi.fn(),
       logGateway: {
         debug: vi.fn(),
         error: vi.fn(),
@@ -131,6 +132,28 @@ describe("deviceHandlers", () => {
     );
   });
 
+  it("invalidates affected clients synchronously before responding to device.pair.remove", async () => {
+    removePairedDeviceMock.mockResolvedValue({ deviceId: "device-1", removedAtMs: 123 });
+    const opts = createOptions("device.pair.remove", { deviceId: "device-1" });
+    const respond = opts.respond as ReturnType<typeof vi.fn>;
+    const invalidate = opts.context.invalidateClientsForDevice as ReturnType<typeof vi.fn>;
+    const disconnect = opts.context.disconnectClientsForDevice as ReturnType<typeof vi.fn>;
+
+    respond.mockImplementation(() => {
+      // At the moment the response is emitted, the client flag must already
+      // be set so any RPCs pipelined in the WS buffer are rejected.
+      expect(invalidate).toHaveBeenCalledWith("device-1", { reason: "device-pair-removed" });
+      // The hard close is deferred to the microtask to let the response flush.
+      expect(disconnect).not.toHaveBeenCalled();
+    });
+
+    await deviceHandlers["device.pair.remove"](opts);
+    await Promise.resolve();
+
+    expect(respond).toHaveBeenCalled();
+    expect(disconnect).toHaveBeenCalledWith("device-1");
+  });
+
   it("does not disconnect clients when device removal fails", async () => {
     removePairedDeviceMock.mockResolvedValue(null);
     const opts = createOptions("device.pair.remove", { deviceId: "device-1" });
@@ -289,6 +312,59 @@ describe("deviceHandlers", () => {
     );
   });
 
+  it("invalidates affected clients synchronously before responding to device.token.rotate", async () => {
+    mockPairedOperatorDevice();
+    mockRotateOperatorTokenSuccess();
+    const opts = createOptions(
+      "device.token.rotate",
+      { deviceId: "device-1", role: "operator", scopes: ["operator.pairing"] },
+      { client: createClient(["operator.pairing"], "device-1", { isDeviceTokenAuth: true }) },
+    );
+    const respond = opts.respond as ReturnType<typeof vi.fn>;
+    const invalidate = opts.context.invalidateClientsForDevice as ReturnType<typeof vi.fn>;
+    const disconnect = opts.context.disconnectClientsForDevice as ReturnType<typeof vi.fn>;
+
+    respond.mockImplementation(() => {
+      expect(invalidate).toHaveBeenCalledWith("device-1", {
+        role: "operator",
+        reason: "device-token-rotated",
+      });
+      expect(disconnect).not.toHaveBeenCalled();
+    });
+
+    await deviceHandlers["device.token.rotate"](opts);
+    await Promise.resolve();
+
+    expect(respond).toHaveBeenCalled();
+    expect(disconnect).toHaveBeenCalledWith("device-1", { role: "operator" });
+  });
+
+  it("invalidates affected clients synchronously before responding to device.token.revoke", async () => {
+    revokeDeviceTokenMock.mockResolvedValue({ role: "operator", revokedAtMs: 456 });
+    const opts = createOptions(
+      "device.token.revoke",
+      { deviceId: "device-1", role: "operator" },
+      { client: createClient(["operator.pairing"], "device-1", { isDeviceTokenAuth: true }) },
+    );
+    const respond = opts.respond as ReturnType<typeof vi.fn>;
+    const invalidate = opts.context.invalidateClientsForDevice as ReturnType<typeof vi.fn>;
+    const disconnect = opts.context.disconnectClientsForDevice as ReturnType<typeof vi.fn>;
+
+    respond.mockImplementation(() => {
+      expect(invalidate).toHaveBeenCalledWith("device-1", {
+        role: "operator",
+        reason: "device-token-revoked",
+      });
+      expect(disconnect).not.toHaveBeenCalled();
+    });
+
+    await deviceHandlers["device.token.revoke"](opts);
+    await Promise.resolve();
+
+    expect(respond).toHaveBeenCalled();
+    expect(disconnect).toHaveBeenCalledWith("device-1", { role: "operator" });
+  });
+
   it("treats normalized device ids as self-owned for token rotation", async () => {
     mockPairedOperatorDevice();
     mockRotateOperatorTokenSuccess();
diff --git a/src/gateway/server-methods/devices.ts b/src/gateway/server-methods/devices.ts
@@ -315,6 +315,13 @@ export const deviceHandlers: GatewayRequestHandlers = {
       return;
     }
     context.logGateway.info(`device pairing removed device=${removed.deviceId}`);
+    // Mark affected clients invalid *before* responding so any RPCs already
+    // pipelined into their WS socket buffer are rejected at the per-request
+    // dispatch check, closing the race between queueMicrotask-scheduled
+    // disconnect and inflight frames.
+    context.invalidateClientsForDevice?.(removed.deviceId, {
+      reason: "device-pair-removed",
+    });
     respond(true, removed, undefined);
     queueMicrotask(() => {
       context.disconnectClientsForDevice?.(removed.deviceId);
@@ -410,6 +417,14 @@ export const deviceHandlers: GatewayRequestHandlers = {
     context.logGateway.info(
       `device token rotated device=${deviceId} role=${entry.role} scopes=${entry.scopes.join(",")}`,
     );
+    // Mark affected clients invalid *before* responding so any RPCs already
+    // pipelined into their WS socket buffer are rejected at the per-request
+    // dispatch check, closing the race between queueMicrotask-scheduled
+    // disconnect and inflight frames.
+    context.invalidateClientsForDevice?.(deviceId.trim(), {
+      role: entry.role,
+      reason: "device-token-rotated",
+    });
     respond(
       true,
       {
@@ -459,6 +474,14 @@ export const deviceHandlers: GatewayRequestHandlers = {
     }
     const normalizedDeviceId = deviceId.trim();
     context.logGateway.info(`device token revoked device=${normalizedDeviceId} role=${entry.role}`);
+    // Mark affected clients invalid *before* responding so any RPCs already
+    // pipelined into their WS socket buffer are rejected at the per-request
+    // dispatch check, closing the race between queueMicrotask-scheduled
+    // disconnect and inflight frames.
+    context.invalidateClientsForDevice?.(normalizedDeviceId, {
+      role: entry.role,
+      reason: "device-token-revoked",
+    });
     respond(
       true,
       {
diff --git a/src/gateway/server-methods/shared-types.ts b/src/gateway/server-methods/shared-types.ts
@@ -59,6 +59,10 @@ export type GatewayRequestContext = {
   hasConnectedMobileNode: () => boolean;
   hasExecApprovalClients?: (excludeConnId?: string) => boolean;
   disconnectClientsForDevice?: (deviceId: string, opts?: { role?: string }) => void;
+  invalidateClientsForDevice?: (
+    deviceId: string,
+    opts?: { role?: string; reason?: string },
+  ) => void;
   disconnectClientsUsingSharedGatewayAuth?: () => void;
   enforceSharedGatewayAuthGenerationForConfigWrite?: (nextConfig: OpenClawConfig) => void;
   nodeRegistry: NodeRegistry;
diff --git a/src/gateway/server-request-context.test.ts b/src/gateway/server-request-context.test.ts
@@ -1,6 +1,72 @@
 import { describe, expect, it, vi } from "vitest";
 import type { GatewayServerLiveState } from "./server-live-state.js";
-import { createGatewayRequestContext } from "./server-request-context.js";
+import {
+  createGatewayRequestContext,
+  type GatewayRequestContextParams,
+} from "./server-request-context.js";
+
+function makeContextParams(
+  overrides: Partial<GatewayRequestContextParams> = {},
+): GatewayRequestContextParams {
+  const runtimeState: Pick<GatewayServerLiveState, "cronState"> = {
+    cronState: {
+      cron: { start: vi.fn(), stop: vi.fn() } as never,
+      storePath: "/tmp/cron",
+      cronEnabled: true,
+    },
+  };
+  return {
+    deps: {} as never,
+    runtimeState,
+    execApprovalManager: undefined,
+    pluginApprovalManager: undefined,
+    loadGatewayModelCatalog: vi.fn(async () => []),
+    getHealthCache: vi.fn(() => null),
+    refreshHealthSnapshot: vi.fn(async () => ({}) as never),
+    logHealth: { error: vi.fn() },
+    logGateway: { warn: vi.fn(), info: vi.fn(), error: vi.fn() } as never,
+    incrementPresenceVersion: vi.fn(() => 1),
+    getHealthVersion: vi.fn(() => 1),
+    broadcast: vi.fn(),
+    broadcastToConnIds: vi.fn(),
+    nodeSendToSession: vi.fn(),
+    nodeSendToAllSubscribed: vi.fn(),
+    nodeSubscribe: vi.fn(),
+    nodeUnsubscribe: vi.fn(),
+    nodeUnsubscribeAll: vi.fn(),
+    hasConnectedMobileNode: vi.fn(() => false),
+    clients: new Set(),
+    enforceSharedGatewayAuthGenerationForConfigWrite: vi.fn(),
+    nodeRegistry: {} as never,
+    agentRunSeq: new Map(),
+    chatAbortControllers: new Map(),
+    chatAbortedRuns: new Map(),
+    chatRunBuffers: new Map(),
+    chatDeltaSentAt: new Map(),
+    chatDeltaLastBroadcastLen: new Map(),
+    addChatRun: vi.fn(),
+    removeChatRun: vi.fn(),
+    subscribeSessionEvents: vi.fn(),
+    unsubscribeSessionEvents: vi.fn(),
+    subscribeSessionMessageEvents: vi.fn(),
+    unsubscribeSessionMessageEvents: vi.fn(),
+    unsubscribeAllSessionEvents: vi.fn(),
+    getSessionEventSubscriberConnIds: vi.fn(() => new Set<string>()),
+    registerToolEventRecipient: vi.fn(),
+    dedupe: new Map(),
+    wizardSessions: new Map(),
+    findRunningWizard: vi.fn(() => null),
+    purgeWizardSession: vi.fn(),
+    getRuntimeSnapshot: vi.fn(() => ({}) as never),
+    startChannel: vi.fn(async () => undefined),
+    stopChannel: vi.fn(async () => undefined),
+    markChannelLoggedOut: vi.fn(),
+    wizardRunner: vi.fn(async () => undefined),
+    broadcastVoiceWakeChanged: vi.fn(),
+    unavailableGatewayMethods: new Set(),
+    ...overrides,
+  };
+}
 
 describe("createGatewayRequestContext", () => {
   it("reads cron state live from runtime state", () => {
@@ -77,4 +143,66 @@ describe("createGatewayRequestContext", () => {
     expect(context.cron).toBe(cronB);
     expect(context.cronStorePath).toBe("/tmp/cron-b");
   });
+
+  it("invalidateClientsForDevice sets the flag on matching clients without closing the socket", () => {
+    const target = {
+      connId: "conn-target",
+      connect: { device: { id: "device-1" }, role: "primary" },
+      socket: { close: vi.fn() },
+    };
+    const unrelated = {
+      connId: "conn-unrelated",
+      connect: { device: { id: "device-2" }, role: "primary" },
+      socket: { close: vi.fn() },
+    };
+    const clients = new Set([target, unrelated] as never);
+
+    const context = createGatewayRequestContext(makeContextParams({ clients }));
+    context.invalidateClientsForDevice?.("device-1", { reason: "device-token-rotated" });
+
+    expect((target as { invalidated?: boolean }).invalidated).toBe(true);
+    expect((target as { invalidatedReason?: string }).invalidatedReason).toBe(
+      "device-token-rotated",
+    );
+    expect(target.socket.close).not.toHaveBeenCalled();
+
+    expect((unrelated as { invalidated?: boolean }).invalidated).toBeUndefined();
+    expect(unrelated.socket.close).not.toHaveBeenCalled();
+  });
+
+  it("disconnectClientsForDevice also marks the invalidated flag before closing", () => {
+    const target = {
+      connId: "conn-target",
+      connect: { device: { id: "device-1" }, role: "primary" },
+      socket: { close: vi.fn() },
+    };
+    const clients = new Set([target] as never);
+
+    const context = createGatewayRequestContext(makeContextParams({ clients }));
+    context.disconnectClientsForDevice?.("device-1");
+
+    expect((target as { invalidated?: boolean }).invalidated).toBe(true);
+    expect((target as { invalidatedReason?: string }).invalidatedReason).toBe("device-removed");
+    expect(target.socket.close).toHaveBeenCalledWith(4001, "device removed");
+  });
+
+  it("invalidateClientsForDevice filters by role when provided", () => {
+    const primary = {
+      connId: "conn-primary",
+      connect: { device: { id: "device-1" }, role: "primary" },
+      socket: { close: vi.fn() },
+    };
+    const secondary = {
+      connId: "conn-secondary",
+      connect: { device: { id: "device-1" }, role: "secondary" },
+      socket: { close: vi.fn() },
+    };
+    const clients = new Set([primary, secondary] as never);
+
+    const context = createGatewayRequestContext(makeContextParams({ clients }));
+    context.invalidateClientsForDevice?.("device-1", { role: "primary" });
+
+    expect((primary as { invalidated?: boolean }).invalidated).toBe(true);
+    expect((secondary as { invalidated?: boolean }).invalidated).toBeUndefined();
+  });
 });
diff --git a/src/gateway/server-request-context.ts b/src/gateway/server-request-context.ts
@@ -6,6 +6,8 @@ import { disconnectAllSharedGatewayAuthClients } from "./server-shared-auth-gene
 type GatewayRequestContextClient = GatewayClient & {
   socket: { close: (code: number, reason: string) => void };
   usesSharedGatewayAuth?: boolean;
+  invalidated?: boolean;
+  invalidatedReason?: string;
 };
 
 export type GatewayRequestContextParams = {
@@ -103,6 +105,19 @@ export function createGatewayRequestContext(
       }
       return false;
     },
+    invalidateClientsForDevice: (deviceId: string, opts?: { role?: string; reason?: string }) => {
+      const reason = opts?.reason ?? "device-invalidated";
+      for (const gatewayClient of params.clients) {
+        if (gatewayClient.connect.device?.id !== deviceId) {
+          continue;
+        }
+        if (opts?.role && gatewayClient.connect.role !== opts.role) {
+          continue;
+        }
+        gatewayClient.invalidated = true;
+        gatewayClient.invalidatedReason = reason;
+      }
+    },
     disconnectClientsForDevice: (deviceId: string, opts?: { role?: string }) => {
       for (const gatewayClient of params.clients) {
         if (gatewayClient.connect.device?.id !== deviceId) {
@@ -111,6 +126,11 @@ export function createGatewayRequestContext(
         if (opts?.role && gatewayClient.connect.role !== opts.role) {
           continue;
         }
+        // Mark before closing so any RPCs already pipelined in the WS buffer
+        // are rejected at the per-request dispatch check, regardless of
+        // whether socket.close() takes effect synchronously.
+        gatewayClient.invalidated = true;
+        gatewayClient.invalidatedReason ??= "device-removed";
         try {
           gatewayClient.socket.close(4001, "device removed");
         } catch {
diff --git a/src/gateway/server/ws-connection/message-handler.ts b/src/gateway/server/ws-connection/message-handler.ts
@@ -1441,6 +1441,15 @@ export function attachGatewayWsMessageHandler(params: {
       }
       const req = parsed;
       logWs("in", "req", { connId, id: req.id, method: req.method });
+      if (client.invalidated) {
+        const reason = client.invalidatedReason ?? "invalidated";
+        setCloseCause("client-invalidated", {
+          reason,
+          method: req.method,
+        });
+        close(4001, `client invalidated: ${reason}`);
+        return;
+      }
       if (client.usesSharedGatewayAuth) {
         const requiredSharedGatewaySessionGeneration =
           getRequiredSharedGatewaySessionGeneration?.();
diff --git a/src/gateway/server/ws-types.ts b/src/gateway/server/ws-types.ts
@@ -13,4 +13,6 @@ export type GatewayWsClient = {
   canvasHostUrl?: string;
   canvasCapability?: string;
   canvasCapabilityExpiresAtMs?: number;
+  invalidated?: boolean;
+  invalidatedReason?: string;
 };
