Shade/packages/shade-crypto-web/tests/worker-streams.test.ts

import { afterEach, describe, expect, test } from 'bun:test';
import { sha256 } from '@noble/hashes/sha2.js';
import {
  createDecryptStream,
  createEncryptStream,
  createWorkerCryptoProvider,
  SubtleCryptoProvider,
  WorkerCryptoProvider,
} from '../src/index.js';

const WORKER_URL = new URL('../src/worker.ts', import.meta.url);
const subtle = new SubtleCryptoProvider();
let provider: WorkerCryptoProvider | null = null;

afterEach(async () => {
  if (provider) {
    await provider.destroy();
    provider = null;
  }
});

async function makeProvider(): Promise<WorkerCryptoProvider> {
  provider = await createWorkerCryptoProvider({ workerUrl: WORKER_URL });
  return provider;
}

async function readAll(rs: ReadableStream<Uint8Array>): Promise<Uint8Array> {
  const reader = rs.getReader();
  const parts: Uint8Array[] = [];
  let total = 0;
  for (;;) {
    const { done, value } = await reader.read();
    if (done) break;
    parts.push(value);
    total += value.byteLength;
  }
  const out = new Uint8Array(total);
  let off = 0;
  for (const p of parts) {
    out.set(p, off);
    off += p.byteLength;
  }
  return out;
}

function streamFromChunks(chunks: Uint8Array[]): ReadableStream<Uint8Array> {
  let i = 0;
  return new ReadableStream<Uint8Array>({
    pull(controller) {
      if (i < chunks.length) controller.enqueue(chunks[i++]!);
      else controller.close();
    },
  });
}

describe('encryptStream / decryptStream — round-trip', () => {
  test('round-trips small payload exactly', async () => {
    const p = await makeProvider();
    const streamId = subtle.randomBytes(16);
    const streamSecret = subtle.randomBytes(32);
    const plaintext = new TextEncoder().encode('hello stream');

    const enc = await createEncryptStream({
      provider: p,
      streamId,
      streamSecret,
      chunkSize: 1024,
    });
    const wireBytes = await readAll(
      streamFromChunks([plaintext]).pipeThrough(enc.stream),
    );

    // Frame: each enqueue is one wire envelope. We can't trivially split
    // a concatenated buffer back into envelopes, but we know how many
    // chunks were emitted (len/chunkSize, plus the final isLast). Easier
    // path: collect them as separate writes through a side channel.
    const chunks: Uint8Array[] = [];
    await streamFromChunks([plaintext])
      .pipeThrough(
        (
          await createEncryptStream({
            provider: p,
            streamId,
            streamSecret,
            chunkSize: 1024,
          })
        ).stream,
      )
      .pipeTo(
        new WritableStream<Uint8Array>({
          write(c) {
            chunks.push(c);
          },
        }),
      );

    const dec = await createDecryptStream({ provider: p, streamId, streamSecret });
    const recovered = await readAll(streamFromChunks(chunks).pipeThrough(dec.stream));
    expect(recovered).toEqual(plaintext);
    expect(wireBytes.byteLength).toBeGreaterThan(plaintext.byteLength); // overhead
  });

  test('round-trips multi-chunk payload with sha256 parity', async () => {
    const p = await makeProvider();
    const streamId = subtle.randomBytes(16);
    const streamSecret = subtle.randomBytes(32);
    const total = 750 * 1024; // 750 KiB → forces 3+ chunks at 256 KiB
    const plaintext = subtle.randomBytes(total);
    const expectedSha = sha256(plaintext);

    const enc = await createEncryptStream({
      provider: p,
      streamId,
      streamSecret,
      chunkSize: 256 * 1024,
    });

    const wireChunks: Uint8Array[] = [];
    await streamFromChunks([plaintext])
      .pipeThrough(enc.stream)
      .pipeTo(
        new WritableStream<Uint8Array>({
          write(c) {
            wireChunks.push(c);
          },
        }),
      );

    // 750 KiB / 256 KiB = 2 full chunks + 1 final (238 KiB, isLast=true)
    expect(wireChunks.length).toBe(3);

    const senderLaneSha = await enc.laneSha256;
    expect(senderLaneSha).toEqual(expectedSha);

    const dec = await createDecryptStream({
      provider: p,
      streamId,
      streamSecret,
    });
    const recovered = await readAll(streamFromChunks(wireChunks).pipeThrough(dec.stream));
    expect(recovered).toEqual(plaintext);
    expect(await dec.laneSha256).toEqual(expectedSha);
  });

  test('fragmented input produces same output as single-shot', async () => {
    const p = await makeProvider();
    const streamId = subtle.randomBytes(16);
    const streamSecret = subtle.randomBytes(32);
    const plaintext = subtle.randomBytes(50_000);

    async function run(parts: Uint8Array[]): Promise<Uint8Array[]> {
      const wire: Uint8Array[] = [];
      const e = await createEncryptStream({
        provider: p!,
        streamId,
        streamSecret,
        chunkSize: 8 * 1024,
      });
      await streamFromChunks(parts)
        .pipeThrough(e.stream)
        .pipeTo(new WritableStream({ write: (c) => void wire.push(c) }));
      return wire;
    }

    const single = await run([plaintext]);
    const split = await run([
      plaintext.subarray(0, 17_000),
      plaintext.subarray(17_000, 33_000),
      plaintext.subarray(33_000),
    ]);
    expect(split.length).toBe(single.length);
    for (let i = 0; i < single.length; i++) {
      // Same chunk size, same lane key, same seq — wire bytes match
      // byte-for-byte (deterministic nonces + AEAD).
      expect(split[i]).toEqual(single[i]!);
    }
  });

  test('100 KiB stream end-to-end completes', async () => {
    const p = await makeProvider();
    const streamId = subtle.randomBytes(16);
    const streamSecret = subtle.randomBytes(32);
    const plaintext = subtle.randomBytes(100 * 1024);

    const enc = await createEncryptStream({
      provider: p,
      streamId,
      streamSecret,
      chunkSize: 16 * 1024,
    });
    const wire: Uint8Array[] = [];
    await streamFromChunks([plaintext])
      .pipeThrough(enc.stream)
      .pipeTo(new WritableStream({ write: (c) => void wire.push(c) }));

    const dec = await createDecryptStream({ provider: p, streamId, streamSecret });
    const out = await readAll(streamFromChunks(wire).pipeThrough(dec.stream));
    expect(out).toEqual(plaintext);
    expect(await dec.laneSha256).toEqual(await enc.laneSha256);
  });

  test('decryptStream rejects out-of-order chunks', async () => {
    const p = await makeProvider();
    const streamId = subtle.randomBytes(16);
    const streamSecret = subtle.randomBytes(32);
    const plaintext = subtle.randomBytes(20_000);

    const enc = await createEncryptStream({
      provider: p,
      streamId,
      streamSecret,
      chunkSize: 4 * 1024,
    });
    const wire: Uint8Array[] = [];
    await streamFromChunks([plaintext])
      .pipeThrough(enc.stream)
      .pipeTo(new WritableStream({ write: (c) => void wire.push(c) }));

    expect(wire.length).toBeGreaterThan(2);
    // Swap first and second chunk
    [wire[0], wire[1]] = [wire[1]!, wire[0]!];

    const dec = await createDecryptStream({ provider: p, streamId, streamSecret });
    await expect(
      streamFromChunks(wire).pipeThrough(dec.stream).pipeTo(
        new WritableStream({ write() {} }),
      ),
    ).rejects.toThrow();
  });
});