WebDrop/packages/core/src/stream.ts

import { Flags } from './constants.js';
import { equalBytes, type Bytes } from './bytes.js';
import { IntegrityError } from './errors.js';
import { MAX_ENC_PAYLOAD, openFrame, parseEnvelope, peekEnvelopeLength } from './frame.js';
import { decodeFirstFrame, type ItemMeta } from './item.js';
import { sha256 } from './kdf.js';
import { frameManifestHash } from './signatures.js';
import type { ReaderContext } from './reader.js';

/** 字节队列:接收任意切分的块,按需取出前 n 个字节。 */
class ByteQueue {
  readonly #parts: Uint8Array[] = [];
  #length = 0;

  get length(): number {
    return this.#length;
  }

  push(chunk: Uint8Array): void {
    if (chunk.length === 0) return;
    this.#parts.push(chunk);
    this.#length += chunk.length;
  }

  /** 复制前 n 个字节,不消耗。 */
  peek(n: number): Bytes {
    const out = new Uint8Array(n);
    let filled = 0;
    for (const part of this.#parts) {
      if (filled >= n) break;
      const take = Math.min(part.length, n - filled);
      out.set(part.subarray(0, take), filled);
      filled += take;
    }
    return out;
  }

  take(n: number): Bytes {
    const out = new Uint8Array(n);
    let filled = 0;
    while (filled < n) {
      const head = this.#parts[0] as Uint8Array;
      const need = n - filled;
      if (head.length <= need) {
        out.set(head, filled);
        filled += head.length;
        this.#parts.shift();
      } else {
        out.set(head.subarray(0, need), filled);
        this.#parts[0] = head.subarray(need);
        filled += need;
      }
    }
    this.#length -= n;
    return out;
  }
}

async function* bytesOf(source: ReadableStream<Uint8Array> | AsyncIterable<Uint8Array>): AsyncGenerator<Uint8Array, void, void> {
  if (typeof ReadableStream !== 'undefined' && source instanceof ReadableStream) {
    const reader = source.getReader();
    let finished = false;
    try {
      for (;;) {
        const { done, value } = await reader.read();
        if (done) { finished = true; return; }
        yield value;
      }
    } finally {
      if (!finished) await reader.cancel().catch(() => {});
      reader.releaseLock();
    }
  }
  yield* source as AsyncIterable<Uint8Array>;
}

/**
 * 把「封装首尾相接」的字节流(GET /items/{id}/frames 的响应体)切成一个个完整封装。
 * 连接中途断开会在末尾留下半个封装:这是正常的网络情况,不是篡改,所以直接丢弃,
 * 调用方根据 decryptItemStream 报告的 nextIndex 续传即可。
 */
export async function* splitEnvelopes(
  source: ReadableStream<Uint8Array> | AsyncIterable<Uint8Array>,
): AsyncGenerator<Bytes, void, void> {
  const queue = new ByteQueue();
  for await (const chunk of bytesOf(source)) {
    queue.push(chunk);
    for (;;) {
      if (queue.length < 26) break;
      const total = peekEnvelopeLength(queue.peek(26), MAX_ENC_PAYLOAD);
      if (total === null || queue.length < total) break;
      yield queue.take(total);
    }
  }
}

export interface DecryptStreamOptions {
  /** 这次响应从哪一帧开始(拉取时的 from 参数),默认 0。 */
  fromIndex?: number;
  /** 期望的条目 id(base64url 解码后的字节)。提供后,混入别的条目的帧会被拒绝。 */
  itemId?: Uint8Array;
}

export interface ItemDecodeResult {
  /** 已收到已认证的 FINAL。 */
  complete: boolean;
  /** 下一个期望的帧序号。没收完时用它作为 from 续传(不会重复计下载次数)。 */
  nextIndex: number;
  /** 第 0 帧里的元数据;从中间帧续传时没有。 */
  meta: ItemMeta | undefined;
  /** Present only for complete reads starting at frame zero. */
  ciphertextManifestHash?: Bytes;
}

/**
 * 流式解密一个条目:边收边解边产出内容,内存占用只有一帧。可以放进 Worker 里运行。
 *
 * 逐帧校验:认证、条目 id、帧序号连续、标志一致、FINAL 之后不再有帧。
 * 任何一项失败都会抛 IntegrityError(流会出错),已产出的内容不可信的部分由调用方丢弃。
 * 流正常结束但没有 FINAL 不是错误(连接可能被断开),此时 complete 为 false。
 */
export async function* decryptItemFrames(
  ctx: ReaderContext,
  source: ReadableStream<Uint8Array> | AsyncIterable<Uint8Array>,
  opts: DecryptStreamOptions = {},
): AsyncGenerator<Bytes, ItemDecodeResult, void> {
  let expected = opts.fromIndex ?? 0;
  let itemId = opts.itemId;
  let flagsMask: number | undefined;
  let finalSeen = false;
  let meta: ItemMeta | undefined;
  const frameDigests: Bytes[] = [];

  for await (const bytes of splitEnvelopes(source)) {
    frameDigests.push(await sha256(bytes));
    const env = parseEnvelope(bytes);

    const isEnc = (env.flags & Flags.Enc) !== 0;
    if (isEnc !== (ctx.mode === 'enc')) throw new IntegrityError('wrong_mode');

    if (!itemId) itemId = env.itemId;
    else if (!equalBytes(itemId, env.itemId)) throw new IntegrityError('item_mismatch');

    if (finalSeen) throw new IntegrityError('frame_after_final');
    if (env.frameIndex !== expected) throw new IntegrityError('gap');

    const itemFlags = env.flags & (Flags.Gated | Flags.Snapshot | Flags.Control);
    if (flagsMask === undefined) flagsMask = itemFlags;
    else if (flagsMask !== itemFlags) throw new IntegrityError('flags_changed');

    let data: Bytes;
    if (ctx.mode === 'enc') {
      data = await openFrame(await ctx.itemKey(env.itemId), ctx.channelId, env);
    } else {
      data = env.payload;
    }
    if (env.frameIndex === 0) {
      const first = decodeFirstFrame(data); meta = first.meta; data = first.content;
    }

    expected++;
    if ((env.flags & Flags.Final) !== 0) finalSeen = true;
    if (data.length > 0) yield data;
  }

  return { complete: finalSeen, nextIndex: expected, meta,
    ...(finalSeen && (opts.fromIndex ?? 0) === 0 ? { ciphertextManifestHash: Uint8Array.from(await frameManifestHash(frameDigests)) } : {}) };
}

/** 把 decryptItemFrames 包成 ReadableStream,便于直接 pipe 到文件。 */
export function decryptItemStream(
  ctx: ReaderContext,
  source: ReadableStream<Uint8Array> | AsyncIterable<Uint8Array>,
  opts: DecryptStreamOptions = {},
): { content: ReadableStream<Uint8Array>; result: Promise<ItemDecodeResult> } {
  const gen = decryptItemFrames(ctx, source, opts);
  let resolve!: (r: ItemDecodeResult) => void;
  let reject!: (e: unknown) => void;
  const result = new Promise<ItemDecodeResult>((res, rej) => {
    resolve = res;
    reject = rej;
  });
  result.catch(() => {}); // 错误通过流本身传播,避免未处理的 rejection

  const content = new ReadableStream<Uint8Array>({
    async pull(controller) {
      try {
        const r = await gen.next();
        if (r.done) {
          resolve(r.value);
          controller.close();
        } else {
          controller.enqueue(r.value);
        }
      } catch (e) {
        reject(e);
        controller.error(e);
      }
    },
    async cancel(reason) {
      reject(reason ?? new DOMException('解密已取消', 'AbortError'));
      await gen.return({ complete: false, nextIndex: opts.fromIndex ?? 0, meta: undefined });
    },
  });
  return { content, result };
}