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
 *   0x11 = StreamChunk
 *
 * All multi-byte integers are big-endian.
 *
 * Length prefixes are 4-byte (u32) since wire VERSION 0x02. (VERSION 0x01
 * used 2-byte/u16 length prefixes which silently truncated payloads larger
 * than 64 KiB — a hard correctness ceiling that blocked inline file ops
 * up to 256 KiB. The bump is incompatible with 0.2.x peers.)
 */

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

const VERSION = 0x02;

const TYPE_PREKEY = 0x01;
const TYPE_RATCHET = 0x02;
export const TYPE_STREAM_CHUNK = 0x11;

// ─── Stream chunk types ──────────────────────────────────────

/**
 * Wire-decoded stream-chunk envelope (type 0x11).
 *
 * See spec §2.2. The nonce is deterministic — derived from
 * (laneId, seq) on both sides — but is also serialized over the wire for
 * self-description and validated by the receiver.
 */
export interface StreamChunkWire {
  streamId: Uint8Array;     // 16 bytes
  laneId: number;           // u32
  seq: number | bigint;     // u64
  isLast: boolean;
  nonce: Uint8Array;        // 12 bytes
  aad: Uint8Array;          // additional bound data (length=0 in v0.2.0, reserved)
  ciphertext: Uint8Array;   // AES-GCM ciphertext including 16-byte tag
}

const STREAM_ID_BYTES = 16;
const STREAM_NONCE_BYTES = 12;

// ─── 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);
}

/**
 * Encode a stream-chunk envelope to wire bytes (type 0x11).
 *
 * Layout: see `shade-streams-spec.md` §2.2.
 */
export function encodeStreamChunk(c: StreamChunkWire): Uint8Array {
  if (c.streamId.length !== STREAM_ID_BYTES) {
    throw new Error(`streamId must be ${STREAM_ID_BYTES} bytes`);
  }
  if (c.nonce.length !== STREAM_NONCE_BYTES) {
    throw new Error(`nonce must be ${STREAM_NONCE_BYTES} bytes`);
  }
  if (!Number.isInteger(c.laneId) || c.laneId < 0 || c.laneId > 0xffff_ffff) {
    throw new Error(`laneId out of u32 range: ${c.laneId}`);
  }
  const seqBig = typeof c.seq === 'bigint' ? c.seq : BigInt(c.seq);
  if (seqBig < 0n || seqBig > 0xffff_ffff_ffff_ffffn) {
    throw new Error(`seq out of u64 range: ${c.seq}`);
  }

  const headerSize = 1 + 1 + STREAM_ID_BYTES + 4 + 8 + 1 + STREAM_NONCE_BYTES + 4 + c.aad.length + 4;
  const out = new Uint8Array(headerSize + c.ciphertext.length);
  const view = new DataView(out.buffer);
  let offset = 0;

  out[offset++] = VERSION;
  out[offset++] = TYPE_STREAM_CHUNK;
  out.set(c.streamId, offset);
  offset += STREAM_ID_BYTES;

  view.setUint32(offset, c.laneId, false);
  offset += 4;

  view.setBigUint64(offset, seqBig, false);
  offset += 8;

  out[offset++] = c.isLast ? 0x01 : 0x00;

  out.set(c.nonce, offset);
  offset += STREAM_NONCE_BYTES;

  view.setUint32(offset, c.aad.length, false);
  offset += 4;
  out.set(c.aad, offset);
  offset += c.aad.length;

  view.setUint32(offset, c.ciphertext.length, false);
  offset += 4;
  out.set(c.ciphertext, offset);

  return out;
}

// ─── 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}`);
}

/**
 * Decode a stream-chunk envelope from wire bytes (type 0x11).
 * Throws if the data is malformed or the type tag is wrong.
 */
export function decodeStreamChunk(data: Uint8Array): StreamChunkWire {
  const minHeaderSize = 2 + STREAM_ID_BYTES + 4 + 8 + 1 + STREAM_NONCE_BYTES + 4 + 4;
  if (data.length < minHeaderSize) {
    throw new Error(`stream-chunk too short: ${data.length} < ${minHeaderSize}`);
  }
  if (data[0] !== VERSION) {
    throw new Error(`Unknown version: ${data[0]}`);
  }
  if (data[1] !== TYPE_STREAM_CHUNK) {
    throw new Error(`Not a stream-chunk: type=${data[1]}`);
  }

  const view = new DataView(data.buffer, data.byteOffset);
  let offset = 2;

  const streamId = data.slice(offset, offset + STREAM_ID_BYTES);
  offset += STREAM_ID_BYTES;

  const laneId = view.getUint32(offset, false);
  offset += 4;

  const seq = view.getBigUint64(offset, false);
  offset += 8;

  const isLast = data[offset] === 0x01;
  offset += 1;

  const nonce = data.slice(offset, offset + STREAM_NONCE_BYTES);
  offset += STREAM_NONCE_BYTES;

  const aadLen = view.getUint32(offset, false);
  offset += 4;
  if (offset + aadLen + 4 > data.length) {
    throw new Error('stream-chunk truncated in aad/ctLen');
  }
  const aad = data.slice(offset, offset + aadLen);
  offset += aadLen;

  const ctLen = view.getUint32(offset, false);
  offset += 4;
  if (offset + ctLen !== data.length) {
    throw new Error(
      `stream-chunk length mismatch: declared ${offset + ctLen}, actual ${data.length}`,
    );
  }
  const ciphertext = data.slice(offset, offset + ctLen);

  return { streamId, laneId, seq, isLast, nonce, aad, ciphertext };
}

/**
 * Inspect the type tag of an arbitrary envelope without full parsing.
 * Returns one of `'prekey' | 'ratchet' | 'stream-chunk' | 'unknown'`.
 */
export function inspectEnvelopeType(
  data: Uint8Array,
): 'prekey' | 'ratchet' | 'stream-chunk' | 'unknown' {
  if (data.length < 2 || data[0] !== VERSION) return 'unknown';
  switch (data[1]) {
    case TYPE_PREKEY:
      return 'prekey';
    case TYPE_RATCHET:
      return 'ratchet';
    case TYPE_STREAM_CHUNK:
      return 'stream-chunk';
    default:
      return 'unknown';
  }
}

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);

  const msg: PreKeyMessage = {
    registrationId,
    signedPreKeyId,
    ephemeralKey: ephemeral.value,
    identityDHKey: identityDH.value,
    message: ratchet.value,
  };
  if (preKeyId !== undefined) msg.preKeyId = preKeyId;
  return msg;
}

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(4);
  new DataView(len.buffer).setUint32(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).getUint32(0, false);
  const value = data.slice(offset + 4, offset + 4 + len);
  return { value, end: offset + 4 + 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;
}