rspace-online/scripts/test-wallet-store.ts

/**
 * test-wallet-store.ts — Test encrypted client-side wallet store.
 *
 * Mocks localStorage since we're running in Bun (not a browser).
 *
 * Usage:
 *   bun run scripts/test-wallet-store.ts
 */

// Mock localStorage before importing anything
const storage = new Map<string, string>();
(globalThis as any).localStorage = {
  getItem: (k: string) => storage.get(k) ?? null,
  setItem: (k: string, v: string) => storage.set(k, v),
  removeItem: (k: string) => storage.delete(k),
  clear: () => storage.clear(),
};

import { WalletStore } from '../src/encryptid/wallet-store';
import { EncryptIDKeyManager } from '../src/encryptid/key-derivation';

let passed = 0;
let failed = 0;

function assert(condition: boolean, msg: string) {
  if (condition) {
    console.log(`  ✓ ${msg}`);
    passed++;
  } else {
    console.error(`  ✗ ${msg}`);
    failed++;
  }
}

async function main() {
  console.log('=== Wallet Store Tests ===\n');

  // Set up key manager with a test PRF
  const prfOutput = new ArrayBuffer(32);
  new Uint8Array(prfOutput).set(Array.from({ length: 32 }, (_, i) => i + 1));

  const km = new EncryptIDKeyManager();
  await km.initFromPRF(prfOutput);
  const keys = await km.getKeys();

  const store = new WalletStore(keys.encryptionKey);

  // Test 1: Empty store
  console.log('[1] Empty store');
  const empty = await store.list();
  assert(empty.length === 0, 'No wallets initially');
  const noDefault = await store.getDefault();
  assert(noDefault === null, 'No default wallet');

  // Test 2: Add first wallet (auto-default)
  console.log('\n[2] Add first wallet');
  const w1 = await store.add({
    safeAddress: '0x1111111111111111111111111111111111111111',
    chainId: 84532,
    eoaAddress: keys.eoaAddress!,
    label: 'Test Treasury',
  });
  assert(w1.id.length > 0, 'Has UUID');
  assert(w1.safeAddress === '0x1111111111111111111111111111111111111111', 'Address correct');
  assert(w1.chainId === 84532, 'Chain correct');
  assert(w1.isDefault === true, 'First wallet is auto-default');
  assert(w1.label === 'Test Treasury', 'Label correct');
  assert(w1.addedAt > 0, 'Has timestamp');

  // Test 3: Data is encrypted in localStorage
  console.log('\n[3] Encrypted at rest');
  const raw = storage.get('encryptid_wallets');
  assert(raw !== undefined, 'Data exists in localStorage');
  const blob = JSON.parse(raw!);
  assert(typeof blob.c === 'string', 'Has ciphertext field');
  assert(typeof blob.iv === 'string', 'Has IV field');
  assert(!raw!.includes('Treasury'), 'Label NOT in plaintext');
  assert(!raw!.includes('1111111'), 'Address NOT in plaintext');

  // Test 4: Add second wallet
  console.log('\n[4] Add second wallet');
  const w2 = await store.add({
    safeAddress: '0x2222222222222222222222222222222222222222',
    chainId: 8453,
    eoaAddress: keys.eoaAddress!,
    label: 'Mainnet Safe',
  });
  assert(w2.isDefault === false, 'Second wallet is not default');
  const all = await store.list();
  assert(all.length === 2, 'Now have 2 wallets');

  // Test 5: Get default
  console.log('\n[5] Get default');
  const def = await store.getDefault();
  assert(def !== null, 'Has default');
  assert(def!.id === w1.id, 'First wallet is still default');

  // Test 6: Get by address + chain
  console.log('\n[6] Get by address + chain');
  const found = await store.get('0x2222222222222222222222222222222222222222', 8453);
  assert(found !== null, 'Found by address+chain');
  assert(found!.label === 'Mainnet Safe', 'Correct wallet');
  const notFound = await store.get('0x2222222222222222222222222222222222222222', 1);
  assert(notFound === null, 'Not found on wrong chain');

  // Test 7: Update label and default
  console.log('\n[7] Update');
  const updated = await store.update(w2.id, { label: 'Base Mainnet', isDefault: true });
  assert(updated !== null, 'Update succeeded');
  assert(updated!.label === 'Base Mainnet', 'Label updated');
  assert(updated!.isDefault === true, 'Now default');
  const newDefault = await store.getDefault();
  assert(newDefault!.id === w2.id, 'Default switched');
  // Old default should be false now
  const allAfter = await store.list();
  const oldW1 = allAfter.find(w => w.id === w1.id);
  assert(oldW1!.isDefault === false, 'Old default cleared');

  // Test 8: Duplicate add = upsert
  console.log('\n[8] Duplicate add (upsert)');
  const w1Updated = await store.add({
    safeAddress: '0x1111111111111111111111111111111111111111',
    chainId: 84532,
    eoaAddress: keys.eoaAddress!,
    label: 'Renamed Treasury',
  });
  assert(w1Updated.label === 'Renamed Treasury', 'Label updated via upsert');
  const afterUpsert = await store.list();
  assert(afterUpsert.length === 2, 'Still 2 wallets (not 3)');

  // Test 9: Persistence — new store instance reads encrypted data
  console.log('\n[9] Persistence across instances');
  const store2 = new WalletStore(keys.encryptionKey);
  const restored = await store2.list();
  assert(restored.length === 2, 'Restored 2 wallets');
  assert(restored.some(w => w.label === 'Renamed Treasury'), 'Labels preserved');
  assert(restored.some(w => w.label === 'Base Mainnet'), 'Both wallets present');

  // Test 10: Wrong key can't decrypt
  console.log('\n[10] Wrong key fails gracefully');
  const km2 = new EncryptIDKeyManager();
  const otherPrf = new ArrayBuffer(32);
  new Uint8Array(otherPrf).set(Array.from({ length: 32 }, (_, i) => 255 - i));
  await km2.initFromPRF(otherPrf);
  const otherKeys = await km2.getKeys();
  const storeWrongKey = new WalletStore(otherKeys.encryptionKey);
  const wrongResult = await storeWrongKey.list();
  assert(wrongResult.length === 0, 'Wrong key returns empty (graceful failure)');
  km2.clear();

  // Test 11: Remove wallet
  console.log('\n[11] Remove');
  const removed = await store.remove(w2.id);
  assert(removed === true, 'Remove succeeded');
  const afterRemove = await store.list();
  assert(afterRemove.length === 1, 'Down to 1 wallet');
  // Removed wallet was default, so remaining should be promoted
  assert(afterRemove[0].isDefault === true, 'Remaining wallet promoted to default');

  // Test 12: Remove non-existent
  const removedAgain = await store.remove('nonexistent-id');
  assert(removedAgain === false, 'Remove non-existent returns false');

  // Test 13: Clear
  console.log('\n[13] Clear');
  await store.clear();
  const afterClear = await store.list();
  assert(afterClear.length === 0, 'Empty after clear');
  assert(!storage.has('encryptid_wallets'), 'localStorage key removed');

  // Cleanup
  km.clear();

  console.log(`\n=== Results: ${passed} passed, ${failed} failed ===`);
  process.exit(failed > 0 ? 1 : 0);
}

main().catch(err => {
  console.error('Fatal:', err);
  process.exit(1);
});