From 39361d13be004584d4b19aa656256b7d74f086ca Mon Sep 17 00:00:00 2001 From: Ayaan Zaidi Date: Fri, 3 Apr 2026 17:57:48 +0530 Subject: [PATCH] fix: restore bootstrap tokens after send failure (#60221) --- .../server/ws-connection/message-handler.ts | 50 +++++++++++++++---- src/infra/device-bootstrap.test.ts | 21 +++++++- src/infra/device-bootstrap.ts | 18 +++++-- 3 files changed, 76 insertions(+), 13 deletions(-) diff --git a/src/gateway/server/ws-connection/message-handler.ts b/src/gateway/server/ws-connection/message-handler.ts index bb91c24cb6b..c2b334f1b95 100644 --- a/src/gateway/server/ws-connection/message-handler.ts +++ b/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 => + await new Promise((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>["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)}`), ); diff --git a/src/infra/device-bootstrap.test.ts b/src/infra/device-bootstrap.test.ts index 98898fd9ac4..ca77ff84fa3 100644 --- a/src/infra/device-bootstrap.test.ts +++ b/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, }); diff --git a/src/infra/device-bootstrap.ts b/src/infra/device-bootstrap.ts index 5a362aa2ef6..c64c91f1092 100644 --- a/src/infra/device-bootstrap.ts +++ b/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 { + return await withLock(async () => { + const state = await loadState(params.baseDir); + state[params.record.token] = params.record; await persistState(state, params.baseDir); - return { removed: true }; }); }