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

import type {
  WorkerCryptoProvider,
  WorkerStreamReceiver,
  WorkerStreamSender,
} from './worker-client.js';

/** Default plaintext chunk size — 256 KiB. Matches `@shade/transfer`. */
export const DEFAULT_STREAM_CHUNK_SIZE = 256 * 1024;

export interface CreateEncryptStreamOptions {
  provider: WorkerCryptoProvider;
  streamId: Uint8Array;
  streamSecret: Uint8Array;
  laneId?: number;
  /**
   * Plaintext bytes per AEAD chunk. Smaller = lower latency per chunk +
   * more postMessage overhead; larger = higher per-chunk RAM in the
   * worker. Default 256 KiB.
   */
  chunkSize?: number;
  /**
   * First sequence number this sender will emit. Default 0.
   * Use for resume.
   */
  startSeq?: number;
  /** First seq this receiver will accept; defaults to 0. */
  signal?: AbortSignal;
}

export interface CreateDecryptStreamOptions {
  provider: WorkerCryptoProvider;
  streamId: Uint8Array;
  streamSecret: Uint8Array;
  laneId?: number;
  startSeq?: number;
  signal?: AbortSignal;
}

/**
 * Build a `TransformStream<Uint8Array, Uint8Array>` that encrypts every
 * passing byte as a stream-chunk wire envelope. The actual AEAD work
 * happens in the worker — the main thread only buffers, slices, and
 * forwards.
 *
 * Output: one wire chunk per `enqueue`. Concatenation is the responsibility
 * of the downstream consumer (typically an HTTP-shipping `TransformStream`).
 */
export function createEncryptStream(opts: CreateEncryptStreamOptions): {
  stream: TransformStream<Uint8Array, Uint8Array>;
  /** Promise that resolves to the final lane sha256 once the stream finishes. */
  laneSha256: Promise<Uint8Array>;
} {
  const chunkSize = opts.chunkSize ?? DEFAULT_STREAM_CHUNK_SIZE;
  if (chunkSize <= 0) throw new Error('chunkSize must be positive');

  // Plaintext slices accumulate here until we have at least `chunkSize`
  // bytes (so we emit fixed-size chunks except for the very last one).
  let pending: Uint8Array = new Uint8Array(0);
  let sender: WorkerStreamSender | null = null;
  let resolveLaneSha: (b: Uint8Array) => void;
  let rejectLaneSha: (e: Error) => void;
  const laneSha256 = new Promise<Uint8Array>((res, rej) => {
    resolveLaneSha = res;
    rejectLaneSha = rej;
  });

  // Cast to `Transformer<I,O>` because some TS lib versions still ship
  // the pre-2023 shape without `cancel`. Runtime supports it (Bun, all
  // modern browsers).
  const transformer = {
    async start(): Promise<void> {
      sender = await opts.provider.createStreamSender({
        streamId: opts.streamId,
        streamSecret: opts.streamSecret,
        laneId: opts.laneId ?? 0,
        startSeq: opts.startSeq ?? 0,
      });
    },
    async transform(
      chunk: Uint8Array,
      controller: TransformStreamDefaultController<Uint8Array>,
    ): Promise<void> {
      if (sender === null) throw new Error('encryptStream: sender not initialized');
      if (chunk.byteLength === 0) return;

      pending = concat(pending, chunk);

      // Emit complete chunks. Hold back the trailing partial — we don't
      // know yet whether it's the last one (which gets isLast=true).
      while (pending.byteLength >= chunkSize) {
        const slice = pending.subarray(0, chunkSize);
        const rest = pending.subarray(chunkSize);
        const out = await sender.encryptChunk(slice, false);
        controller.enqueue(out.bytes);
        // Detach `rest` from the larger backing buffer so it can be GCed.
        pending = new Uint8Array(rest);
      }
    },
    async flush(controller: TransformStreamDefaultController<Uint8Array>): Promise<void> {
      if (sender === null) throw new Error('encryptStream: sender not initialized');
      try {
        // Always emit a final chunk with isLast=true. Even if `pending`
        // is empty: receivers rely on a trailing isLast envelope to
        // mark stream completion.
        const out = await sender.encryptChunk(pending, true);
        controller.enqueue(out.bytes);
        pending = new Uint8Array(0);
        const sha = await sender.getLaneSha256();
        resolveLaneSha(sha);
      } catch (err) {
        rejectLaneSha(err instanceof Error ? err : new Error(String(err)));
        throw err;
      } finally {
        await sender.destroy();
        sender = null;
      }
    },
    async cancel(reason: unknown): Promise<void> {
      try {
        if (sender !== null) await sender.destroy();
      } finally {
        sender = null;
        rejectLaneSha(reason instanceof Error ? reason : new Error(String(reason)));
      }
    },
  };
  const stream = new TransformStream<Uint8Array, Uint8Array>(
    transformer as unknown as Transformer<Uint8Array, Uint8Array>,
  );

  if (opts.signal) {
    const abort = (): void => {
      stream.writable.abort(opts.signal!.reason).catch(() => {});
    };
    if (opts.signal.aborted) abort();
    else opts.signal.addEventListener('abort', abort, { once: true });
  }

  return { stream, laneSha256 };
}

/**
 * Build a `TransformStream<Uint8Array, Uint8Array>` that decrypts wire
 * stream-chunk envelopes back into plaintext. The input chunks must be
 * complete envelopes — the caller is responsible for framing on the wire
 * (one envelope per write).
 */
export function createDecryptStream(opts: CreateDecryptStreamOptions): {
  stream: TransformStream<Uint8Array, Uint8Array>;
  /** Promise that resolves to the final lane sha256 once decryption finishes. */
  laneSha256: Promise<Uint8Array>;
} {
  let receiver: WorkerStreamReceiver | null = null;
  let resolveLaneSha: (b: Uint8Array) => void;
  let rejectLaneSha: (e: Error) => void;
  const laneSha256 = new Promise<Uint8Array>((res, rej) => {
    resolveLaneSha = res;
    rejectLaneSha = rej;
  });

  const transformer = {
    async start(): Promise<void> {
      receiver = await opts.provider.createStreamReceiver({
        streamId: opts.streamId,
        streamSecret: opts.streamSecret,
        laneId: opts.laneId ?? 0,
        startSeq: opts.startSeq ?? 0,
      });
    },
    async transform(
      chunk: Uint8Array,
      controller: TransformStreamDefaultController<Uint8Array>,
    ): Promise<void> {
      if (receiver === null) throw new Error('decryptStream: receiver not initialized');
      const dec = await receiver.decryptChunk(chunk);
      if (dec.plaintext.byteLength > 0) controller.enqueue(dec.plaintext);
      if (dec.isLast) {
        const sha = await receiver.getLaneSha256();
        resolveLaneSha(sha);
      }
    },
    async flush(): Promise<void> {
      if (receiver !== null) await receiver.destroy();
      receiver = null;
    },
    async cancel(reason: unknown): Promise<void> {
      try {
        if (receiver !== null) await receiver.destroy();
      } finally {
        receiver = null;
        rejectLaneSha(reason instanceof Error ? reason : new Error(String(reason)));
      }
    },
  };
  const stream = new TransformStream<Uint8Array, Uint8Array>(
    transformer as unknown as Transformer<Uint8Array, Uint8Array>,
  );

  if (opts.signal) {
    const abort = (): void => {
      stream.writable.abort(opts.signal!.reason).catch(() => {});
    };
    if (opts.signal.aborted) abort();
    else opts.signal.addEventListener('abort', abort, { once: true });
  }

  return { stream, laneSha256 };
}

function concat(a: Uint8Array, b: Uint8Array): Uint8Array {
  if (a.byteLength === 0) return b;
  if (b.byteLength === 0) return a;
  const out = new Uint8Array(a.byteLength + b.byteLength);
  out.set(a, 0);
  out.set(b, a.byteLength);
  return out;
}