Shade/packages/shade-proto/tests/wire.test.ts

import { describe, test, expect } from 'bun:test';
import {
  encodeEnvelope,
  decodeEnvelope,
  encodeRatchetMessage,
  encodeStreamChunk,
  decodeStreamChunk,
  inspectEnvelopeType,
  TYPE_STREAM_CHUNK,
} from '../src/index.js';
import type { StreamChunkWire } from '../src/index.js';
import type { PreKeyMessage, RatchetMessage, ShadeEnvelope } from '@shade/core';

function randBytes(n: number): Uint8Array {
  const buf = new Uint8Array(n);
  crypto.getRandomValues(buf);
  return buf;
}

function makeRatchetMessage(): RatchetMessage {
  return {
    dhPublicKey: randBytes(32),
    previousCounter: 42,
    counter: 7,
    ciphertext: randBytes(64),
    nonce: randBytes(12),
  };
}

function makePreKeyMessage(): PreKeyMessage {
  return {
    registrationId: 12345,
    preKeyId: 100,
    signedPreKeyId: 1,
    ephemeralKey: randBytes(32),
    identityDHKey: randBytes(32),
    message: makeRatchetMessage(),
  };
}

describe('Wire Format', () => {
  // ─── RatchetMessage ────────────────────────────────────────

  describe('RatchetMessage', () => {
    test('encode/decode roundtrip', () => {
      const msg = makeRatchetMessage();
      const envelope: ShadeEnvelope = {
        type: 'ratchet',
        content: msg,
        timestamp: Date.now(),
        senderAddress: 'alice',
      };

      const encoded = encodeEnvelope(envelope);
      const decoded = decodeEnvelope(encoded);

      expect(decoded.type).toBe('ratchet');
      const rm = decoded.content as RatchetMessage;
      expect(rm.dhPublicKey).toEqual(msg.dhPublicKey);
      expect(rm.previousCounter).toBe(42);
      expect(rm.counter).toBe(7);
      expect(rm.ciphertext).toEqual(msg.ciphertext);
      expect(rm.nonce).toEqual(msg.nonce);
    });

    test('compact size (smaller than JSON)', () => {
      const msg = makeRatchetMessage();
      const envelope: ShadeEnvelope = {
        type: 'ratchet',
        content: msg,
        timestamp: 0,
        senderAddress: '',
      };

      const binary = encodeEnvelope(envelope);
      const json = new TextEncoder().encode(JSON.stringify(envelope));

      // Binary should be significantly smaller
      expect(binary.length).toBeLessThan(json.length);
    });
  });

  // ─── PreKeyMessage ─────────────────────────────────────────

  describe('PreKeyMessage', () => {
    test('encode/decode roundtrip with preKeyId', () => {
      const msg = makePreKeyMessage();
      const envelope: ShadeEnvelope = {
        type: 'prekey',
        content: msg,
        timestamp: Date.now(),
        senderAddress: 'alice',
      };

      const encoded = encodeEnvelope(envelope);
      const decoded = decodeEnvelope(encoded);

      expect(decoded.type).toBe('prekey');
      const pm = decoded.content as PreKeyMessage;
      expect(pm.registrationId).toBe(12345);
      expect(pm.preKeyId).toBe(100);
      expect(pm.signedPreKeyId).toBe(1);
      expect(pm.ephemeralKey).toEqual(msg.ephemeralKey);
      expect(pm.identityDHKey).toEqual(msg.identityDHKey);

      // Nested ratchet message
      expect(pm.message.dhPublicKey).toEqual(msg.message.dhPublicKey);
      expect(pm.message.counter).toBe(msg.message.counter);
      expect(pm.message.ciphertext).toEqual(msg.message.ciphertext);
      expect(pm.message.nonce).toEqual(msg.message.nonce);
    });

    test('encode/decode roundtrip without preKeyId', () => {
      const msg = makePreKeyMessage();
      msg.preKeyId = undefined;

      const envelope: ShadeEnvelope = {
        type: 'prekey',
        content: msg,
        timestamp: 0,
        senderAddress: '',
      };

      const encoded = encodeEnvelope(envelope);
      const decoded = decodeEnvelope(encoded);

      const pm = decoded.content as PreKeyMessage;
      expect(pm.preKeyId).toBeUndefined();
    });
  });

  // ─── Edge Cases ────────────────────────────────────────────

  describe('edge cases', () => {
    test('empty ciphertext', () => {
      const msg: RatchetMessage = {
        dhPublicKey: randBytes(32),
        previousCounter: 0,
        counter: 0,
        ciphertext: new Uint8Array(0),
        nonce: randBytes(12),
      };

      const encoded = encodeRatchetMessage(msg);
      const decoded = decodeEnvelope(encoded);
      expect((decoded.content as RatchetMessage).ciphertext.length).toBe(0);
    });

    test('large ciphertext (10KB)', () => {
      const msg: RatchetMessage = {
        dhPublicKey: randBytes(32),
        previousCounter: 0,
        counter: 0,
        ciphertext: randBytes(10240),
        nonce: randBytes(12),
      };

      const encoded = encodeRatchetMessage(msg);
      const decoded = decodeEnvelope(encoded);
      expect((decoded.content as RatchetMessage).ciphertext).toEqual(msg.ciphertext);
    });

    test('max counter values', () => {
      const msg: RatchetMessage = {
        dhPublicKey: randBytes(32),
        previousCounter: 0xFFFFFFFF - 1,
        counter: 0xFFFFFFFF - 1,
        ciphertext: randBytes(16),
        nonce: randBytes(12),
      };

      const encoded = encodeRatchetMessage(msg);
      const decoded = decodeEnvelope(encoded);
      const rm = decoded.content as RatchetMessage;
      expect(rm.previousCounter).toBe(0xFFFFFFFF - 1);
      expect(rm.counter).toBe(0xFFFFFFFF - 1);
    });

    test('rejects unknown version', () => {
      const data = new Uint8Array([0xFF, 0x01]);
      expect(() => decodeEnvelope(data)).toThrow('Unknown version');
    });

    test('rejects unknown type', () => {
      const data = new Uint8Array([0x02, 0xFF]);
      expect(() => decodeEnvelope(data)).toThrow('Unknown type');
    });

    test('rejects too-short data', () => {
      expect(() => decodeEnvelope(new Uint8Array([0x02]))).toThrow('Too short');
      expect(() => decodeEnvelope(new Uint8Array([]))).toThrow('Too short');
    });
  });

  // ─── StreamChunk (type 0x11) ───────────────────────────────

  describe('StreamChunk', () => {
    function makeChunk(overrides: Partial<StreamChunkWire> = {}): StreamChunkWire {
      return {
        streamId: randBytes(16),
        laneId: 0,
        seq: 0n,
        isLast: false,
        nonce: randBytes(12),
        aad: new Uint8Array(0),
        ciphertext: randBytes(64),
        ...overrides,
      };
    }

    test('encode/decode roundtrip', () => {
      const c = makeChunk({ laneId: 7, seq: 42n, isLast: true });
      const encoded = encodeStreamChunk(c);
      const decoded = decodeStreamChunk(encoded);
      expect(decoded.streamId).toEqual(c.streamId);
      expect(decoded.laneId).toBe(7);
      expect(decoded.seq).toBe(42n);
      expect(decoded.isLast).toBe(true);
      expect(decoded.nonce).toEqual(c.nonce);
      expect(decoded.aad.length).toBe(0);
      expect(decoded.ciphertext).toEqual(c.ciphertext);
    });

    test('emits the correct type tag', () => {
      const encoded = encodeStreamChunk(makeChunk());
      expect(encoded[0]).toBe(0x02); // version
      expect(encoded[1]).toBe(TYPE_STREAM_CHUNK);
    });

    test('handles empty ciphertext', () => {
      const c = makeChunk({ ciphertext: new Uint8Array(0) });
      const decoded = decodeStreamChunk(encodeStreamChunk(c));
      expect(decoded.ciphertext.length).toBe(0);
    });

    test('handles 16 MiB ciphertext (max chunk + tag)', () => {
      const c = makeChunk({ ciphertext: randBytes(16 * 1024 * 1024 + 16) });
      const encoded = encodeStreamChunk(c);
      const decoded = decodeStreamChunk(encoded);
      expect(decoded.ciphertext.length).toBe(c.ciphertext.length);
    });

    test('handles MAX_SEQ', () => {
      const max = 0xffff_ffff_ffff_ffffn;
      const c = makeChunk({ seq: max });
      const decoded = decodeStreamChunk(encodeStreamChunk(c));
      expect(decoded.seq).toBe(max);
    });

    test('rejects wrong-length streamId', () => {
      expect(() => encodeStreamChunk(makeChunk({ streamId: randBytes(15) }))).toThrow();
    });

    test('rejects wrong-length nonce', () => {
      expect(() => encodeStreamChunk(makeChunk({ nonce: randBytes(11) }))).toThrow();
    });

    test('rejects out-of-range laneId', () => {
      expect(() => encodeStreamChunk(makeChunk({ laneId: 0x1_0000_0000 }))).toThrow();
      expect(() => encodeStreamChunk(makeChunk({ laneId: -1 }))).toThrow();
    });

    test('decode rejects wrong type tag', () => {
      const valid = encodeStreamChunk(makeChunk());
      const tampered = valid.slice();
      tampered[1] = 0x02; // ratchet
      expect(() => decodeStreamChunk(tampered)).toThrow();
    });

    test('decode rejects wrong version', () => {
      const valid = encodeStreamChunk(makeChunk());
      const tampered = valid.slice();
      tampered[0] = 0x01; // bumped to 0x02 in v0.3.0; 0x01 is the legacy version
      expect(() => decodeStreamChunk(tampered)).toThrow();
    });

    test('decode rejects truncated body', () => {
      const valid = encodeStreamChunk(makeChunk());
      expect(() => decodeStreamChunk(valid.slice(0, valid.length - 5))).toThrow();
    });

    test('decode rejects body extended beyond ctLen', () => {
      const valid = encodeStreamChunk(makeChunk());
      const tampered = new Uint8Array(valid.length + 1);
      tampered.set(valid);
      expect(() => decodeStreamChunk(tampered)).toThrow();
    });

    test('inspectEnvelopeType identifies stream-chunk', () => {
      expect(inspectEnvelopeType(encodeStreamChunk(makeChunk()))).toBe('stream-chunk');
    });

    test('inspectEnvelopeType identifies ratchet/prekey/unknown', () => {
      const ratchet: ShadeEnvelope = {
        type: 'ratchet',
        content: makeRatchetMessage(),
        timestamp: 0,
        senderAddress: '',
      };
      expect(inspectEnvelopeType(encodeEnvelope(ratchet))).toBe('ratchet');
      expect(inspectEnvelopeType(new Uint8Array([0x02, 0x99]))).toBe('unknown');
      expect(inspectEnvelopeType(new Uint8Array([]))).toBe('unknown');
    });
  });
});