Shade/packages/shade-streams/tests/sender-receiver.test.ts

import { describe, test, expect } from 'bun:test';
import { SubtleCryptoProvider } from '@shade/crypto-web';
import {
  StreamSender,
  StreamReceiver,
  StreamFinishedError,
  generateStreamId,
  generateStreamSecret,
} from '../src/index.js';

const crypto = new SubtleCryptoProvider();

function hex(b: Uint8Array): string {
  return Array.from(b, (x) => x.toString(16).padStart(2, '0')).join('');
}

async function makePair(opts?: { laneId?: number; startSeq?: number }) {
  const streamId = generateStreamId(crypto);
  const streamSecret = generateStreamSecret(crypto);
  const laneId = opts?.laneId ?? 0;
  const sender = await StreamSender.create({
    crypto,
    streamId,
    streamSecret,
    laneId,
    ...(opts?.startSeq !== undefined ? { startSeq: opts.startSeq } : {}),
  });
  const receiver = await StreamReceiver.create({
    crypto,
    streamId,
    streamSecret,
    laneId,
    ...(opts?.startSeq !== undefined ? { startSeq: opts.startSeq } : {}),
  });
  return { sender, receiver, streamId, streamSecret };
}

describe('Single-lane sender/receiver roundtrip', () => {
  test('basic single-chunk transfer', async () => {
    const { sender, receiver } = await makePair();
    const plaintext = new TextEncoder().encode('hello shade');
    const { bytes } = await sender.encryptChunk(plaintext, true);
    const decrypted = await receiver.decryptChunk(bytes);
    expect(new TextDecoder().decode(decrypted.plaintext)).toBe('hello shade');
    expect(decrypted.seq).toBe(0);
    expect(decrypted.isLast).toBe(true);
  });

  test('multi-chunk transfer with monotonic seq', async () => {
    const { sender, receiver } = await makePair();
    const chunks = ['alpha', 'beta', 'gamma', 'delta'];
    for (let i = 0; i < chunks.length; i++) {
      const isLast = i === chunks.length - 1;
      const { bytes, seq } = await sender.encryptChunk(
        new TextEncoder().encode(chunks[i]!),
        isLast,
      );
      expect(seq).toBe(i);
      const dec = await receiver.decryptChunk(bytes);
      expect(new TextDecoder().decode(dec.plaintext)).toBe(chunks[i]);
      expect(dec.seq).toBe(i);
      expect(dec.isLast).toBe(isLast);
    }
  });

  test('lane sha256 matches between sender and receiver', async () => {
    const { sender, receiver } = await makePair();
    const data = [
      crypto.randomBytes(1024),
      crypto.randomBytes(2048),
      crypto.randomBytes(512),
    ];
    for (let i = 0; i < data.length; i++) {
      const { bytes } = await sender.encryptChunk(data[i]!, i === data.length - 1);
      await receiver.decryptChunk(bytes);
    }
    expect(hex(sender.getLaneSha256Digest())).toBe(hex(receiver.getLaneSha256Digest()));
  });

  test('handles empty chunks', async () => {
    const { sender, receiver } = await makePair();
    const { bytes } = await sender.encryptChunk(new Uint8Array(0), true);
    const dec = await receiver.decryptChunk(bytes);
    expect(dec.plaintext.length).toBe(0);
    expect(dec.isLast).toBe(true);
  });

  test('ship-gate: ~10 MiB roundtrip preserves byte-for-byte content', async () => {
    const { sender, receiver } = await makePair();
    const total = 10 * 1024 * 1024;
    const chunkSize = 256 * 1024;
    const allBytes = crypto.randomBytes(total);

    const reconstructed: Uint8Array[] = [];
    for (let off = 0; off < total; off += chunkSize) {
      const slice = allBytes.subarray(off, Math.min(off + chunkSize, total));
      const isLast = off + chunkSize >= total;
      const { bytes } = await sender.encryptChunk(slice, isLast);
      const dec = await receiver.decryptChunk(bytes);
      reconstructed.push(dec.plaintext);
    }

    let off = 0;
    for (const piece of reconstructed) {
      for (let i = 0; i < piece.length; i++) {
        if (piece[i] !== allBytes[off + i]) {
          throw new Error(`mismatch at byte ${off + i}`);
        }
      }
      off += piece.length;
    }
    expect(off).toBe(total);
    expect(hex(sender.getLaneSha256Digest())).toBe(hex(receiver.getLaneSha256Digest()));
  });

  test('byte counters track encrypted/decrypted plaintext', async () => {
    const { sender, receiver } = await makePair();
    const a = crypto.randomBytes(100);
    const b = crypto.randomBytes(250);
    const { bytes: ab } = await sender.encryptChunk(a, false);
    const { bytes: bb } = await sender.encryptChunk(b, true);
    await receiver.decryptChunk(ab);
    await receiver.decryptChunk(bb);
    expect(sender.bytesSent).toBe(350n);
    expect(receiver.bytesReceived).toBe(350n);
  });

  test('finished flag set after isLast', async () => {
    const { sender, receiver } = await makePair();
    const { bytes } = await sender.encryptChunk(new Uint8Array(8), true);
    expect(sender.isFinished).toBe(true);
    await receiver.decryptChunk(bytes);
    expect(receiver.isFinished).toBe(true);
  });

  test('encryptChunk after finish throws StreamFinishedError', async () => {
    const { sender } = await makePair();
    await sender.encryptChunk(new Uint8Array(0), true);
    await expect(sender.encryptChunk(new Uint8Array(0), false)).rejects.toThrow(
      StreamFinishedError,
    );
  });

  test('decryptChunk after finish throws StreamFinishedError', async () => {
    const { sender, receiver } = await makePair();
    const { bytes: a } = await sender.encryptChunk(new Uint8Array(8), true);
    await receiver.decryptChunk(a);
    // Try to feed another chunk — sender wouldn't normally produce one, but
    // simulate an attacker sending bytes after the legitimate isLast.
    const sender2 = await StreamSender.create({
      crypto,
      streamId: (sender as unknown as { streamId: Uint8Array }).streamId,
      streamSecret: new Uint8Array(32),
      laneId: 0,
    });
    const { bytes: extra } = await sender2.encryptChunk(new Uint8Array(8), false);
    await expect(receiver.decryptChunk(extra)).rejects.toThrow(StreamFinishedError);
  });

  test('destroy zeroes the lane key (subsequent calls throw)', async () => {
    const { sender, receiver } = await makePair();
    sender.destroy();
    receiver.destroy();
    await expect(sender.encryptChunk(new Uint8Array(0), false)).rejects.toThrow(
      StreamFinishedError,
    );
  });

  test('startSeq enables resume from arbitrary offset', async () => {
    const { sender, receiver } = await makePair({ startSeq: 100 });
    const { bytes, seq } = await sender.encryptChunk(new TextEncoder().encode('mid'), false);
    expect(seq).toBe(100);
    const dec = await receiver.decryptChunk(bytes);
    expect(dec.seq).toBe(100);
  });
});