WebDrop/packages/core/src/frame.ts

import { AAD_FRAME, Flags, HEADER_SIZE, ID_SIZE, MAX_FRAME_DATA, MAX_FRAME_PLAINTEXT, TAG_SIZE, TIERS } from './constants.js';
import { concat, readU32, utf8, writeU32, type Bytes } from './bytes.js';
import { IntegrityError } from './errors.js';

// ───────────────────────── padding(protocol §6) ─────────────────────────

/** 不小于 total 的最小档位。 */
export function tierFor(total: number): number {
  for (const t of TIERS) if (t >= total) return t;
  throw new RangeError('超出单帧上限');
}

/** 明文 = u32 data_len ‖ data ‖ 零填充,总长为档位。 */
export function padPlaintext(data: Uint8Array): Bytes {
  if (data.length > MAX_FRAME_DATA) throw new RangeError('单帧 data 最多 1,048,572 字节');
  const tier = tierFor(4 + data.length);
  const out = new Uint8Array(tier);
  out.set(writeU32(data.length), 0);
  out.set(data, 4);
  return out;
}

/** 接收方校验:长度字段不越界,总长是档位,填充全为零。 */
export function unpadPlaintext(plain: Uint8Array): Bytes {
  if (plain.length < 4 || !(TIERS as readonly number[]).includes(plain.length)) throw new IntegrityError('bad_padding');
  const len = readU32(plain, 0);
  if (len > plain.length - 4) throw new IntegrityError('bad_padding');
  for (let i = 4 + len; i < plain.length; i++) {
    if (plain[i] !== 0) throw new IntegrityError('bad_padding');
  }
  return plain.slice(4, 4 + len);
}

// ───────────────────────── 封装(protocol §4) ─────────────────────────

export interface Envelope {
  flags: number;
  itemId: Bytes;
  frameIndex: number;
  payloadLength: number;
  /** 前 26 字节。 */
  header: Bytes;
  payload: Bytes;
  /** 完整封装。 */
  bytes: Bytes;
}

/** 严格解析:版本、保留位、长度必须全部合法。与服务端 Envelope.TryParse 一致。 */
export function parseEnvelope(data: Uint8Array): Envelope {
  if (data.length < HEADER_SIZE || data[0] !== 1) throw new IntegrityError('bad_envelope');
  const flags = data[1] as number;
  if ((flags & 0xe0) !== 0) throw new IntegrityError('bad_envelope');
  const payloadLength = readU32(data, 22);
  if (data.length !== HEADER_SIZE + payloadLength) throw new IntegrityError('bad_envelope');

  const bytes = data.slice();
  return {
    flags,
    itemId: bytes.slice(2, 2 + ID_SIZE),
    frameIndex: readU32(bytes, 18),
    payloadLength,
    header: bytes.slice(0, HEADER_SIZE),
    payload: bytes.slice(HEADER_SIZE),
    bytes,
  };
}

export function encodeEnvelope(flags: number, itemId: Uint8Array, frameIndex: number, payload: Uint8Array): Bytes {
  if (itemId.length !== ID_SIZE) throw new RangeError('item_id 必须是 16 字节');
  const header = concat(new Uint8Array([1, flags]), itemId, writeU32(frameIndex), writeU32(payload.length));
  return concat(header, payload);
}

/** 流式解析时用:从缓冲区起点读出下一个封装的总长度,数据不够返回 null。 */
export function peekEnvelopeLength(buf: Uint8Array, maxPayload: number): number | null {
  if (buf.length < HEADER_SIZE) return null;
  if (buf[0] !== 1) throw new IntegrityError('bad_envelope');
  const payloadLength = readU32(buf, 22);
  if (payloadLength > maxPayload) throw new IntegrityError('bad_envelope');
  return HEADER_SIZE + payloadLength;
}

export const MAX_ENC_PAYLOAD = MAX_FRAME_PLAINTEXT + TAG_SIZE;

// ───────────────────────── 加解密(protocol §5) ─────────────────────────

/** IV = 8 字节零 ‖ frame_index(u32)。每个条目有独立的密钥,所以 frame_index 唯一即保证 IV 唯一。 */
export function frameIv(frameIndex: number): Bytes {
  const iv = new Uint8Array(12);
  iv.set(writeU32(frameIndex), 8);
  return iv;
}

/** AAD = "WD1" ‖ channel_id ‖ envelope[0..22),绑定频道、条目、序号和全部标志位。 */
export function frameAad(channelId: Uint8Array, header: Uint8Array): Bytes {
  return concat(utf8(AAD_FRAME), channelId, header.subarray(0, 22));
}

export interface SealParams {
  flags: number;
  itemId: Uint8Array;
  index: number;
  data: Uint8Array;
}

/**
 * 加密一帧并返回完整封装。
 * 注意:同一 (itemId, index) 只能加密一份固定明文,否则会造成 IV 重用。重试必须重发相同字节。
 */
export async function sealFrame(itemKey: CryptoKey, channelId: Uint8Array, p: SealParams): Promise<Bytes> {
  if ((p.flags & Flags.Enc) === 0) throw new RangeError('enc 帧必须带 ENC 标志');
  if ((p.flags & 0xe0) !== 0) throw new RangeError('保留标志位必须为 0');

  const plaintext = padPlaintext(p.data);
  const header = concat(new Uint8Array([1, p.flags]), p.itemId, writeU32(p.index), writeU32(plaintext.length + TAG_SIZE));

  const cipher = await globalThis.crypto.subtle.encrypt(
    { name: 'AES-GCM', iv: frameIv(p.index), additionalData: frameAad(channelId, header), tagLength: 128 },
    itemKey,
    plaintext,
  );
  return concat(header, new Uint8Array(cipher));
}

/** 认证并解密一帧,返回去掉 padding 的 data。任何失败都抛 IntegrityError。 */
export async function openFrame(itemKey: CryptoKey, channelId: Uint8Array, env: Envelope): Promise<Bytes> {
  if ((env.flags & Flags.Enc) === 0) throw new IntegrityError('wrong_mode');
  if (env.payloadLength < TAG_SIZE) throw new IntegrityError('bad_envelope');

  let plain: Bytes;
  try {
    plain = new Uint8Array(
      await globalThis.crypto.subtle.decrypt(
        { name: 'AES-GCM', iv: frameIv(env.frameIndex), additionalData: frameAad(channelId, env.header), tagLength: 128 },
        itemKey,
        env.payload,
      ),
    );
  } catch {
    throw new IntegrityError('auth_failed');
  }
  return unpadPlaintext(plain);
}