From 2756033bf97be6185336c148ce3e39cc3572e911 Mon Sep 17 00:00:00 2001
From: adlee-was-taken <aaron.daniel.lee@gmail.com>
Date: Mon, 27 Apr 2026 18:34:35 -0400
Subject: [PATCH] feat(ext/setup): unified device registration in Step 5; fixes
 silent dropped pubkey

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 extension/src/setup/setup.ts | 110 +++++++++++++++++++++--------------
 1 file changed, 67 insertions(+), 43 deletions(-)

diff --git a/extension/src/setup/setup.ts b/extension/src/setup/setup.ts
index 8c0fa6c..e250189 100644
--- a/extension/src/setup/setup.ts
+++ b/extension/src/setup/setup.ts
@@ -7,6 +7,7 @@
 /// Step 5: Finish (download reference image, push config to extension or copy JSON)
 
 import { createGitHost, uint8ArrayToBase64 } from '../service-worker/git-host';
+import { addDevice } from '../service-worker/devices';
 import { probeVault } from './probe';
 import type { VaultConfig } from '../shared/types';
 import type { SessionHandle } from 'relicario-wasm';
@@ -1042,32 +1043,38 @@ function renderStep5(): string {
     repoPath: state.repoPath,
     apiToken: state.apiToken,
   };
-
   const configJson = JSON.stringify(config, null, 2);
+  const isAttach = state.mode === 'attach';
 
   return `
     <div class="wizard-step">
       <div class="success-box">
-        <h3>vault created</h3>
-        <p class="secondary">Your vault has been initialized and pushed to the repository.</p>
+        <h3>${isAttach ? 'device verified' : 'vault created'}</h3>
+        <p class="secondary">
+          ${isAttach
+            ? 'Your passphrase and reference image decrypt the vault successfully.'
+            : 'Your vault has been initialized and pushed to the repository.'}
+        </p>
       </div>
 
-      <div class="form-group">
-        <label class="label">reference image</label>
-        <p class="muted" style="margin-bottom:8px;">
-          Download and store this image securely. It is your second factor for decryption.
-          Without it, you cannot unlock the vault.
-        </p>
-        <button class="btn btn-primary" id="download-ref-btn">download reference.jpg</button>
-      </div>
+      ${isAttach ? '' : `
+        <div class="form-group">
+          <label class="label">reference image</label>
+          <p class="muted" style="margin-bottom:8px;">
+            Download and store this image securely. It is your second factor for decryption.
+            Without it, you cannot unlock the vault.
+          </p>
+          <button class="btn btn-primary" id="download-ref-btn">download reference.jpg</button>
+        </div>
+      `}
 
       ${state.extensionDetected ? `
         <div class="form-group" style="margin-top:16px;">
-          <label class="label">extension configuration</label>
+          <label class="label">register this device</label>
           <button class="btn btn-primary" id="push-config-btn" ${state.configPushed ? 'disabled' : ''}>
-            ${state.configPushed ? 'config saved to extension' : 'save config to extension'}
+            ${state.configPushed ? 'device registered' : (isAttach ? 'attach this device' : 'register this device')}
           </button>
-          ${state.configPushed ? '<span class="test-result pass" style="margin-left:8px;">saved</span>' : ''}
+          ${state.configPushed ? '<span class="test-result pass" style="margin-left:8px;">done</span>' : ''}
         </div>
       ` : `
         <div class="form-group" style="margin-top:16px;">
@@ -1098,7 +1105,8 @@ function attachStep5(): void {
   });
 
   document.getElementById('push-config-btn')?.addEventListener('click', async () => {
-    if (!state.referenceImageBytes) return;
+    state.error = null;
+    render();
 
     const config: VaultConfig = {
       hostType: state.hostType,
@@ -1107,39 +1115,55 @@ function attachStep5(): void {
       apiToken: state.apiToken,
     };
 
-    const imageBase64 = uint8ArrayToBase64(state.referenceImageBytes);
-
     try {
-      chrome.runtime.sendMessage(
-        { type: 'save_setup', config, imageBase64 },
-        async (response: { ok: boolean; error?: string }) => {
-          if (response?.ok) {
-            state.configPushed = true;
+      const w = await loadWasm();
+      const keypair = JSON.parse(w.generate_device_keypair()) as {
+        public_key_hex: string; private_key_base64: string;
+      };
 
-            // Generate device keypair and register
-            const w = await loadWasm();
-            const keypair = JSON.parse(w.generate_device_keypair()) as { public_key_hex: string; private_key_base64: string };
+      // 1) Save private key + name locally.
+      await chrome.storage.local.set({
+        device_name: state.deviceName,
+        device_private_key: keypair.private_key_base64,
+      });
 
-            // Store private key locally
-            await chrome.storage.local.set({
-              device_name: state.deviceName,
-              device_private_key: keypair.private_key_base64,
-            });
+      // 2) Save vault config + reference image to extension storage.
+      const imageBytes = state.referenceImageBytes ?? state.referenceImageBytesAttach;
+      const imageBase64 = imageBytes ? uint8ArrayToBase64(imageBytes) : '';
+      const saveOk = await new Promise<boolean>((resolve) => {
+        chrome.runtime.sendMessage(
+          { type: 'save_setup', config, imageBase64 },
+          (response: { ok: boolean; error?: string }) => {
+            if (!response?.ok) {
+              state.error = response?.error ?? 'Failed to save config to extension';
+              resolve(false); return;
+            }
+            resolve(true);
+          },
+        );
+      });
+      if (!saveOk) { render(); return; }
 
-            // Register device with vault
-            chrome.runtime.sendMessage({
-              type: 'add_device',
-              name: state.deviceName,
-              public_key: keypair.public_key_hex,
-            });
-          } else {
-            state.error = response?.error ?? 'Failed to save config to extension';
-          }
-          render();
-        },
-      );
+      // 3) Register device on the remote (read-modify-write devices.json).
+      const hostUrl = state.hostType === 'github' ? 'https://api.github.com' : state.hostUrl;
+      const host = createGitHost(state.hostType, hostUrl, state.repoPath, state.apiToken);
+      await addDevice(host, {
+        name: state.deviceName,
+        public_key: keypair.public_key_hex,
+        added_at: Math.floor(Date.now() / 1000),
+      });
+
+      // 4) Release any attach-mode WASM handle.
+      if (state.verifiedHandle !== null) {
+        try { w.lock(state.verifiedHandle); } catch { /* best effort */ }
+        state.verifiedHandle = null;
+      }
+
+      state.configPushed = true;
+      render();
     } catch (err: unknown) {
-      state.error = `Failed to communicate with extension: ${err instanceof Error ? err.message : String(err)}`;
+      console.error('[relicario setup] register device failed:', err);
+      state.error = `Failed to register device: ${err instanceof Error ? err.message : String(err)}`;
       render();
     }
   });