fix: detect KeepKey hot-swap + drop stale cache on device change (#49)

BitHighlander · claude · web-flow · commit c08ec7bfab72 · 2026-04-24T16:22:30.000-05:00
* fix: detect KeepKey hot-swap + drop stale cache on device change

The extension treated any reachable vault as a reachable device, with no
validation that the device currently paired with the vault is the same
one whose pubkeys we cached. That made three failure modes silent:

  1. Hot-swap on same vault — checkKeepKey kept state.deviceConnected=true
     and never re-probed, so addresses/balances pinned to the old device
     indefinitely.
  2. Extension reload with a new device — init() loaded the old pubkey
     cache without comparing deviceId against the freshly-probed one, so
     cached addresses survived even a full BEX reload.
  3. Vault auth rejection (stale apiKey after re-pair) — fetchPubkeys
     caught the error and silently fell back to cached pubkeys from the
     previous device, masking the re-pair requirement.

Four narrow changes:

- wallet.init() now compares cached deviceInfo.deviceId to the freshly-
  probed deviceId when a device is reachable. Mismatch → clear storage
  before fetchPubkeys runs. (Covers case 2.)

- wallet.handleDeviceSwitch() primitive clears in-memory pubkeys/paths,
  resets the deviceConnected flag, and wipes pubkeyStorage so the next
  fetch rebuilds from the new device. Keeps SDK alive — only the
  device-specific state is dropped.

- background/index.ts checkKeepKey now runs a periodic (30s, throttled)
  getFeatures re-probe while deviceConnected. Compares before/after
  deviceId; on mismatch calls handleDeviceSwitch which clears balance
  cache, per-chain address caches (Solana/Tron/TON), and triggers a
  fresh pubkey+balance refresh. (Covers case 1.)

- fetchPubkeys no longer silently falls back to cache on auth errors.
  Detects 401 / "unauthorized" / "not paired" shapes, clears the stored
  apiKey, and throws a user-facing message for the sidebar to surface.
  Non-auth failures (network blip, device busy) still fall back to cache
  as view-only — that path is correct. (Covers case 3.)

No user-visible UX change on the happy path. Recovery from device swap
is now automatic within ~30s on steady state, immediate on extension
reload.

Co-Authored-By: Claude Opus 4.7 (1M context) &lt;noreply@anthropic.com&gt;

* fix: rebuild default paths on device switch so refreshPubkeys has a batch

handleDeviceSwitch cleared state.paths to [], but the caller
(background checkKeepKey) then calls wallet.refreshPubkeys() which
only probes + fetches — it does not repopulate paths. fetchPubkeys
maps state.paths into the batch sent to the device, so an empty
paths array meant an empty batch, meaning the hot-swap recovery
finished with zero pubkeys and zero balances instead of the new
device's view.

Reset to getDefaultPaths() on switch so refreshPubkeys has something
to query against.

Co-Authored-By: Claude Opus 4.7 (1M context) &lt;noreply@anthropic.com&gt;

* fix: re-prefetch Solana/Tron/TON pubkeys on device switch

refreshPubkeys only re-runs the default-path batch via
wallet.getPublicKeys — that batch covers BTC / LTC / ETH / Cosmos /
etc., but *not* SOL / TRX / TON. Those three are derived dynamically
at onStart via prefetchSolanaPubkey / prefetchTronPubkey /
prefetchTonAddress (each calls its chain's `*GetAddress` method and
stashes the result in a per-chain cache + adds a pubkey to state).

After a hot-swap:
  - handleDeviceSwitch clears the per-chain caches (resetXState)
  - refreshPubkeys repopulates the default batch
  - but nothing re-runs the three prefetches
  → SOL/TRX/TON vanish from the network dropdown / asset list / balance
    fetch until the user reloads the extension or manually navigates to
    those chains, which racy-triggers the lazy getAddress in each handler.

Fan out all three prefetches in parallel via Promise.allSettled after
refreshPubkeys. Non-throwing by design, so if one chain's derivation
fails (minFirmware mismatch, etc.) the other two still recover.

Co-Authored-By: Claude Opus 4.7 (1M context) &lt;noreply@anthropic.com&gt;

---------

Co-authored-by: Claude Opus 4.7 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/chrome-extension/src/background/index.ts b/chrome-extension/src/background/index.ts
@@ -110,6 +110,66 @@ function pushStateChangeEvent() {
 let lastDeviceProbeAt = 0;
 const DEVICE_PROBE_INTERVAL_MS = 15_000;
 
+// Separate, longer throttle for the "already-connected, verify deviceId hasn't
+// changed" re-probe. Vault + device hot-swap is rare enough that 30s latency
+// on detection is fine; keeping this longer than the view-only probe avoids
+// doubling the getFeatures traffic on the steady-state path.
+let lastDeviceVerifyAt = 0;
+const DEVICE_VERIFY_INTERVAL_MS = 30_000;
+
+/**
+ * Called when we detect that the vault is now paired with a different
+ * KeepKey than the one we had cached. Clears every piece of state keyed
+ * to the previous device and re-fetches from the new one.
+ */
+async function handleDeviceSwitch(newDeviceInfo: any) {
+  const tag = TAG + ' | handleDeviceSwitch | ';
+  console.warn(tag, 'Device swap detected. Purging caches and re-fetching.');
+
+  // In-memory + storage cache owned by wallet.ts
+  await wallet.handleDeviceSwitch(newDeviceInfo);
+
+  // Balance caches — pubkey-keyed, so they're poisoned by the old device
+  cachedBalances = [];
+  balancesFetchInProgress = null;
+
+  // Per-chain address caches (Solana/Tron/TON each keep their own lookup
+  // cache above the pubkey layer)
+  resetSolanaState();
+  resetTronState();
+  resetTonState();
+
+  // Re-fetch against the new device. refreshPubkeys re-probes and pulls
+  // a fresh pubkey batch, then updates state.initialized.
+  try {
+    await wallet.refreshPubkeys();
+
+    // refreshPubkeys only hits getDefaultPaths() — the big batched
+    // derivation. Solana, Tron, and TON addresses are *dynamically*
+    // added at onStart via these prefetches (SOL needs solanaGetAddress,
+    // TRX needs tronGetAddress, TON needs tonGetAddress — none of which
+    // are in the xpub.getPublicKeys batch). Without re-running them on
+    // a device swap, those three chains end up with stale per-chain
+    // caches (zeroed out by resetXState above) and no pubkeys, so
+    // they'd vanish from the network dropdown until a user manually
+    // visited the asset or reloaded the extension.
+    //
+    // Fire in parallel — each is non-throwing, so an individual chain
+    // failure won't take the others down.
+    await Promise.allSettled([prefetchSolanaPubkey(), prefetchTronPubkey(), prefetchTonAddress()]);
+
+    pushStateChangeEvent();
+    pushBalancesUpdated();
+    // Kick a fresh balance fetch in the background so the dashboard
+    // swaps to the new device's balances without waiting for the next
+    // user-triggered refresh.
+    fetchBalancesFromPioneer(true).catch(e => console.warn(tag, 'Post-switch balance fetch failed:', e));
+  } catch (e) {
+    console.error(tag, 'Failed to re-fetch pubkeys after device switch:', (e as Error)?.message || e);
+    pushStateChangeEvent();
+  }
+}
+
 async function checkKeepKey() {
   const prevState = KEEPKEY_STATE;
   try {
@@ -139,6 +199,28 @@ async function checkKeepKey() {
         // First-run case: init failed earlier (no device, no cache) — retry.
         lastDeviceProbeAt = now;
         onStart();
+      } else if (wallet.isInitialized() && wallet.isDeviceConnected()) {
+        // Steady state: vault-up, device-connected. Periodically re-probe
+        // features to verify the same physical device is still paired. A
+        // hot-swap to a different KeepKey will look identical to the /docs
+        // endpoint but returns a different device_id from getFeatures.
+        const mayVerify = now - lastDeviceVerifyAt >= DEVICE_VERIFY_INTERVAL_MS;
+        if (mayVerify) {
+          lastDeviceVerifyAt = now;
+          const beforeId = wallet.getDeviceId();
+          wallet
+            .probeDevice()
+            .then(ok => {
+              if (!ok) return; // probe failure — next tick will set state=4
+              const afterId = wallet.getDeviceId();
+              if (beforeId && afterId && beforeId !== afterId) {
+                handleDeviceSwitch(wallet.getDeviceInfo()).catch(e =>
+                  console.error(TAG, 'handleDeviceSwitch failed:', e),
+                );
+              }
+            })
+            .catch(e => console.warn(TAG, 'Feature re-probe failed:', (e as Error)?.message || e));
+        }
       }
     }
   } catch (error: any) {
diff --git a/chrome-extension/src/background/wallet.ts b/chrome-extension/src/background/wallet.ts
@@ -98,6 +98,22 @@ export async function init(): Promise<WalletState> {
       if (!state.deviceInfo) {
         state.deviceInfo = { label: 'KeepKey', model: 'KeepKey', deviceId: 'unknown' };
       }
+    } else {
+      // Device is reachable — validate that any cached pubkeys match this
+      // specific device. Service worker restarts and hot-swaps can both
+      // leave the pubkey cache pointing at a previously-connected device;
+      // reusing it would surface the wrong addresses and stale balances.
+      try {
+        const cached = await pubkeyStorage.loadPubkeys();
+        const cachedId = cached?.deviceInfo?.deviceId;
+        const currentId = state.deviceInfo?.deviceId;
+        if (cachedId && currentId && cachedId !== 'unknown' && cachedId !== currentId) {
+          console.warn(tag, `Device changed: cache=${cachedId} → probed=${currentId}. Invalidating cache.`);
+          await pubkeyStorage.clearPubkeys();
+        }
+      } catch (e) {
+        console.warn(tag, 'Cache validation failed:', (e as Error)?.message || e);
+      }
     }
 
     // Set up paths
@@ -205,8 +221,28 @@ async function fetchPubkeys(): Promise<void> {
     }
   } catch (e) {
     console.error(tag, 'Error fetching pubkeys from device:', e);
-    // Device call failed — downgrade to view-only and use cache if available
     state.deviceConnected = false;
+
+    // Auth errors mean the vault's pairing with this extension is broken
+    // (e.g. user re-paired with a different device, or vault rotated the
+    // key). Silently falling back to cached pubkeys from the *previous*
+    // device would mask the re-pair requirement — the user would see
+    // stale addresses with no explanation. Throw so the sidebar can
+    // prompt re-pair and the cached pubkeys are only used when we've
+    // chosen view-only mode, not when auth is rejecting us.
+    if (isAuthError(e)) {
+      console.warn(tag, 'Auth error from vault — clearing stale API key and failing fast');
+      try {
+        await keepKeyApiKeyStorage.saveApiKey('');
+      } catch {
+        /* ignore */
+      }
+      throw new Error('Vault auth failed. Your KeepKey needs to be re-paired with the extension.');
+    }
+
+    // Non-auth failure (network blip, device busy, etc.) — view-only via
+    // cache is acceptable. We keep deviceConnected=false so signing paths
+    // will re-probe.
     if (cachedPubkeys.length > 0) {
       console.warn(tag, 'Falling back to', cachedPubkeys.length, 'cached pubkeys (view-only)');
       state.pubkeys = cachedPubkeys;
@@ -216,6 +252,26 @@ async function fetchPubkeys(): Promise<void> {
   }
 }
 
+/**
+ * Heuristic for "the vault rejected our credentials." Covers the
+ * error shapes we've seen from the vault-sdk client: HTTP 401 wrapped
+ * in a fetch error, plus bare messages that include those substrings.
+ * Conservative — misclassifying a non-auth failure as auth would just
+ * wipe the API key and force re-pairing, which is recoverable; the
+ * opposite direction (missing a real auth error and silently serving
+ * stale pubkeys) is the bug we're fixing.
+ */
+function isAuthError(e: unknown): boolean {
+  const msg = ((e as Error)?.message || String(e)).toLowerCase();
+  return (
+    msg.includes('401') ||
+    msg.includes('unauthorized') ||
+    msg.includes('auth') ||
+    msg.includes('invalid api key') ||
+    msg.includes('not paired')
+  );
+}
+
 /**
  * Get the SDK instance, re-initializing if needed.
  */
@@ -366,6 +422,46 @@ export function getDeviceInfo() {
   return state.deviceInfo;
 }
 
+/**
+ * Current device id, or null if unknown. Used by background polling to
+ * detect device hot-swaps without duplicating the probe logic.
+ */
+export function getDeviceId(): string | null {
+  const id = state.deviceInfo?.deviceId;
+  return id && id !== 'unknown' ? id : null;
+}
+
+/**
+ * Tear down everything keyed to the previously-connected device so the
+ * next `refreshPubkeys()` / `fetchPubkeys()` starts clean:
+ *   - in-memory pubkeys (cleared)
+ *   - paths (rebuilt to defaults — refreshPubkeys does NOT repopulate
+ *     them, so leaving paths empty would silently send an empty batch
+ *     to the device and return zero pubkeys)
+ *   - persisted pubkey cache (wiped)
+ *   - `deviceConnected` flag (forces re-probe on next call)
+ *
+ * We intentionally keep `state.sdk` alive — the vault REST client can
+ * happily serve the new device once pubkeys are re-fetched — and we
+ * keep `state.deviceInfo` populated with the freshly-probed device's
+ * info (passed in by the caller) so UI can show the new label
+ * immediately without waiting for the re-fetch to finish.
+ */
+export async function handleDeviceSwitch(newDeviceInfo: WalletState['deviceInfo']): Promise<void> {
+  const tag = TAG + ' | handleDeviceSwitch | ';
+  console.warn(tag, 'Device switch detected — clearing caches');
+  state.pubkeys = [];
+  state.paths = getDefaultPaths();
+  state.deviceInfo = newDeviceInfo;
+  state.deviceConnected = false;
+  state.initialized = false;
+  try {
+    await pubkeyStorage.clearPubkeys();
+  } catch (e) {
+    console.warn(tag, 'Failed to clear pubkey storage:', e);
+  }
+}
+
 /**
  * Get the full wallet state (for backward compat with APP references).
  */
