synor/apps/web/src/lib/wasm-crypto.ts

/**
 * WASM Crypto Module Loader
 *
 * This module provides lazy-loading of the synor-crypto-wasm module
 * for client-side Dilithium3 post-quantum signatures.
 *
 * ## Why Lazy Loading?
 *
 * The WASM module is ~2MB, so we only load it when needed:
 * - User opts into client-side signing
 * - Hardware wallet mode (all keys local)
 * - Offline signing scenarios
 *
 * ## Usage
 *
 * ```typescript
 * import { loadWasmCrypto, isWasmLoaded } from './wasm-crypto';
 *
 * // Load the WASM module (only loads once)
 * const wasm = await loadWasmCrypto();
 *
 * // Create Dilithium keypair from seed
 * const dilithiumKey = wasm.DilithiumSigningKey.fromSeed(seed);
 * const signature = dilithiumKey.sign(message);
 * ```
 */

// Re-export types from the generated WASM module
// These are auto-generated by wasm-bindgen and match synor-crypto-wasm exports

import type {
  DilithiumSigningKey as WasmDilithiumSigningKeyType,
  Keypair as WasmKeypairType,
  Mnemonic as WasmMnemonicType,
} from '../wasm/synor_crypto';

// Export the types for external use
export type WasmDilithiumSigningKey = WasmDilithiumSigningKeyType;
export type WasmKeypair = WasmKeypairType;
export type WasmMnemonic = WasmMnemonicType;

// Interface for the loaded WASM module
export interface SynorCryptoWasm {
  // Classes
  Keypair: typeof WasmKeypairType;
  DilithiumSigningKey: typeof WasmDilithiumSigningKeyType;
  Mnemonic: typeof WasmMnemonicType;

  // Standalone functions
  verifyWithPublicKey(publicKey: Uint8Array, message: Uint8Array, signature: Uint8Array): boolean;
  dilithiumVerify(signature: Uint8Array, message: Uint8Array, publicKey: Uint8Array): boolean;
  dilithiumSizes(): { publicKey: number; secretKey: number; signature: number };
  sha3_256(data: Uint8Array): Uint8Array;
  blake3(data: Uint8Array): Uint8Array;
  deriveKey(inputKey: Uint8Array, salt: Uint8Array, info: Uint8Array, outputLen: number): Uint8Array;
  validateAddress(address: string): boolean;
  decodeAddress(address: string): unknown;
}

// Singleton instance of the loaded WASM module
let wasmModule: SynorCryptoWasm | null = null;
let loadingPromise: Promise<SynorCryptoWasm> | null = null;

/**
 * Check if the WASM module is already loaded.
 */
export function isWasmLoaded(): boolean {
  return wasmModule !== null;
}

/**
 * Load the WASM crypto module.
 *
 * This function is idempotent - calling it multiple times will return
 * the same module instance.
 *
 * @returns Promise resolving to the WASM module
 * @throws Error if WASM loading fails
 */
export async function loadWasmCrypto(): Promise<SynorCryptoWasm> {
  // Return cached module if already loaded
  if (wasmModule) {
    return wasmModule;
  }

  // Return existing loading promise to avoid parallel loads
  if (loadingPromise) {
    return loadingPromise;
  }

  // Start loading the WASM module
  loadingPromise = (async () => {
    try {
      // Dynamic import for code splitting - Vite handles WASM bundling
      const wasm = await import('../wasm/synor_crypto');

      // For wasm-bindgen bundler target, the default export is the init function
      // Some bundlers auto-init, others require explicit init call
      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      const initFn = (wasm as any).default;
      if (typeof initFn === 'function') {
        await initFn();
      }

      // Cast to our interface type
      wasmModule = {
        Keypair: wasm.Keypair,
        DilithiumSigningKey: wasm.DilithiumSigningKey,
        Mnemonic: wasm.Mnemonic,
        verifyWithPublicKey: wasm.verifyWithPublicKey,
        dilithiumVerify: wasm.dilithiumVerify,
        dilithiumSizes: wasm.dilithiumSizes as () => { publicKey: number; secretKey: number; signature: number },
        sha3_256: wasm.sha3_256,
        blake3: wasm.blake3,
        deriveKey: wasm.deriveKey,
        validateAddress: wasm.validateAddress,
        decodeAddress: wasm.decodeAddress,
      };

      return wasmModule;
    } catch (error) {
      loadingPromise = null; // Reset so it can be retried
      throw new Error(
        `Failed to load WASM crypto module: ${error instanceof Error ? error.message : 'Unknown error'}. ` +
        'Make sure the WASM module is built and copied to apps/web/src/wasm/'
      );
    }
  })();

  return loadingPromise;
}

/**
 * Unload the WASM module to free memory.
 * This is useful for single-page apps that want to reclaim memory
 * after signing operations are complete.
 */
export function unloadWasmCrypto(): void {
  wasmModule = null;
  loadingPromise = null;
}

// ==================== High-Level WASM Crypto Functions ====================

/**
 * Create a Dilithium3 keypair from a 32-byte seed.
 *
 * The seed should be derived from the same mnemonic as the Ed25519 key
 * to maintain key correlation for the hybrid signature scheme.
 *
 * @param seed - 32-byte seed (typically from BIP-39 mnemonic)
 * @returns Dilithium signing key
 */
export async function createDilithiumKeyFromSeed(
  seed: Uint8Array
): Promise<WasmDilithiumSigningKey> {
  const wasm = await loadWasmCrypto();
  return wasm.DilithiumSigningKey.fromSeed(seed);
}

/**
 * Sign a message with Dilithium3.
 *
 * @param message - Message to sign
 * @param seed - 32-byte seed for key derivation
 * @returns 3293-byte Dilithium signature
 */
export async function signWithDilithium(
  message: Uint8Array,
  seed: Uint8Array
): Promise<Uint8Array> {
  const key = await createDilithiumKeyFromSeed(seed);
  try {
    return key.sign(message);
  } finally {
    key.free(); // Clean up WASM memory
  }
}

/**
 * Verify a Dilithium3 signature.
 *
 * @param message - Original message
 * @param signature - Dilithium signature (3293 bytes)
 * @param publicKey - Dilithium public key (1952 bytes)
 * @returns true if signature is valid
 */
export async function verifyDilithiumSignature(
  message: Uint8Array,
  signature: Uint8Array,
  publicKey: Uint8Array
): Promise<boolean> {
  const wasm = await loadWasmCrypto();
  return wasm.dilithiumVerify(signature, message, publicKey);
}

/**
 * Get Dilithium3 key and signature sizes.
 */
export async function getDilithiumSizes(): Promise<{
  publicKey: number;
  secretKey: number;
  signature: number;
}> {
  const wasm = await loadWasmCrypto();
  return wasm.dilithiumSizes();
}

// ==================== Ed25519 WASM Functions ====================

/**
 * Create an Ed25519 keypair from seed using WASM.
 *
 * This is an alternative to @noble/ed25519 that keeps all crypto in WASM.
 * Useful for consistency or when noble packages aren't available.
 */
export async function createEd25519KeyFromSeed(
  seed: Uint8Array
): Promise<WasmKeypair> {
  const wasm = await loadWasmCrypto();
  return wasm.Keypair.fromSeed(seed);
}

/**
 * Sign a message with Ed25519 using WASM.
 */
export async function signWithEd25519Wasm(
  message: Uint8Array,
  seed: Uint8Array
): Promise<Uint8Array> {
  const key = await createEd25519KeyFromSeed(seed);
  try {
    return key.sign(message);
  } finally {
    key.free();
  }
}