Shade/packages/shade-streams/tests/coordinator.test.ts

import { describe, test, expect } from 'bun:test';
import { SubtleCryptoProvider } from '@shade/crypto-web';
import { ValidationError } from '@shade/core';
import {
  MultiLaneSender,
  MultiLaneReceiver,
  StreamProtocolError,
  generateStreamId,
  generateStreamSecret,
  planRangePartition,
  planRoundRobinPartition,
  chunkRange,
  sha256Once,
} 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('');
}

/**
 * Roundtrip a fixed input through `laneCount` lanes using range partitioning.
 * Returns the per-side overall sha256 + the reconstructed plaintext.
 */
async function roundtripRange(input: Uint8Array, laneCount: number, chunkSize: number) {
  const streamId = generateStreamId(crypto);
  const streamSecret = generateStreamSecret(crypto);
  const lanes = planRangePartition(input.length, laneCount);

  const sender = await MultiLaneSender.create({ crypto, streamId, streamSecret, lanes });
  const receiver = await MultiLaneReceiver.create({ crypto, streamId, streamSecret, lanes });

  // Append the entire input to the sender's overall hasher in original order
  // (range mode: lane i's slice is contiguous in original order).
  sender.appendOverall(input);

  // Encrypt all chunks for all lanes (interleaved as a real consumer would).
  const wireChunks: Array<{ laneId: number; bytes: Uint8Array }> = [];
  for (const lane of lanes) {
    if (lane.partition.kind !== 'range') throw new Error('expected range');
    const slices = chunkRange(lane.partition.startByte, lane.partition.endByte, chunkSize);
    for (let i = 0; i < slices.length; i++) {
      const s = slices[i]!;
      const isLast = i === slices.length - 1;
      const plaintext = input.subarray(s.start, s.end);
      const { bytes } = await sender.encryptForLane(lane.laneId, plaintext, isLast);
      wireChunks.push({ laneId: lane.laneId, bytes });
    }
  }

  // Receiver decrypts. Range mode: gather lane outputs in laneId order.
  const laneBuffers = new Map<number, Uint8Array[]>();
  for (const { bytes } of wireChunks) {
    const dec = await receiver.decryptChunk(bytes);
    if (!laneBuffers.has(dec.laneId)) laneBuffers.set(dec.laneId, []);
    laneBuffers.get(dec.laneId)!.push(dec.plaintext);
  }
  // Concatenate lane outputs in laneId order to rebuild original byte order.
  const reconstructed: Uint8Array[] = [];
  for (let i = 0; i < laneCount; i++) {
    for (const piece of laneBuffers.get(i) ?? []) reconstructed.push(piece);
  }
  // Feed receiver's overall hasher in original byte order.
  for (const piece of reconstructed) receiver.appendOverall(piece);

  return {
    sender,
    receiver,
    senderOverall: sender.getOverallSha256(),
    receiverOverall: receiver.getOverallSha256(),
    reconstructed: concat(reconstructed),
  };
}

/** Roundtrip via round-robin partitioning. Chunk i goes to lane (i mod L). */
async function roundtripRoundRobin(
  input: Uint8Array,
  laneCount: number,
  chunkSize: number,
) {
  const streamId = generateStreamId(crypto);
  const streamSecret = generateStreamSecret(crypto);
  const lanes = planRoundRobinPartition(laneCount);

  const sender = await MultiLaneSender.create({ crypto, streamId, streamSecret, lanes });
  const receiver = await MultiLaneReceiver.create({ crypto, streamId, streamSecret, lanes });

  // Append in original order.
  sender.appendOverall(input);

  // Slice into chunks; round-robin assignment.
  const slices = chunkRange(0, input.length, chunkSize);
  // Determine `isLast` for each lane (last chunk this lane sees).
  const lastChunkByLane = new Map<number, number>();
  for (let i = 0; i < slices.length; i++) {
    lastChunkByLane.set(i % laneCount, i);
  }
  const wireChunks: Array<{ chunkIndex: number; bytes: Uint8Array }> = [];
  for (let i = 0; i < slices.length; i++) {
    const s = slices[i]!;
    const laneId = i % laneCount;
    const isLast = lastChunkByLane.get(laneId) === i;
    const plaintext = input.subarray(s.start, s.end);
    const { bytes } = await sender.encryptForLane(laneId, plaintext, isLast);
    wireChunks.push({ chunkIndex: i, bytes });
  }

  // Receiver: collect chunks; reorder by chunkIndex (the original-order index).
  const decoded = new Map<number, Uint8Array>();
  for (const { chunkIndex, bytes } of wireChunks) {
    const dec = await receiver.decryptChunk(bytes);
    decoded.set(chunkIndex, dec.plaintext);
  }
  const reconstructed: Uint8Array[] = [];
  for (let i = 0; i < slices.length; i++) {
    reconstructed.push(decoded.get(i)!);
  }
  for (const piece of reconstructed) receiver.appendOverall(piece);

  return {
    sender,
    receiver,
    senderOverall: sender.getOverallSha256(),
    receiverOverall: receiver.getOverallSha256(),
    reconstructed: concat(reconstructed),
  };
}

function concat(parts: Uint8Array[]): Uint8Array {
  const total = parts.reduce((s, p) => s + p.length, 0);
  const out = new Uint8Array(total);
  let off = 0;
  for (const p of parts) {
    out.set(p, off);
    off += p.length;
  }
  return out;
}

describe('MultiLaneSender / MultiLaneReceiver — basic shape', () => {
  test('rejects empty lanes array', async () => {
    const streamId = generateStreamId(crypto);
    const streamSecret = generateStreamSecret(crypto);
    await expect(
      MultiLaneSender.create({ crypto, streamId, streamSecret, lanes: [] }),
    ).rejects.toThrow(ValidationError);
  });

  test('rejects duplicate laneIds', async () => {
    const streamId = generateStreamId(crypto);
    const streamSecret = generateStreamSecret(crypto);
    await expect(
      MultiLaneSender.create({
        crypto,
        streamId,
        streamSecret,
        lanes: [
          { laneId: 0, partition: { kind: 'round-robin', lane: 0, count: 2 } },
          { laneId: 0, partition: { kind: 'round-robin', lane: 1, count: 2 } },
        ],
      }),
    ).rejects.toThrow(ValidationError);
  });

  test('encryptForLane on unknown laneId throws StreamProtocolError', async () => {
    const sender = await MultiLaneSender.create({
      crypto,
      streamId: generateStreamId(crypto),
      streamSecret: generateStreamSecret(crypto),
      lanes: planRoundRobinPartition(2),
    });
    await expect(sender.encryptForLane(99, new Uint8Array(0), false)).rejects.toThrow(
      StreamProtocolError,
    );
  });
});

describe('Range-partition roundtrip', () => {
  test('1 KB / 4 lanes / 256 B chunk', async () => {
    const input = crypto.randomBytes(1024);
    const r = await roundtripRange(input, 4, 256);
    expect(r.reconstructed).toEqual(input);
    expect(hex(r.senderOverall)).toBe(hex(r.receiverOverall));
    expect(hex(r.senderOverall)).toBe(hex(sha256Once(input)));
  });

  test('exactly chunkSize-aligned input', async () => {
    const input = crypto.randomBytes(8 * 256);
    const r = await roundtripRange(input, 4, 256);
    expect(r.reconstructed).toEqual(input);
    expect(hex(r.senderOverall)).toBe(hex(r.receiverOverall));
  });

  test('input smaller than chunkSize × laneCount', async () => {
    const input = crypto.randomBytes(50);
    const r = await roundtripRange(input, 4, 64);
    expect(r.reconstructed).toEqual(input);
    expect(hex(r.senderOverall)).toBe(hex(r.receiverOverall));
  });
});

describe('Round-robin partition roundtrip', () => {
  test('4 lanes, 1 KB / 128 B chunks', async () => {
    const input = crypto.randomBytes(1024);
    const r = await roundtripRoundRobin(input, 4, 128);
    expect(r.reconstructed).toEqual(input);
    expect(hex(r.senderOverall)).toBe(hex(r.receiverOverall));
  });
});

describe('Lane-parity ship-gate (1 / 4 / 16 lanes → same overallSha256)', () => {
  const sizes = [
    { label: '1 KiB', bytes: 1024 },
    { label: '256 KiB', bytes: 256 * 1024 },
    { label: '2 MiB', bytes: 2 * 1024 * 1024 },
  ];

  for (const { label, bytes } of sizes) {
    test(`${label} input — same sha256 across {1, 4, 16} lanes (range)`, async () => {
      const input = crypto.randomBytes(bytes);
      const expected = hex(sha256Once(input));
      for (const laneCount of [1, 4, 16]) {
        const r = await roundtripRange(input, laneCount, 64 * 1024);
        expect(r.reconstructed).toEqual(input);
        expect(hex(r.senderOverall)).toBe(expected);
        expect(hex(r.receiverOverall)).toBe(expected);
      }
    });
  }

  test('1 MiB input — same sha256 across {1, 4, 16} lanes (round-robin)', async () => {
    const input = crypto.randomBytes(1024 * 1024);
    const expected = hex(sha256Once(input));
    for (const laneCount of [1, 4, 16]) {
      const r = await roundtripRoundRobin(input, laneCount, 32 * 1024);
      expect(r.reconstructed).toEqual(input);
      expect(hex(r.senderOverall)).toBe(expected);
      expect(hex(r.receiverOverall)).toBe(expected);
    }
  });

  test('range and round-robin produce the same overall sha256 for the same input', async () => {
    const input = crypto.randomBytes(128 * 1024);
    const a = await roundtripRange(input, 4, 16 * 1024);
    const b = await roundtripRoundRobin(input, 4, 16 * 1024);
    expect(hex(a.senderOverall)).toBe(hex(b.senderOverall));
    expect(hex(a.receiverOverall)).toBe(hex(b.receiverOverall));
  });
});

describe('Per-lane fingerprints', () => {
  test('match between sender and receiver after roundtrip', async () => {
    const input = crypto.randomBytes(64 * 1024);
    const r = await roundtripRange(input, 4, 8 * 1024);
    const senderFps = r.sender.getLaneFingerprints();
    const receiverFps = r.receiver.getLaneFingerprints();
    expect(senderFps.length).toBe(4);
    for (let i = 0; i < 4; i++) {
      expect(hex(senderFps[i]!.sha256)).toBe(hex(receiverFps[i]!.sha256));
      expect(senderFps[i]!.byteCount).toBe(receiverFps[i]!.byteCount);
      expect(senderFps[i]!.chunkCount).toBe(receiverFps[i]!.chunkCount);
    }
  });

  test('byteCount across all lanes equals total input', async () => {
    const input = crypto.randomBytes(99 * 1024); // intentionally non-divisible
    const r = await roundtripRange(input, 4, 8 * 1024);
    const total = r.sender
      .getLaneFingerprints()
      .reduce((s, l) => s + l.byteCount, 0);
    expect(total).toBe(input.length);
  });

  test('allLanesFinished reflects per-lane completion', async () => {
    const input = crypto.randomBytes(1024);
    const r = await roundtripRange(input, 2, 256);
    expect(r.sender.allLanesFinished).toBe(true);
    expect(r.receiver.allLanesFinished).toBe(true);
  });
});