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

/**
 * Dedicated Web Worker entry. Bundle this as a module worker:
 *
 *   // Vite / modern Webpack / Rspack
 *   const w = new Worker(new URL('@shade/crypto-web/worker', import.meta.url),
 *                        { type: 'module' });
 *
 * The main thread talks to this worker via the protocol in
 * `worker-protocol.ts`. All heavy crypto (AES-GCM, HKDF, HMAC, X25519,
 * Ed25519) and stream state (per-lane keys + seq counters + running
 * sha256) live here so the main thread is never blocked.
 *
 * Lifecycle: every request is one `postMessage` round-trip. Sender /
 * receiver state is keyed by numeric ids handed in by the main thread —
 * the worker never invents ids. `destroy*` calls zeroize lane keys.
 */

import { StreamSender, StreamReceiver } from '@shade/streams';
import { SubtleCryptoProvider } from './provider.js';
import {
  WORKER_PROTOCOL_VERSION,
  fromTransferable,
  toTransferable,
  transferablesOf,
  type WorkerRequest,
  type WorkerResponse,
  type WorkerResult,
} from './worker-protocol.js';

interface DedicatedWorkerScope {
  postMessage(data: unknown, transfer?: ArrayBuffer[]): void;
  addEventListener(
    type: 'message',
    listener: (ev: { data: WorkerRequest }) => void,
  ): void;
}

const scope = globalThis as unknown as DedicatedWorkerScope;
const provider = new SubtleCryptoProvider();
const senders = new Map<number, StreamSender>();
const receivers = new Map<number, StreamReceiver>();

scope.addEventListener('message', (ev) => {
  void handle(ev.data);
});

async function handle(req: WorkerRequest): Promise<void> {
  try {
    const result = await dispatch(req);
    const transfer = transferablesOf(result);
    const res: WorkerResponse = { id: req.id, ok: true, result };
    scope.postMessage(res, transfer);
  } catch (err) {
    const error = err instanceof Error
      ? { name: err.name, message: err.message }
      : { name: 'Error', message: String(err) };
    const res: WorkerResponse = { id: req.id, ok: false, error };
    scope.postMessage(res);
  }
}

async function dispatch(req: WorkerRequest): Promise<WorkerResult> {
  switch (req.method) {
    case 'init': {
      if (req.protocolVersion !== WORKER_PROTOCOL_VERSION) {
        throw new Error(
          `worker protocol version mismatch: main=${req.protocolVersion} worker=${WORKER_PROTOCOL_VERSION}`,
        );
      }
      return { kind: 'ack' };
    }
    case 'ping':
      return { kind: 'pong' };

    // ─── crypto.* ─────────────────────────────────────────
    case 'crypto.generateX25519KeyPair': {
      const kp = await provider.generateX25519KeyPair();
      return {
        kind: 'keypair',
        publicKey: toTransferable(kp.publicKey),
        privateKey: toTransferable(kp.privateKey),
      };
    }
    case 'crypto.x25519': {
      const out = await provider.x25519(
        fromTransferable(req.privateKey),
        fromTransferable(req.publicKey),
      );
      return { kind: 'bytes', bytes: toTransferable(out) };
    }
    case 'crypto.generateEd25519KeyPair': {
      const kp = await provider.generateEd25519KeyPair();
      return {
        kind: 'keypair',
        publicKey: toTransferable(kp.publicKey),
        privateKey: toTransferable(kp.privateKey),
      };
    }
    case 'crypto.sign': {
      const sig = await provider.sign(
        fromTransferable(req.privateKey),
        fromTransferable(req.message),
      );
      return { kind: 'bytes', bytes: toTransferable(sig) };
    }
    case 'crypto.verify': {
      const valid = await provider.verify(
        fromTransferable(req.publicKey),
        fromTransferable(req.message),
        fromTransferable(req.signature),
      );
      return { kind: 'verify', valid };
    }
    case 'crypto.aesGcmEncrypt': {
      const out = await provider.aesGcmEncrypt(
        fromTransferable(req.key),
        fromTransferable(req.plaintext),
        req.aad ? fromTransferable(req.aad) : undefined,
      );
      return {
        kind: 'aead-encrypt',
        ciphertext: toTransferable(out.ciphertext),
        nonce: toTransferable(out.nonce),
      };
    }
    case 'crypto.aesGcmDecrypt': {
      const out = await provider.aesGcmDecrypt(
        fromTransferable(req.key),
        fromTransferable(req.ciphertext),
        fromTransferable(req.nonce),
        req.aad ? fromTransferable(req.aad) : undefined,
      );
      return { kind: 'bytes', bytes: toTransferable(out) };
    }
    case 'crypto.hkdf': {
      const out = await provider.hkdf(
        fromTransferable(req.ikm),
        fromTransferable(req.salt),
        fromTransferable(req.info),
        req.length,
      );
      return { kind: 'bytes', bytes: toTransferable(out) };
    }
    case 'crypto.hmacSha256': {
      const out = await provider.hmacSha256(
        fromTransferable(req.key),
        fromTransferable(req.data),
      );
      return { kind: 'bytes', bytes: toTransferable(out) };
    }

    // ─── stream.* (sender) ────────────────────────────────
    case 'stream.createSender': {
      if (senders.has(req.senderId)) {
        throw new Error(`senderId ${req.senderId} already exists`);
      }
      const sender = await StreamSender.create({
        crypto: provider,
        streamId: fromTransferable(req.streamId),
        streamSecret: fromTransferable(req.streamSecret),
        laneId: req.laneId,
        startSeq: req.startSeq,
      });
      senders.set(req.senderId, sender);
      return { kind: 'ack' };
    }
    case 'stream.encryptChunk': {
      const sender = senders.get(req.senderId);
      if (sender === undefined) throw new Error(`unknown senderId ${req.senderId}`);
      const chunk = await sender.encryptChunk(fromTransferable(req.plaintext), req.isLast);
      return {
        kind: 'chunk-encrypt',
        bytes: toTransferable(chunk.bytes),
        seq: chunk.seq,
      };
    }
    case 'stream.getSenderLaneSha256': {
      const sender = senders.get(req.senderId);
      if (sender === undefined) throw new Error(`unknown senderId ${req.senderId}`);
      return { kind: 'bytes', bytes: toTransferable(sender.getLaneSha256Digest()) };
    }
    case 'stream.destroySender': {
      const sender = senders.get(req.senderId);
      if (sender !== undefined) {
        sender.destroy();
        senders.delete(req.senderId);
      }
      return { kind: 'ack' };
    }

    // ─── stream.* (receiver) ──────────────────────────────
    case 'stream.createReceiver': {
      if (receivers.has(req.receiverId)) {
        throw new Error(`receiverId ${req.receiverId} already exists`);
      }
      const receiver = await StreamReceiver.create({
        crypto: provider,
        streamId: fromTransferable(req.streamId),
        streamSecret: fromTransferable(req.streamSecret),
        laneId: req.laneId,
        startSeq: req.startSeq,
      });
      receivers.set(req.receiverId, receiver);
      return { kind: 'ack' };
    }
    case 'stream.decryptChunk': {
      const receiver = receivers.get(req.receiverId);
      if (receiver === undefined) throw new Error(`unknown receiverId ${req.receiverId}`);
      const dec = await receiver.decryptChunk(fromTransferable(req.wireBytes));
      return {
        kind: 'chunk-decrypt',
        plaintext: toTransferable(dec.plaintext),
        seq: dec.seq,
        isLast: dec.isLast,
      };
    }
    case 'stream.getReceiverLaneSha256': {
      const receiver = receivers.get(req.receiverId);
      if (receiver === undefined) throw new Error(`unknown receiverId ${req.receiverId}`);
      return { kind: 'bytes', bytes: toTransferable(receiver.getLaneSha256Digest()) };
    }
    case 'stream.destroyReceiver': {
      const receiver = receivers.get(req.receiverId);
      if (receiver !== undefined) {
        receiver.destroy();
        receivers.delete(req.receiverId);
      }
      return { kind: 'ack' };
    }
  }
}