relicario/extension/src/service-worker/__tests__/vault.test.ts

import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import * as vault from '../vault';
import * as session from '../session';
import type { PopupState } from '../router/popup-only';
import type { GitHost } from '../git-host';
import * as gitHostMod from '../git-host';

// --- Mock git-host module ---
// createGitHost is called internally by handleCreateVault / handleAttachVault;
// we need to intercept it and return a fake GitHost. uint8ArrayToBase64 must
// still work — vault.ts calls it for the imageBase64 storage value.

// Shared factory used both inside vi.mock and in beforeEach re-wire.
function makeHostMock(): GitHost & { _calls: Record<string, unknown[][]> } {
  const calls: Record<string, unknown[][]> = {
    writeFileCreateOnly: [],
    writeFile: [],
    readFile: [],
  };
  return {
    _calls: calls,
    readFile: vi.fn().mockImplementation(async (path: string) => {
      // Serve the vault-meta files needed by fetchVaultMeta + attach flow.
      if (path === '.relicario/salt') return new Uint8Array(32);
      if (path === '.relicario/params.json') {
        return new TextEncoder().encode('{"argon2_m":65536,"argon2_t":3,"argon2_p":4}');
      }
      if (path === 'manifest.enc') return new Uint8Array([0xab, 0xcd]);
      // .relicario/devices.json throws so readDevices falls back to [].
      throw new Error(`404: ${path}`);
    }),
    writeFile: vi.fn().mockImplementation(async (...args: unknown[]) => {
      calls.writeFile.push(args);
    }),
    writeFileCreateOnly: vi.fn().mockImplementation(async (...args: unknown[]) => {
      calls.writeFileCreateOnly.push(args);
    }),
    deleteFile: vi.fn(),
    listDir: vi.fn().mockResolvedValue([]),
    lastCommit: vi.fn().mockResolvedValue(null),
    putBlob: vi.fn(),
    getBlob: vi.fn(),
    deleteBlob: vi.fn(),
  };
}

vi.mock('../git-host', async () => {
  const actual = await vi.importActual<typeof import('../git-host')>('../git-host');

  // Expose a handle so tests can grab the last-created fake host.
  (globalThis as { __lastFakeGitHost?: ReturnType<typeof makeHostMock> | null }).__lastFakeGitHost = null;

  return {
    ...actual,
    createGitHost: vi.fn().mockImplementation(() => {
      const h = makeHostMock();
      (globalThis as { __lastFakeGitHost?: ReturnType<typeof makeHostMock> | null }).__lastFakeGitHost = h;
      return h;
    }),
  };
});

// --- Chrome storage mock ---

function mockChromeStorage(initial: Record<string, unknown> = {}) {
  const store: Record<string, unknown> = { ...initial };
  (global as { chrome: unknown }).chrome = {
    storage: {
      local: {
        get: vi.fn((keys: string | string[]) => {
          const arr = Array.isArray(keys) ? keys : [keys];
          const out: Record<string, unknown> = {};
          for (const k of arr) if (k in store) out[k] = store[k];
          return Promise.resolve(out);
        }),
        set: vi.fn((kv: Record<string, unknown>) => {
          Object.assign(store, kv);
          return Promise.resolve();
        }),
      },
    },
  } as never;
  return store;
}

// --- Helpers ---

function makeFakeHandle() {
  return { free: vi.fn() };
}

function makeWasm(overrides: Record<string, unknown> = {}) {
  const fakeHandle = makeFakeHandle();
  return {
    _handle: fakeHandle,
    embed_image_secret: vi.fn(() => new Uint8Array([1, 2, 3])),
    unlock: vi.fn(() => fakeHandle),
    manifest_encrypt: vi.fn(() => new Uint8Array([9])),
    manifest_decrypt: vi.fn(() => ({ schema_version: 2, items: {} })),
    default_vault_settings_json: vi.fn(() => '{}'),
    settings_encrypt: vi.fn(() => new Uint8Array([8])),
    register_device: vi.fn(() => ({ signing_public_key: 'pk', deploy_public_key: 'dk' })),
    lock: vi.fn(),
    ...overrides,
  };
}

function makeState(wasm: ReturnType<typeof makeWasm>): PopupState {
  return {
    manifest: null,
    gitHost: null,
    wasm,
  };
}

const BASE_MSG = {
  config: { hostType: 'gitea' as const, hostUrl: 'https://g', repoPath: 'u/v', apiToken: 't' },
  passphrase: 'pw',
  carrierImageBytes: new Uint8Array([0, 0, 0]).buffer,
  deviceName: 'Dev',
};

// --- Tests ---

describe('handleCreateVault', () => {
  let setCurrent: ReturnType<typeof vi.spyOn>;

  beforeEach(() => {
    mockChromeStorage();
    setCurrent = vi.spyOn(session, 'setCurrent').mockImplementation(() => {});
  });

  afterEach(() => {
    vi.restoreAllMocks();
  });

  it('Test 1 (happy path): returns ok:true with expected data and correct side effects', async () => {
    const wasm = makeWasm();
    const state = makeState(wasm);

    const resp = await vault.handleCreateVault(BASE_MSG, state);

    expect(resp.ok).toBe(true);
    if (!resp.ok) throw new Error('expected ok:true');

    // Response shape
    expect(resp.data.referenceImageBytes).toBeInstanceOf(Uint8Array);
    expect(resp.data.deviceName).toBe('Dev');
    expect(resp.data.recoveryQrAvailable).toBe(true);

    // Fake GitHost captures the four writeFileCreateOnly calls
    const fakeHost = (globalThis as { __lastFakeGitHost?: { writeFileCreateOnly: ReturnType<typeof vi.fn> } }).__lastFakeGitHost;
    expect(fakeHost).not.toBeNull();
    const wfco = fakeHost!.writeFileCreateOnly as ReturnType<typeof vi.fn>;
    const paths = wfco.mock.calls.map((c: unknown[]) => c[0]);
    expect(paths).toContain('.relicario/salt');
    expect(paths).toContain('.relicario/params.json');
    expect(paths).toContain('manifest.enc');
    expect(paths).toContain('settings.enc');

    // register_device called with the device name
    expect(wasm.register_device).toHaveBeenCalledWith('Dev');

    // chrome.storage.local.set called with vaultConfig + imageBase64 + device_name
    const chromeSets = (global as { chrome: { storage: { local: { set: ReturnType<typeof vi.fn> } } } })
      .chrome.storage.local.set.mock.calls;
    const merged: Record<string, unknown> = {};
    for (const [kv] of chromeSets) Object.assign(merged, kv);
    expect(merged).toHaveProperty('vaultConfig');
    expect(merged).toHaveProperty('imageBase64');
    expect(merged).toHaveProperty('device_name', 'Dev');

    // session.setCurrent was called (ownership transferred — handle NOT freed)
    expect(setCurrent).toHaveBeenCalled();
    expect(wasm._handle.free).not.toHaveBeenCalled();
  });

  it('Test 2 (failure path — early throw): ok:false, no writeFileCreateOnly calls', async () => {
    const wasm = makeWasm({
      embed_image_secret: vi.fn(() => { throw new Error('embed failed'); }),
    });
    const state = makeState(wasm);

    const resp = await vault.handleCreateVault(BASE_MSG, state);

    expect(resp.ok).toBe(false);
    if (resp.ok) throw new Error('expected ok:false');
    expect(resp.error).toBeTruthy();
    expect(resp.error.length).toBeGreaterThan(0);

    const fakeHost = (globalThis as { __lastFakeGitHost?: { writeFileCreateOnly: ReturnType<typeof vi.fn> } }).__lastFakeGitHost;
    // No GitHost was created at all (failed before createGitHost call), OR
    // if somehow created, no writeFileCreateOnly calls happened.
    if (fakeHost) {
      expect((fakeHost.writeFileCreateOnly as ReturnType<typeof vi.fn>).mock.calls).toHaveLength(0);
    }
  });

  it('Test 3 (handle cleanup on mid-flight failure): lock + free called, ok:false', async () => {
    const wasm = makeWasm({
      manifest_encrypt: vi.fn(() => { throw new Error('encrypt failed'); }),
    });
    const state = makeState(wasm);

    const resp = await vault.handleCreateVault(BASE_MSG, state);

    expect(resp.ok).toBe(false);

    // unlock succeeded (handle was acquired), manifest_encrypt failed after that.
    // Finally block must: lock(handle) then handle.free().
    expect(wasm.lock).toHaveBeenCalledWith(wasm._handle);
    expect(wasm._handle.free).toHaveBeenCalled();

    // Ownership was NOT transferred — setCurrent must NOT have been called.
    expect(setCurrent).not.toHaveBeenCalled();
  });
});

// --- attach_vault ---

const ATTACH_MSG = {
  config: { hostType: 'gitea' as const, hostUrl: 'https://g', repoPath: 'u/v', apiToken: 't' },
  passphrase: 'pw',
  referenceImageBytes: new Uint8Array([1, 2, 3]).buffer,
  deviceName: 'Dev2',
};

describe('handleAttachVault', () => {
  let setCurrent: ReturnType<typeof vi.spyOn>;

  beforeEach(() => {
    mockChromeStorage();
    // Re-wire createGitHost: vi.restoreAllMocks() in the create-vault afterEach
    // strips the mockImplementation from the vi.fn(), leaving it returning undefined.
    // We re-establish it here so each attach test starts with a fresh fake host.
    vi.mocked(gitHostMod.createGitHost).mockImplementation(() => {
      const h = makeHostMock();
      (globalThis as { __lastFakeGitHost?: ReturnType<typeof makeHostMock> | null }).__lastFakeGitHost = h;
      return h;
    });
    setCurrent = vi.spyOn(session, 'setCurrent').mockImplementation(() => {});
  });

  afterEach(() => {
    vi.restoreAllMocks();
  });

  it('Test 1 (happy path): returns ok:true, state populated, handle ownership transferred', async () => {
    const wasm = makeWasm();
    const state = makeState(wasm);

    const resp = await vault.handleAttachVault(ATTACH_MSG, state);

    expect(resp.ok).toBe(true);
    if (!resp.ok) throw new Error('expected ok:true');
    expect(resp.data.deviceName).toBe('Dev2');

    // WASM calls in order: unlock → manifest_decrypt (verification) → register_device
    expect(wasm.unlock).toHaveBeenCalled();
    expect(wasm.manifest_decrypt).toHaveBeenCalled();
    expect(wasm.register_device).toHaveBeenCalledWith('Dev2');

    // chrome.storage.local.set received vaultConfig + imageBase64 + device_name
    const chromeSets = (global as { chrome: { storage: { local: { set: ReturnType<typeof vi.fn> } } } })
      .chrome.storage.local.set.mock.calls;
    const merged: Record<string, unknown> = {};
    for (const [kv] of chromeSets) Object.assign(merged, kv);
    expect(merged).toHaveProperty('vaultConfig');
    expect(merged).toHaveProperty('imageBase64');
    expect(merged).toHaveProperty('device_name', 'Dev2');

    // session.setCurrent called — ownership transferred; handle NOT freed
    expect(setCurrent).toHaveBeenCalled();
    expect(wasm._handle.free).not.toHaveBeenCalled();

    // State wired up
    expect(state.manifest).not.toBeNull();
    expect(state.gitHost).not.toBeNull();
  });

  it('Test 2 (wrong credentials — manifest_decrypt throws): ok:false, handle locked+freed, no side effects', async () => {
    const wasm = makeWasm({
      manifest_decrypt: vi.fn(() => { throw new Error('AEAD verification failed'); }),
    });
    const state = makeState(wasm);

    const resp = await vault.handleAttachVault(ATTACH_MSG, state);

    expect(resp.ok).toBe(false);
    if (resp.ok) throw new Error('expected ok:false');
    expect(resp.error).toBeTruthy();
    expect(resp.error.length).toBeGreaterThan(0);

    // register_device must NOT be called (we failed before it)
    expect(wasm.register_device).not.toHaveBeenCalled();

    // Finally block must lock then free the handle we own
    expect(wasm.lock).toHaveBeenCalledWith(wasm._handle);
    expect(wasm._handle.free).toHaveBeenCalled();

    // Session must NOT have been updated
    expect(setCurrent).not.toHaveBeenCalled();
  });
});