Shade/packages/shade-crypto-web/src/worker-client.ts

import type { CryptoProvider } from '@shade/core';
import {
  WORKER_PROTOCOL_VERSION,
  fromTransferable,
  toTransferableCopy,
  type WorkerRequest,
  type WorkerResponse,
  type WorkerResult,
} from './worker-protocol.js';

/** Distributive omit of `id` from each variant of {@link WorkerRequest}. */
type WorkerRequestBody = WorkerRequest extends infer T
  ? T extends { id: number }
    ? Omit<T, 'id'>
    : never
  : never;

/**
 * Minimal Worker shape we depend on. Lets the main-thread proxy work
 * against both browser `Worker` and Bun's `Worker` without dragging in
 * `lib.dom.d.ts`.
 */
export interface WorkerLike {
  postMessage(message: unknown, transfer?: ArrayBuffer[]): void;
  addEventListener(
    type: 'message',
    listener: (ev: { data: WorkerResponse }) => void,
  ): void;
  removeEventListener(
    type: 'message',
    listener: (ev: { data: WorkerResponse }) => void,
  ): void;
  addEventListener(
    type: 'error',
    listener: (ev: { message?: string; error?: unknown }) => void,
  ): void;
  terminate(): void;
}

export interface WorkerCryptoProviderOptions {
  /**
   * URL of the bundled worker entry. Required because every bundler
   * resolves worker URLs differently — supply yours and stop guessing.
   *
   *   // Vite / Webpack 5 / Rspack:
   *   workerUrl: new URL('@shade/crypto-web/worker', import.meta.url)
   */
  workerUrl: URL | string;
  /**
   * How long the worker may stay idle before it self-terminates. Set to
   * `Infinity` to opt out (e.g. for SharedArrayBuffer / persistent UI).
   * Default: 30_000 ms — matches the V3.8 acceptance criterion.
   */
  idleTimeoutMs?: number;
  /**
   * Override the worker factory. Useful in tests with `bun:test`'s
   * `Worker` global, or to inject a polyfill.
   */
  spawn?: (url: URL | string) => WorkerLike;
}

/**
 * Public factory. Resolves once the worker has acknowledged the protocol
 * version handshake — so a stale bundle blows up here rather than in the
 * middle of an encrypt call.
 */
export async function createWorkerCryptoProvider(
  opts: WorkerCryptoProviderOptions,
): Promise<WorkerCryptoProvider> {
  const provider = new WorkerCryptoProvider(opts);
  await provider.handshake();
  return provider;
}

let SENDER_SEQ = 1;
let RECEIVER_SEQ = 1;

/**
 * `CryptoProvider` implementation that forwards every async call to a
 * dedicated Web Worker. Sync methods (`randomBytes`, `randomUint32`,
 * `constantTimeEqual`, `zeroize`) execute on the calling thread — they
 * are pure and instantaneous, so a worker round-trip would be silly.
 *
 * The worker is spawned on construction (lazy: see `createWorkerCryptoProvider`)
 * and terminated automatically after `idleTimeoutMs` of inactivity.
 * Subsequent calls re-spawn transparently.
 *
 * Stream sender/receiver state lives on the worker — the provider
 * exposes `createStreamSender` / `createStreamReceiver` factories that
 * return main-thread proxies (`WorkerStreamSender` / `WorkerStreamReceiver`).
 */
export class WorkerCryptoProvider implements CryptoProvider {
  private worker: WorkerLike | null = null;
  private nextRequestId = 1;
  private readonly inflight = new Map<
    number,
    { resolve: (r: WorkerResult) => void; reject: (e: Error) => void }
  >();
  private idleTimer: ReturnType<typeof setTimeout> | null = null;
  private destroyed = false;
  private readonly idleTimeoutMs: number;

  constructor(private readonly opts: WorkerCryptoProviderOptions) {
    this.idleTimeoutMs = opts.idleTimeoutMs ?? 30_000;
  }

  // ─── lifecycle ───────────────────────────────────────────

  /** Force-spawn + complete the protocol handshake. Idempotent. */
  async handshake(): Promise<void> {
    await this.ensureWorker();
    await this.send({ method: 'init', protocolVersion: WORKER_PROTOCOL_VERSION });
  }

  /** Permanently terminate the worker. After this, every call rejects. */
  async destroy(): Promise<void> {
    this.destroyed = true;
    this.terminateWorker(new Error('WorkerCryptoProvider destroyed'));
  }

  /**
   * Tear down the current worker and (lazily) spawn a fresh one. Use
   * after rotating identity keys so leaked-state worst case is bounded
   * by one rotation interval.
   */
  async rotate(): Promise<void> {
    this.terminateWorker(new Error('WorkerCryptoProvider rotated'));
  }

  // ─── async CryptoProvider methods ────────────────────────

  async generateX25519KeyPair(): Promise<{ publicKey: Uint8Array; privateKey: Uint8Array }> {
    const r = await this.send({ method: 'crypto.generateX25519KeyPair' });
    if (r.kind !== 'keypair') throw new Error('protocol: expected keypair');
    return { publicKey: fromTransferable(r.publicKey), privateKey: fromTransferable(r.privateKey) };
  }

  async x25519(privateKey: Uint8Array, publicKey: Uint8Array): Promise<Uint8Array> {
    const r = await this.send(
      {
        method: 'crypto.x25519',
        privateKey: toTransferableCopy(privateKey),
        publicKey: toTransferableCopy(publicKey),
      },
      // No transferables from caller's owned buffers — we copied above.
      [],
    );
    if (r.kind !== 'bytes') throw new Error('protocol: expected bytes');
    return fromTransferable(r.bytes);
  }

  async generateEd25519KeyPair(): Promise<{ publicKey: Uint8Array; privateKey: Uint8Array }> {
    const r = await this.send({ method: 'crypto.generateEd25519KeyPair' });
    if (r.kind !== 'keypair') throw new Error('protocol: expected keypair');
    return { publicKey: fromTransferable(r.publicKey), privateKey: fromTransferable(r.privateKey) };
  }

  async sign(privateKey: Uint8Array, message: Uint8Array): Promise<Uint8Array> {
    const r = await this.send({
      method: 'crypto.sign',
      privateKey: toTransferableCopy(privateKey),
      message: toTransferableCopy(message),
    });
    if (r.kind !== 'bytes') throw new Error('protocol: expected bytes');
    return fromTransferable(r.bytes);
  }

  async verify(
    publicKey: Uint8Array,
    message: Uint8Array,
    signature: Uint8Array,
  ): Promise<boolean> {
    const r = await this.send({
      method: 'crypto.verify',
      publicKey: toTransferableCopy(publicKey),
      message: toTransferableCopy(message),
      signature: toTransferableCopy(signature),
    });
    if (r.kind !== 'verify') throw new Error('protocol: expected verify');
    return r.valid;
  }

  async aesGcmEncrypt(
    key: Uint8Array,
    plaintext: Uint8Array,
    aad?: Uint8Array,
  ): Promise<{ ciphertext: Uint8Array; nonce: Uint8Array }> {
    const r = await this.send({
      method: 'crypto.aesGcmEncrypt',
      key: toTransferableCopy(key),
      plaintext: toTransferableCopy(plaintext),
      aad: aad ? toTransferableCopy(aad) : null,
    });
    if (r.kind !== 'aead-encrypt') throw new Error('protocol: expected aead-encrypt');
    return {
      ciphertext: fromTransferable(r.ciphertext),
      nonce: fromTransferable(r.nonce),
    };
  }

  async aesGcmDecrypt(
    key: Uint8Array,
    ciphertext: Uint8Array,
    nonce: Uint8Array,
    aad?: Uint8Array,
  ): Promise<Uint8Array> {
    const r = await this.send({
      method: 'crypto.aesGcmDecrypt',
      key: toTransferableCopy(key),
      ciphertext: toTransferableCopy(ciphertext),
      nonce: toTransferableCopy(nonce),
      aad: aad ? toTransferableCopy(aad) : null,
    });
    if (r.kind !== 'bytes') throw new Error('protocol: expected bytes');
    return fromTransferable(r.bytes);
  }

  async hkdf(
    ikm: Uint8Array,
    salt: Uint8Array,
    info: Uint8Array,
    length: number,
  ): Promise<Uint8Array> {
    const r = await this.send({
      method: 'crypto.hkdf',
      ikm: toTransferableCopy(ikm),
      salt: toTransferableCopy(salt),
      info: toTransferableCopy(info),
      length,
    });
    if (r.kind !== 'bytes') throw new Error('protocol: expected bytes');
    return fromTransferable(r.bytes);
  }

  async hmacSha256(key: Uint8Array, data: Uint8Array): Promise<Uint8Array> {
    const r = await this.send({
      method: 'crypto.hmacSha256',
      key: toTransferableCopy(key),
      data: toTransferableCopy(data),
    });
    if (r.kind !== 'bytes') throw new Error('protocol: expected bytes');
    return fromTransferable(r.bytes);
  }

  // ─── sync — local execution (no round-trip) ──────────────

  randomBytes(length: number): Uint8Array {
    const buf = new Uint8Array(length);
    globalThis.crypto.getRandomValues(buf);
    return buf;
  }

  constantTimeEqual(a: Uint8Array, b: Uint8Array): boolean {
    if (a.length !== b.length) return false;
    let diff = 0;
    for (let i = 0; i < a.length; i++) diff |= a[i]! ^ b[i]!;
    return diff === 0;
  }

  zeroize(buf: Uint8Array): void {
    buf.fill(0);
  }

  randomUint32(): number {
    const buf = this.randomBytes(4);
    return new DataView(buf.buffer, buf.byteOffset, 4).getUint32(0, false);
  }

  // ─── stream factories ────────────────────────────────────

  async createStreamSender(opts: {
    streamId: Uint8Array;
    streamSecret: Uint8Array;
    laneId: number;
    startSeq?: number;
  }): Promise<WorkerStreamSender> {
    const senderId = SENDER_SEQ++;
    await this.send({
      method: 'stream.createSender',
      senderId,
      streamId: toTransferableCopy(opts.streamId),
      streamSecret: toTransferableCopy(opts.streamSecret),
      laneId: opts.laneId,
      startSeq: opts.startSeq ?? 0,
    });
    return new WorkerStreamSender(this, senderId);
  }

  async createStreamReceiver(opts: {
    streamId: Uint8Array;
    streamSecret: Uint8Array;
    laneId: number;
    startSeq?: number;
  }): Promise<WorkerStreamReceiver> {
    const receiverId = RECEIVER_SEQ++;
    await this.send({
      method: 'stream.createReceiver',
      receiverId,
      streamId: toTransferableCopy(opts.streamId),
      streamSecret: toTransferableCopy(opts.streamSecret),
      laneId: opts.laneId,
      startSeq: opts.startSeq ?? 0,
    });
    return new WorkerStreamReceiver(this, receiverId);
  }

  // ─── internals ───────────────────────────────────────────

  /** @internal — used by `WorkerStreamSender` / `WorkerStreamReceiver`. */
  async send(
    body: WorkerRequestBody,
    extraTransferables?: ArrayBuffer[],
  ): Promise<WorkerResult> {
    if (this.destroyed) throw new Error('WorkerCryptoProvider destroyed');
    await this.ensureWorker();

    const id = this.nextRequestId++;
    const req = { id, ...body } as WorkerRequest;

    // Auto-collect transferable buffers from the request payload — every
    // ArrayBuffer-typed field is fair game.
    const transferables = extraTransferables ?? collectArrayBuffers(req);

    return new Promise<WorkerResult>((resolve, reject) => {
      this.inflight.set(id, { resolve, reject });
      try {
        this.worker!.postMessage(req, transferables);
      } catch (err) {
        this.inflight.delete(id);
        reject(err instanceof Error ? err : new Error(String(err)));
        return;
      }
      this.bumpIdleTimer();
    });
  }

  private async ensureWorker(): Promise<void> {
    if (this.destroyed) throw new Error('WorkerCryptoProvider destroyed');
    if (this.worker !== null) return;
    const spawn = this.opts.spawn ?? defaultSpawn;
    const w = spawn(this.opts.workerUrl);
    w.addEventListener('message', this.onMessage);
    w.addEventListener('error', this.onError);
    this.worker = w;
  }

  private readonly onMessage = (ev: { data: WorkerResponse }): void => {
    const res = ev.data;
    const pending = this.inflight.get(res.id);
    if (pending === undefined) return;
    this.inflight.delete(res.id);
    if (res.ok) pending.resolve(res.result);
    else {
      const err = new Error(res.error.message);
      err.name = res.error.name;
      pending.reject(err);
    }
    this.bumpIdleTimer();
  };

  private readonly onError = (ev: { message?: string; error?: unknown }): void => {
    const msg = ev.message ?? (ev.error instanceof Error ? ev.error.message : String(ev.error));
    this.terminateWorker(new Error(`worker error: ${msg}`));
  };

  private bumpIdleTimer(): void {
    if (this.idleTimer !== null) clearTimeout(this.idleTimer);
    if (!isFinite(this.idleTimeoutMs)) return;
    if (this.inflight.size > 0) return;
    this.idleTimer = setTimeout(() => {
      // No outstanding work — recycle the worker. Calls after this
      // re-spawn lazily.
      this.terminateWorker(null);
    }, this.idleTimeoutMs);
    // Don't keep node-style event loops alive solely on this timer.
    const t = this.idleTimer as unknown as { unref?: () => void };
    if (typeof t.unref === 'function') t.unref();
  }

  private terminateWorker(reason: Error | null): void {
    if (this.idleTimer !== null) {
      clearTimeout(this.idleTimer);
      this.idleTimer = null;
    }
    const w = this.worker;
    this.worker = null;

    // Reject every in-flight request so callers don't hang.
    if (reason !== null) {
      for (const [, pending] of this.inflight) pending.reject(reason);
    }
    this.inflight.clear();

    if (w !== null) {
      try {
        w.removeEventListener('message', this.onMessage);
      } catch {
        // ignore
      }
      try {
        w.terminate();
      } catch {
        // ignore
      }
    }
  }
}

/**
 * Walk the request body, collecting every `ArrayBuffer` reference so we
 * can hand them to `postMessage(_, transfer)`. Cheap because the request
 * objects are flat — at most a handful of fields.
 */
function collectArrayBuffers(req: WorkerRequest): ArrayBuffer[] {
  const out: ArrayBuffer[] = [];
  for (const v of Object.values(req as Record<string, unknown>)) {
    if (v instanceof ArrayBuffer) out.push(v);
  }
  return out;
}

function defaultSpawn(url: URL | string): WorkerLike {
  const ctor = (globalThis as unknown as { Worker?: new (u: URL | string, o?: unknown) => unknown })
    .Worker;
  if (typeof ctor !== 'function') {
    throw new Error('Worker is not available in this runtime');
  }
  return new ctor(url, { type: 'module' }) as WorkerLike;
}

/**
 * Main-thread handle on a `StreamSender` that lives in the worker. The
 * lane key never crosses thread boundaries — this proxy only ever ships
 * plaintext slices and receives wire bytes.
 */
export class WorkerStreamSender {
  private destroyed = false;
  constructor(
    private readonly provider: WorkerCryptoProvider,
    private readonly senderId: number,
  ) {}

  async encryptChunk(
    plaintext: Uint8Array,
    isLast: boolean,
  ): Promise<{ bytes: Uint8Array; seq: number }> {
    if (this.destroyed) throw new Error('WorkerStreamSender destroyed');
    const r = await this.provider.send({
      method: 'stream.encryptChunk',
      senderId: this.senderId,
      plaintext: toTransferableCopy(plaintext),
      isLast,
    });
    if (r.kind !== 'chunk-encrypt') throw new Error('protocol: expected chunk-encrypt');
    return { bytes: fromTransferable(r.bytes), seq: r.seq };
  }

  async getLaneSha256(): Promise<Uint8Array> {
    const r = await this.provider.send({
      method: 'stream.getSenderLaneSha256',
      senderId: this.senderId,
    });
    if (r.kind !== 'bytes') throw new Error('protocol: expected bytes');
    return fromTransferable(r.bytes);
  }

  async destroy(): Promise<void> {
    if (this.destroyed) return;
    this.destroyed = true;
    await this.provider.send({ method: 'stream.destroySender', senderId: this.senderId });
  }
}

export class WorkerStreamReceiver {
  private destroyed = false;
  constructor(
    private readonly provider: WorkerCryptoProvider,
    private readonly receiverId: number,
  ) {}

  async decryptChunk(
    wireBytes: Uint8Array,
  ): Promise<{ plaintext: Uint8Array; seq: number; isLast: boolean }> {
    if (this.destroyed) throw new Error('WorkerStreamReceiver destroyed');
    const r = await this.provider.send({
      method: 'stream.decryptChunk',
      receiverId: this.receiverId,
      wireBytes: toTransferableCopy(wireBytes),
    });
    if (r.kind !== 'chunk-decrypt') throw new Error('protocol: expected chunk-decrypt');
    return {
      plaintext: fromTransferable(r.plaintext),
      seq: r.seq,
      isLast: r.isLast,
    };
  }

  async getLaneSha256(): Promise<Uint8Array> {
    const r = await this.provider.send({
      method: 'stream.getReceiverLaneSha256',
      receiverId: this.receiverId,
    });
    if (r.kind !== 'bytes') throw new Error('protocol: expected bytes');
    return fromTransferable(r.bytes);
  }

  async destroy(): Promise<void> {
    if (this.destroyed) return;
    this.destroyed = true;
    await this.provider.send({ method: 'stream.destroyReceiver', receiverId: this.receiverId });
  }
}