fix: restore bootstrap tokens after send failure (#60221)

src/gateway/server/ws-connection/message-handler.ts

@@ -5,6 +5,7 @@ import { loadConfig } from "../../../config/config.js";
 import {
   getDeviceBootstrapTokenProfile,
   redeemDeviceBootstrapTokenProfile,
+  restoreDeviceBootstrapToken,
   revokeDeviceBootstrapToken,
   verifyDeviceBootstrapToken,
 } from "../../../infra/device-bootstrap.js";
@@ -215,6 +216,17 @@ export function attachGatewayWsMessageHandler(params: {
     logWsControl,
   } = params;
 
+  const sendFrame = async (obj: unknown): Promise<void> =>
+    await new Promise<void>((resolve, reject) => {
+      socket.send(JSON.stringify(obj), (err) => {
+        if (err) {
+          reject(err);
+          return;
+        }
+        resolve();
+      });
+    });
+
   const configSnapshot = loadConfig();
   const trustedProxies = configSnapshot.gateway?.trustedProxies ?? [];
   const allowRealIpFallback = configSnapshot.gateway?.allowRealIpFallback === true;
@@ -1150,14 +1162,9 @@ export function attachGatewayWsMessageHandler(params: {
             );
         }
 
-        logWs("out", "hello-ok", {
-          connId,
-          methods: gatewayMethods.length,
-          events: events.length,
-          presence: snapshot.presence.length,
-          stateVersion: snapshot.stateVersion.presence,
-        });
-
+        let consumedBootstrapTokenRecord:
+          | Awaited<ReturnType<typeof revokeDeviceBootstrapToken>>["record"]
+          | undefined;
         if (
           authMethod === "bootstrap-token" &&
           bootstrapProfile &&
@@ -1174,6 +1181,7 @@ export function attachGatewayWsMessageHandler(params: {
               const revoked = await revokeDeviceBootstrapToken({
                 token: bootstrapTokenCandidate,
               });
+              consumedBootstrapTokenRecord = revoked.record;
               if (!revoked.removed) {
                 logGateway.warn(
                   `bootstrap token revoke skipped after profile redemption device=${device.id}`,
@@ -1186,7 +1194,31 @@ export function attachGatewayWsMessageHandler(params: {
             );
           }
         }
-        send({ type: "res", id: frame.id, ok: true, payload: helloOk });
+        try {
+          await sendFrame({ type: "res", id: frame.id, ok: true, payload: helloOk });
+        } catch (err) {
+          if (consumedBootstrapTokenRecord) {
+            try {
+              await restoreDeviceBootstrapToken({
+                record: consumedBootstrapTokenRecord,
+              });
+            } catch (restoreErr) {
+              logGateway.warn(
+                `bootstrap token restore failed after hello send error device=${device?.id ?? "unknown"}: ${formatForLog(restoreErr)}`,
+              );
+            }
+          }
+          setCloseCause("hello-send-failed", { error: formatForLog(err) });
+          close();
+          return;
+        }
+        logWs("out", "hello-ok", {
+          connId,
+          methods: gatewayMethods.length,
+          events: events.length,
+          presence: snapshot.presence.length,
+          stateVersion: snapshot.stateVersion.presence,
+        });
         void refreshGatewayHealthSnapshot({ probe: true }).catch((err) =>
           logHealth.error(`post-connect health refresh failed: ${formatError(err)}`),
         );

src/infra/device-bootstrap.test.ts

@@ -8,6 +8,7 @@ import {
   getDeviceBootstrapTokenProfile,
   issueDeviceBootstrapToken,
   redeemDeviceBootstrapTokenProfile,
+  restoreDeviceBootstrapToken,
   revokeDeviceBootstrapToken,
   verifyDeviceBootstrapToken,
 } from "./device-bootstrap.js";
@@ -163,12 +164,30 @@ describe("device bootstrap tokens", () => {
     });
   });
 
+  it("restores a revoked bootstrap token record after send failure recovery", async () => {
+    const baseDir = await createTempDir();
+    const issued = await issueDeviceBootstrapToken({ baseDir });
+
+    await expect(verifyBootstrapToken(baseDir, issued.token)).resolves.toEqual({ ok: true });
+    const revoked = await revokeDeviceBootstrapToken({ baseDir, token: issued.token });
+    expect(revoked.removed).toBe(true);
+    expect(revoked.record?.token).toBe(issued.token);
+
+    if (!revoked.record) {
+      throw new Error("expected revoked bootstrap token record");
+    }
+    await restoreDeviceBootstrapToken({ baseDir, record: revoked.record });
+    await expect(verifyBootstrapToken(baseDir, issued.token)).resolves.toEqual({ ok: true });
+  });
+
   it("revokes a specific bootstrap token", async () => {
     const baseDir = await createTempDir();
     const first = await issueDeviceBootstrapToken({ baseDir });
     const second = await issueDeviceBootstrapToken({ baseDir });
 
-    await expect(revokeDeviceBootstrapToken({ baseDir, token: first.token })).resolves.toEqual({
+    await expect(
+      revokeDeviceBootstrapToken({ baseDir, token: first.token }),
+    ).resolves.toMatchObject({
       removed: true,
     });
 

src/infra/device-bootstrap.ts

@@ -192,7 +192,7 @@ export async function clearDeviceBootstrapTokens(
 export async function revokeDeviceBootstrapToken(params: {
   token: string;
   baseDir?: string;
-}): Promise<{ removed: boolean }> {
+}): Promise<{ removed: boolean; record?: DeviceBootstrapTokenRecord }> {
   return await withLock(async () => {
     const providedToken = params.token.trim();
     if (!providedToken) {
@@ -205,9 +205,21 @@ export async function revokeDeviceBootstrapToken(params: {
     if (!found) {
       return { removed: false };
     }
-    delete state[found[0]];
+    const [tokenKey, record] = found;
+    delete state[tokenKey];
+    await persistState(state, params.baseDir);
+    return { removed: true, record };
+  });
+}
+
+export async function restoreDeviceBootstrapToken(params: {
+  record: DeviceBootstrapTokenRecord;
+  baseDir?: string;
+}): Promise<void> {
+  return await withLock(async () => {
+    const state = await loadState(params.baseDir);
+    state[params.record.token] = params.record;
     await persistState(state, params.baseDir);
-    return { removed: true };
   });
 }