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

import { describe, test, expect } from 'bun:test';
import { encodeEnvelope, decodeEnvelope, encodePreKeyMessage, encodeRatchetMessage } 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([0x01, 0xFF]);
      expect(() => decodeEnvelope(data)).toThrow('Unknown type');
    });

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