Shade/packages/shade-proto/src/wire.ts

/**
 * Shade Wire Format — compact binary encoding for protocol messages.
 *
 * Format: [version:1][type:1][payload...]
 *
 * Types:
 *   0x01 = PreKeyMessage
 *   0x02 = RatchetMessage
 *
 * All multi-byte integers are big-endian.
 * All byte arrays are length-prefixed (2-byte length + data).
 */

import type { PreKeyMessage, RatchetMessage, ShadeEnvelope } from '@shade/core';

const VERSION = 0x01;

const TYPE_PREKEY = 0x01;
const TYPE_RATCHET = 0x02;

// ─── Encode ──────────────────────────────────────────────────

export function encodeEnvelope(envelope: ShadeEnvelope): Uint8Array {
  if (envelope.type === 'prekey') {
    return encodePreKeyMessage(envelope.content as PreKeyMessage);
  }
  return encodeRatchetMessage(envelope.content as RatchetMessage);
}

export function encodePreKeyMessage(msg: PreKeyMessage): Uint8Array {
  const ratchetBytes = encodeRatchetMessageInner(msg.message);
  const parts: Uint8Array[] = [];

  // Header
  parts.push(new Uint8Array([VERSION, TYPE_PREKEY]));

  // registrationId (4 bytes)
  parts.push(uint32(msg.registrationId));

  // preKeyId (4 bytes, 0xFFFFFFFF = none)
  parts.push(uint32(msg.preKeyId ?? 0xFFFFFFFF));

  // signedPreKeyId (4 bytes)
  parts.push(uint32(msg.signedPreKeyId));

  // ephemeralKey (length-prefixed)
  parts.push(lpBytes(msg.ephemeralKey));

  // identityDHKey (length-prefixed)
  parts.push(lpBytes(msg.identityDHKey));

  // embedded ratchet message (length-prefixed)
  parts.push(lpBytes(ratchetBytes));

  return concat(parts);
}

export function encodeRatchetMessage(msg: RatchetMessage): Uint8Array {
  const parts: Uint8Array[] = [];
  parts.push(new Uint8Array([VERSION, TYPE_RATCHET]));
  parts.push(encodeRatchetMessageInner(msg));
  return concat(parts);
}

function encodeRatchetMessageInner(msg: RatchetMessage): Uint8Array {
  const parts: Uint8Array[] = [];
  parts.push(lpBytes(msg.dhPublicKey));
  parts.push(uint32(msg.previousCounter));
  parts.push(uint32(msg.counter));
  parts.push(lpBytes(msg.ciphertext));
  parts.push(lpBytes(msg.nonce));
  return concat(parts);
}

// ─── Decode ──────────────────────────────────────────────────

export function decodeEnvelope(data: Uint8Array): ShadeEnvelope {
  if (data.length < 2) throw new Error('Too short');
  const version = data[0];
  if (version !== VERSION) throw new Error(`Unknown version: ${version}`);

  const type = data[1];
  const payload = data.slice(2);

  if (type === TYPE_PREKEY) {
    const msg = decodePreKeyMessageInner(payload);
    return { type: 'prekey', content: msg, timestamp: 0, senderAddress: '' };
  }
  if (type === TYPE_RATCHET) {
    const msg = decodeRatchetMessageInner(payload, 0).value;
    return { type: 'ratchet', content: msg, timestamp: 0, senderAddress: '' };
  }
  throw new Error(`Unknown type: ${type}`);
}

function decodePreKeyMessageInner(data: Uint8Array): PreKeyMessage {
  let offset = 0;

  const registrationId = readUint32(data, offset); offset += 4;
  const preKeyIdRaw = readUint32(data, offset); offset += 4;
  const preKeyId = preKeyIdRaw === 0xFFFFFFFF ? undefined : preKeyIdRaw;
  const signedPreKeyId = readUint32(data, offset); offset += 4;

  const ephemeral = readLP(data, offset); offset = ephemeral.end;
  const identityDH = readLP(data, offset); offset = identityDH.end;
  const ratchetData = readLP(data, offset); offset = ratchetData.end;

  const ratchet = decodeRatchetMessageInner(ratchetData.value, 0);

  return {
    registrationId,
    preKeyId,
    signedPreKeyId,
    ephemeralKey: ephemeral.value,
    identityDHKey: identityDH.value,
    message: ratchet.value,
  };
}

function decodeRatchetMessageInner(data: Uint8Array, offset: number): { value: RatchetMessage; end: number } {
  const dhPub = readLP(data, offset); offset = dhPub.end;
  const prevCounter = readUint32(data, offset); offset += 4;
  const counter = readUint32(data, offset); offset += 4;
  const ciphertext = readLP(data, offset); offset = ciphertext.end;
  const nonce = readLP(data, offset); offset = nonce.end;

  return {
    value: {
      dhPublicKey: dhPub.value,
      previousCounter: prevCounter,
      counter,
      ciphertext: ciphertext.value,
      nonce: nonce.value,
    },
    end: offset,
  };
}

// ─── Helpers ─────────────────────────────────────────────────

function uint32(n: number): Uint8Array {
  const buf = new Uint8Array(4);
  new DataView(buf.buffer).setUint32(0, n, false);
  return buf;
}

function lpBytes(data: Uint8Array): Uint8Array {
  const len = new Uint8Array(2);
  new DataView(len.buffer).setUint16(0, data.length, false);
  return concat([len, data]);
}

function readUint32(data: Uint8Array, offset: number): number {
  return new DataView(data.buffer, data.byteOffset + offset).getUint32(0, false);
}

function readLP(data: Uint8Array, offset: number): { value: Uint8Array; end: number } {
  const len = new DataView(data.buffer, data.byteOffset + offset).getUint16(0, false);
  const value = data.slice(offset + 2, offset + 2 + len);
  return { value, end: offset + 2 + len };
}

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