fix: detect KeepKey hot-swap + drop stale cache on device change

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) {

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).
  */