Shade/packages/shade-sdk/src/broadcast.ts

/**
 * V4.6 — Broadcast channels for one-to-many fan-out.
 *
 * Wraps the Signal-style sender-key primitives (`@shade/core/sender-keys`)
 * in a persistent, app-friendly handle. The crypto is unchanged from
 * `sender-keys.ts`; this module wires it into:
 *
 *  - `StorageProvider.{save,get,list,remove}BroadcastChannel` for state
 *  - `Shade.send` for distribution (sealed by the bilateral ratchet)
 *  - `@shade/proto.encodeBroadcast` for the on-wire broadcast envelope
 *
 * The model is **strictly one-to-many**: the channel owner is the only
 * sender. Members hold a tracking copy of the chain key (no signing
 * privkey) and can decrypt — they cannot send into the channel. This
 * matches the Prism use case (PC desktop fans output frames out to
 * paired peers); a future symmetric-group API would build on the same
 * persistence layer.
 */

import {
  type CryptoProvider,
  type StorageProvider,
  type BroadcastChannelRecord,
  type BroadcastMemberRecord,
  type GroupSession,
  type SenderKeyState,
  buildDistribution,
  createSenderKey,
  installDistribution,
  senderKeyEncrypt,
  senderKeyDecrypt,
  toBase64,
  fromBase64,
} from '@shade/core';
import {
  encodeBroadcast,
  decodeBroadcast,
  inspectEnvelopeType,
  type BroadcastWire,
} from '@shade/proto';
import type { ShadeEnvelope } from '@shade/core';

/**
 * Magic prefix used to embed broadcast control messages inside a regular
 * bilateral plaintext (so they ride the existing `shade.send` /
 * `shade.receive` path without a new ratchet wire type). The leading
 * NULs make a collision with user-app text astronomically unlikely.
 */
const CONTROL_MAGIC = 'SHADE-BC/v1:';

/** Discriminator used in `Shade.onMessage`'s third arg. */
export type MessageMeta =
  | { kind: 'direct' }
  | { kind: 'broadcast'; channelId: string; sender: string; generation: number; iteration: number };

/** Snapshot returned by `Shade.listBroadcastChannels()`. */
export interface BroadcastChannelSummary {
  id: string;
  ownerRole: 'sender' | 'receiver';
  ownerAddress: string;
  label?: string;
  generation: number;
  members: readonly string[];
}

// ─── Control-message framing ─────────────────────────────────

interface DistributionControl {
  kind: 'distribution';
  channelId: string;
  ownerAddress: string;
  label?: string;
  generation: number;
  chainKey: string;          // base64
  iteration: number;
  signingPublicKey: string;  // base64
}

interface RevocationControl {
  kind: 'revocation';
  channelId: string;
}

type ControlMessage = DistributionControl | RevocationControl;

export function isControlPlaintext(plaintext: string): boolean {
  return plaintext.startsWith(CONTROL_MAGIC);
}

function packControl(c: ControlMessage): string {
  return CONTROL_MAGIC + JSON.stringify(c);
}

function unpackControl(plaintext: string): ControlMessage {
  if (!plaintext.startsWith(CONTROL_MAGIC)) {
    throw new Error('not a broadcast control plaintext');
  }
  const json = plaintext.slice(CONTROL_MAGIC.length);
  return JSON.parse(json) as ControlMessage;
}

// ─── BroadcastChannel handle (sender-side) ───────────────────

/**
 * Sender-side handle for a broadcast channel. Returned by
 * `Shade.createBroadcastChannel()` and `Shade.getBroadcastChannel()` for
 * channels where this device is the owner.
 *
 * Receiving devices DON'T expose a `BroadcastChannel` handle — they get
 * a `meta.kind === 'broadcast'` callback into `onMessage` instead. There
 * is nothing for the receiver to do beyond decrypting; the SDK manages
 * the chain-key state under the hood.
 */
export interface BroadcastChannel {
  /** Stable, opaque channel identifier. */
  readonly id: string;
  /** Snapshot of currently active member addresses (excludes revoked). */
  members(): Promise<readonly string[]>;
  /**
   * Add `peerAddress` to the channel and return a control envelope to
   * deliver to that peer (sealed under the bilateral ratchet). The peer
   * must already have a verified bilateral session before this is called
   * — typically from the app's pair handler after `markPeerVerified`.
   */
  addMember(peerAddress: string): Promise<{ envelope: ShadeEnvelope }>;
  /**
   * Remove `peerAddress` and rotate the sender-key. Returns one envelope
   * per surviving member carrying the new chain key — deliver each to its
   * `to` address.
   */
  removeMember(peerAddress: string): Promise<{
    rotations: Array<{ to: string; envelope: ShadeEnvelope }>;
  }>;
  /**
   * Encrypt a plaintext once with the current sender chain. Returns a
   * single broadcast envelope (the same byte sequence is delivered to
   * every member by the caller's transport) plus the snapshot of active
   * members.
   */
  broadcast(plaintext: string | Uint8Array): Promise<{
    envelope: Uint8Array;
    members: readonly string[];
  }>;
}

/**
 * Internal hooks the SDK exposes to BroadcastChannel implementations.
 * Lets us avoid a circular import at the value level — the Shade class
 * stays the one place that owns sessions + delivery.
 */
export interface BroadcastSdkHooks {
  /** Bilateral-ratchet `send` for delivering control envelopes. */
  bilateralSend(peerAddress: string, plaintext: string): Promise<ShadeEnvelope>;
  /** This device's own address (= channel owner address for sender-side). */
  myAddress(): string;
  /** Crypto + storage handles (for chain-key advancement and persistence). */
  crypto: CryptoProvider;
  storage: StorageProvider;
}

class BroadcastChannelImpl implements BroadcastChannel {
  constructor(
    readonly id: string,
    private readonly hooks: BroadcastSdkHooks,
  ) {}

  async members(): Promise<readonly string[]> {
    const rows = await this.hooks.storage.getBroadcastMembers!(this.id);
    return rows.filter((m) => m.removedAt === null).map((m) => m.peerAddress);
  }

  async addMember(peerAddress: string): Promise<{ envelope: ShadeEnvelope }> {
    const channel = await loadChannel(this.hooks.storage, this.id, 'sender');
    const now = Date.now();
    const member: BroadcastMemberRecord = {
      channelId: this.id,
      peerAddress,
      joinedAt: now,
      removedAt: null,
    };
    await this.hooks.storage.saveBroadcastMember!(member);
    const control: DistributionControl = {
      kind: 'distribution',
      channelId: this.id,
      ownerAddress: channel.ownerAddress,
      generation: channel.generation,
      chainKey: toBase64(channel.chainKey),
      iteration: channel.iteration,
      signingPublicKey: toBase64(channel.signingPublicKey),
    };
    if (channel.label !== undefined) control.label = channel.label;
    const envelope = await this.hooks.bilateralSend(peerAddress, packControl(control));
    return { envelope };
  }

  async removeMember(peerAddress: string): Promise<{
    rotations: Array<{ to: string; envelope: ShadeEnvelope }>;
  }> {
    const channel = await loadChannel(this.hooks.storage, this.id, 'sender');
    const now = Date.now();

    // Mark target removed.
    await this.hooks.storage.saveBroadcastMember!({
      channelId: this.id,
      peerAddress,
      joinedAt: now,
      removedAt: now,
    });

    // Rotate: fresh chain + signing keypair, generation++.
    const fresh = await createSenderKey(this.hooks.crypto);
    if (channel.signingPrivateKey !== undefined) {
      this.hooks.crypto.zeroize(channel.signingPrivateKey);
    }
    this.hooks.crypto.zeroize(channel.chainKey);

    const rotated: BroadcastChannelRecord = {
      ...channel,
      generation: channel.generation + 1,
      chainKey: fresh.chainKey,
      iteration: 0,
      signingPublicKey: fresh.signingPublicKey,
      ...(fresh.signingPrivateKey !== undefined ? { signingPrivateKey: fresh.signingPrivateKey } : {}),
      updatedAt: now,
    };
    await this.hooks.storage.saveBroadcastChannel!(rotated);

    // Tell the revoked peer they're gone (best-effort; rotation already
    // means they can't decrypt new traffic, but the explicit signal lets
    // their UI reflect the change immediately).
    const revocation: RevocationControl = { kind: 'revocation', channelId: this.id };
    try {
      await this.hooks.bilateralSend(peerAddress, packControl(revocation));
    } catch {
      // ignore — the revoked peer might be offline; rotation has already
      // happened locally so they're cryptographically out either way.
    }

    // Distribute new chain to surviving members.
    const survivors = (await this.members()).filter((m) => m !== peerAddress);
    const rotations: Array<{ to: string; envelope: ShadeEnvelope }> = [];
    for (const peer of survivors) {
      const control: DistributionControl = {
        kind: 'distribution',
        channelId: this.id,
        ownerAddress: rotated.ownerAddress,
        generation: rotated.generation,
        chainKey: toBase64(rotated.chainKey),
        iteration: rotated.iteration,
        signingPublicKey: toBase64(rotated.signingPublicKey),
      };
      if (rotated.label !== undefined) control.label = rotated.label;
      const envelope = await this.hooks.bilateralSend(peer, packControl(control));
      rotations.push({ to: peer, envelope });
    }
    return { rotations };
  }

  async broadcast(plaintext: string | Uint8Array): Promise<{
    envelope: Uint8Array;
    members: readonly string[];
  }> {
    const channel = await loadChannel(this.hooks.storage, this.id, 'sender');
    const session = recordToGroupSession(channel);

    const ptBytes = typeof plaintext === 'string'
      ? new TextEncoder().encode(plaintext)
      : plaintext;

    const msg = await senderKeyEncrypt(
      this.hooks.crypto,
      session,
      channel.ownerAddress,
      ptBytes,
    );

    // Persist updated chain state (the sender-key encrypt mutated the
    // SenderKeyState in `session.senderKeys`).
    const advanced = session.senderKeys.get(channel.ownerAddress)!;
    const updated: BroadcastChannelRecord = {
      ...channel,
      chainKey: advanced.chainKey,
      iteration: advanced.iteration,
      updatedAt: Date.now(),
    };
    await this.hooks.storage.saveBroadcastChannel!(updated);

    const wire: BroadcastWire = {
      channelId: this.id,
      senderAddress: channel.ownerAddress,
      generation: channel.generation,
      iteration: msg.iteration,
      nonce: msg.nonce,
      signature: msg.signature,
      ciphertext: msg.ciphertext,
    };
    return {
      envelope: encodeBroadcast(wire),
      members: await this.members(),
    };
  }
}

// ─── Receiver-side: handle an incoming broadcast envelope ────

/**
 * Decrypt an inbound broadcast wire envelope. Returns `{ plaintext, meta }`
 * the SDK then dispatches to `onMessage` handlers. Throws if no matching
 * channel exists or the ciphertext fails verification.
 *
 * Stale generations (older than the receiver's currently-installed
 * sender-key) are silently dropped: returns `null`.
 */
export async function acceptBroadcastEnvelope(
  hooks: BroadcastSdkHooks,
  envelopeBytes: Uint8Array,
): Promise<{ plaintext: string; meta: MessageMeta } | null> {
  const kind = inspectEnvelopeType(envelopeBytes);
  if (kind !== 'broadcast') {
    throw new Error(`acceptBroadcastEnvelope: not a broadcast (type=${kind})`);
  }
  const wire = decodeBroadcast(envelopeBytes);
  const channel = await hooks.storage.getBroadcastChannel?.(wire.channelId);
  if (!channel) {
    throw new Error(`unknown broadcast channel: ${wire.channelId}`);
  }
  // Only act on broadcasts from the recorded owner.
  if (channel.ownerAddress !== wire.senderAddress) {
    throw new Error(
      `broadcast sender ${wire.senderAddress} does not match channel owner ${channel.ownerAddress}`,
    );
  }
  if (wire.generation < channel.generation) {
    // Stale: silently drop (sender's old chain post-rotation).
    return null;
  }
  if (wire.generation > channel.generation) {
    throw new Error(
      `broadcast generation ${wire.generation} ahead of installed ${channel.generation} — missing distribution`,
    );
  }

  const session = recordToGroupSession(channel);
  const plaintextBytes = await senderKeyDecrypt(hooks.crypto, session, channel.channelId, {
    senderAddress: channel.ownerAddress,
    iteration: wire.iteration,
    ciphertext: wire.ciphertext,
    nonce: wire.nonce,
    signature: wire.signature,
  });
  // Persist advanced chain state.
  const advanced = session.senderKeys.get(channel.ownerAddress)!;
  await hooks.storage.saveBroadcastChannel!({
    ...channel,
    chainKey: advanced.chainKey,
    iteration: advanced.iteration,
    updatedAt: Date.now(),
  });
  const plaintext = new TextDecoder().decode(plaintextBytes);
  return {
    plaintext,
    meta: {
      kind: 'broadcast',
      channelId: wire.channelId,
      sender: wire.senderAddress,
      generation: wire.generation,
      iteration: wire.iteration,
    },
  };
}

/**
 * Detect whether an inbound bilateral plaintext is a broadcast control
 * message (distribution / revocation) and consume it. Returns `true`
 * when the plaintext was consumed (caller MUST NOT dispatch it to
 * onMessage handlers); `false` for normal direct plaintexts.
 */
export async function maybeHandleControlPlaintext(
  hooks: BroadcastSdkHooks,
  from: string,
  plaintext: string,
): Promise<boolean> {
  if (!isControlPlaintext(plaintext)) return false;
  const ctrl = unpackControl(plaintext);
  if (ctrl.kind === 'distribution') {
    await installIncomingDistribution(hooks, from, ctrl);
  } else if (ctrl.kind === 'revocation') {
    await applyRevocation(hooks, ctrl.channelId);
  }
  return true;
}

async function installIncomingDistribution(
  hooks: BroadcastSdkHooks,
  from: string,
  d: DistributionControl,
): Promise<void> {
  // The control message arrived over a bilateral session with `from`.
  // For trust, require that the channel owner address matches the
  // bilateral peer that delivered it — otherwise a paired peer could
  // smuggle a fake distribution claiming to come from another address.
  if (d.ownerAddress !== from) {
    throw new Error(
      `broadcast distribution owner ${d.ownerAddress} does not match bilateral sender ${from}`,
    );
  }
  const existing = await hooks.storage.getBroadcastChannel?.(d.channelId);
  if (existing && existing.generation > d.generation) {
    // Stale distribution — keep what we have.
    return;
  }
  const now = Date.now();
  const record: BroadcastChannelRecord = {
    channelId: d.channelId,
    ownerRole: 'receiver',
    ownerAddress: d.ownerAddress,
    generation: d.generation,
    chainKey: fromBase64(d.chainKey),
    iteration: d.iteration,
    signingPublicKey: fromBase64(d.signingPublicKey),
    createdAt: existing?.createdAt ?? now,
    updatedAt: now,
  };
  if (d.label !== undefined) record.label = d.label;
  await hooks.storage.saveBroadcastChannel!(record);
}

async function applyRevocation(
  hooks: BroadcastSdkHooks,
  channelId: string,
): Promise<void> {
  // Drop our local copy entirely. The owner has rotated; without a fresh
  // distribution we couldn't decrypt new traffic anyway.
  await hooks.storage.removeBroadcastChannel?.(channelId);
}

// ─── Channel creation + summaries ────────────────────────────

export async function createBroadcastChannelImpl(
  hooks: BroadcastSdkHooks,
  opts: { label?: string } = {},
): Promise<BroadcastChannel> {
  ensureBroadcastSupport(hooks.storage);
  const channelId = generateChannelId(hooks.crypto);
  const seed = await createSenderKey(hooks.crypto);
  const now = Date.now();
  const record: BroadcastChannelRecord = {
    channelId,
    ownerRole: 'sender',
    ownerAddress: hooks.myAddress(),
    generation: 0,
    chainKey: seed.chainKey,
    iteration: seed.iteration,
    signingPublicKey: seed.signingPublicKey,
    ...(seed.signingPrivateKey !== undefined ? { signingPrivateKey: seed.signingPrivateKey } : {}),
    createdAt: now,
    updatedAt: now,
  };
  if (opts.label !== undefined) record.label = opts.label;
  await hooks.storage.saveBroadcastChannel!(record);
  return new BroadcastChannelImpl(channelId, hooks);
}

export async function getBroadcastChannelImpl(
  hooks: BroadcastSdkHooks,
  channelId: string,
): Promise<BroadcastChannel | null> {
  ensureBroadcastSupport(hooks.storage);
  const record = await hooks.storage.getBroadcastChannel!(channelId);
  if (!record || record.ownerRole !== 'sender') return null;
  return new BroadcastChannelImpl(channelId, hooks);
}

export async function listBroadcastChannelsImpl(
  hooks: BroadcastSdkHooks,
): Promise<readonly BroadcastChannelSummary[]> {
  if (hooks.storage.listBroadcastChannels === undefined) return [];
  const records = await hooks.storage.listBroadcastChannels();
  const out: BroadcastChannelSummary[] = [];
  for (const r of records) {
    const memberRows = (await hooks.storage.getBroadcastMembers?.(r.channelId)) ?? [];
    const summary: BroadcastChannelSummary = {
      id: r.channelId,
      ownerRole: r.ownerRole,
      ownerAddress: r.ownerAddress,
      generation: r.generation,
      members: memberRows.filter((m) => m.removedAt === null).map((m) => m.peerAddress),
    };
    if (r.label !== undefined) summary.label = r.label;
    out.push(summary);
  }
  return out;
}

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

function ensureBroadcastSupport(storage: StorageProvider): void {
  if (
    storage.saveBroadcastChannel === undefined ||
    storage.getBroadcastChannel === undefined ||
    storage.saveBroadcastMember === undefined ||
    storage.getBroadcastMembers === undefined
  ) {
    throw new Error(
      'StorageProvider does not implement broadcast-channel APIs. Use a backend ≥ Shade 4.6 (memory/sqlite/indexeddb/postgres + their encrypted variants).',
    );
  }
}

async function loadChannel(
  storage: StorageProvider,
  channelId: string,
  expectRole: 'sender' | 'receiver',
): Promise<BroadcastChannelRecord> {
  const c = await storage.getBroadcastChannel?.(channelId);
  if (!c) throw new Error(`broadcast channel ${channelId} not found`);
  if (c.ownerRole !== expectRole) {
    throw new Error(
      `broadcast channel ${channelId} role mismatch: have ${c.ownerRole}, expected ${expectRole}`,
    );
  }
  return c;
}

function recordToGroupSession(record: BroadcastChannelRecord): GroupSession {
  // Reconstruct a minimal GroupSession with one entry (the channel owner).
  const state: SenderKeyState = {
    chainKey: new Uint8Array(record.chainKey),
    iteration: record.iteration,
    signingPublicKey: new Uint8Array(record.signingPublicKey),
  };
  if (record.signingPrivateKey !== undefined) {
    state.signingPrivateKey = new Uint8Array(record.signingPrivateKey);
  }
  return {
    groupId: record.channelId,
    senderKeys: new Map([[record.ownerAddress, state]]),
  };
}

function generateChannelId(crypto: CryptoProvider): string {
  // 16 bytes = 22-char base64url. Stable across restarts because we
  // persist it before returning.
  const bytes = crypto.randomBytes(16);
  return bytesToBase64Url(bytes);
}

function bytesToBase64Url(bytes: Uint8Array): string {
  let bin = '';
  for (let i = 0; i < bytes.length; i++) bin += String.fromCharCode(bytes[i]!);
  return btoa(bin).replace(/\+/g, '-').replace(/\//g, '_').replace(/=+$/, '');
}

// Re-export `buildDistribution` + `installDistribution` for tests + advanced
// callers — internal implementations use them already.
export { buildDistribution, installDistribution };