restore.test.ts

1/**2 * restore.test.ts — coverage for ChunkBuffer, parseRawScan, reconstructAndVerify.3 *4 * Tests build synthetic envelopes via the same helpers the generator uses5 * (Subagent 1's `helpers.ts` + `erasure.ts`). We avoid the full6 * `generateCodex` pipeline so these tests don't depend on PDF rendering7 * or Aztec encoding — they only exercise the restoration math.8 */910import { beforeEach, describe, expect, it } from "vitest";11import { deflate } from "pako";1213import {14  CodexError,15  deriveShardCounts,16  type ChunkEnvelope,17  type CodexPayload,18} from "../../shared/codex-types";19import {20  base64UrlToBytes,21  bytesToBase64Url,22  buildChunkEnvelope,23  sha256Hex,24} from "./helpers";25import { encodeWithErasure } from "./erasure";26import { _resetReedSolomonCacheForTest } from "./rs-loader";27import {28  ChunkBuffer,29  parseRawScan,30  reconstructAndVerify,31} from "./restore";3233// Reset between tests: the WASM Reed-Solomon instance carries internal state34// across encode/reconstruct calls that breaks subsequent reconstructions in35// the same process. Until that's fixed in erasure.ts / rs-loader.ts, isolating36// each test keeps the suite deterministic.37beforeEach(() => {38  _resetReedSolomonCacheForTest();39});4041const VERIFY_CODE = "PX-TEST-ABCD";4243/** Build envelopes for an arbitrary CodexPayload, mirroring the real generator. */44async function buildEnvelopesForPayload(45  payload: CodexPayload,46): Promise<ChunkEnvelope[]> {47  const json = JSON.stringify(payload);48  const compressed = deflate(new TextEncoder().encode(json), { level: 9 });49  // encodeWithErasure shards at CHUNK_SIZE internally — no need to pass it.50  const { dataChunks, parityChunks, totalChunks } =51    await encodeWithErasure(compressed);52  const all = [...dataChunks, ...parityChunks];53  return Promise.all(54    all.map((bytes, i) =>55      buildChunkEnvelope(56        bytes,57        i,58        totalChunks,59        i >= dataChunks.length,60        payload.verifyCode,61      ),62    ),63  );64}6566async function makePayload(seedLength = 12_000): Promise<CodexPayload> {67  // A vaultBlob big enough to span several chunks after deflate. We seed with68  // a deterministic but non-repeating pattern (xorshift-ish) so pako cannot69  // compress it to a single chunk — without that, dataShards collapses to 170  // and the "lose 2 chunks" test has nothing to reconstruct.71  const blobBytes = new Uint8Array(seedLength);72  let s = 0x9e3779b9;73  for (let i = 0; i < seedLength; i++) {74    s = (s ^ (s << 13)) >>> 0;75    s = (s ^ (s >>> 17)) >>> 0;76    s = (s ^ (s << 5)) >>> 0;77    blobBytes[i] = s & 0xff;78  }79  const vaultBlob = bytesToBase64Url(blobBytes);80  const vaultFingerprint = (81    await sha256Hex(new TextEncoder().encode(vaultBlob))82  )83    .slice(0, 16)84    .toUpperCase();85  return {86    version: 2,87    verifyCode: VERIFY_CODE,88    vaultBlob,89    timestamp: 1_700_000_000_000,90    entryCount: 42,91    vaultFingerprint,92    generatedAt: 1_700_000_001_000,93  };94}9596async function makeSingleEnvelope(opts: {97  index: number;98  total: number;99  isParity?: boolean;100  verifyCode?: string;101  bytes?: Uint8Array;102}): Promise<ChunkEnvelope> {103  return buildChunkEnvelope(104    opts.bytes ?? new Uint8Array([1, 2, 3, 4, 5]),105    opts.index,106    opts.total,107    opts.isParity ?? false,108    opts.verifyCode ?? VERIFY_CODE,109  );110}111112describe("ChunkBuffer", () => {113  it("accepts a valid chunk and exposes context", async () => {114    const buf = new ChunkBuffer();115    const env = await makeSingleEnvelope({ index: 0, total: 5 });116    const res = await buf.add(env);117    expect(res.accepted).toBe(true);118    expect(buf.size()).toBe(1);119    expect(buf.getContext()).toEqual({120      verifyCode: VERIFY_CODE,121      total: 5,122      dataShards: deriveShardCounts(5).dataShards,123    });124  });125126  it("rejects a chunk from a different verifyCode", async () => {127    const buf = new ChunkBuffer();128    await buf.add(await makeSingleEnvelope({ index: 0, total: 5 }));129    const foreign = await makeSingleEnvelope({130      index: 1,131      total: 5,132      verifyCode: "PX-OTHR-XXXX",133    });134    const res = await buf.add(foreign);135    expect(res).toEqual({ accepted: false, reason: "foreign_codex" });136    expect(buf.size()).toBe(1);137  });138139  it("rejects a chunk with bad integrity hash", async () => {140    const buf = new ChunkBuffer();141    const env = await makeSingleEnvelope({ index: 0, total: 5 });142    // Corrupt the data field without updating the hash — verifyChunkIntegrity143    // must catch this.144    const corrupted: ChunkEnvelope = {145      ...env,146      d: bytesToBase64Url(new Uint8Array([9, 9, 9, 9])),147    };148    const res = await buf.add(corrupted);149    expect(res).toEqual({ accepted: false, reason: "bad_integrity" });150    expect(buf.size()).toBe(0);151  });152153  it("rejects a chunk with wrong envelope version", async () => {154    const buf = new ChunkBuffer();155    const env = await makeSingleEnvelope({ index: 0, total: 5 });156    const wrongVersion = { ...env, v: 1 } as unknown as ChunkEnvelope;157    const res = await buf.add(wrongVersion);158    expect(res.accepted).toBe(false);159    expect(res.reason).toBe("parse_error");160  });161162  it("deduplicates: adding the same index twice keeps only one", async () => {163    const buf = new ChunkBuffer();164    const env = await makeSingleEnvelope({ index: 2, total: 5 });165    expect((await buf.add(env)).accepted).toBe(true);166    const second = await buf.add(env);167    expect(second).toEqual({ accepted: false, reason: "duplicate" });168    expect(buf.size()).toBe(1);169  });170171  it("isReady() flips true at dataShards count", async () => {172    const buf = new ChunkBuffer();173    const total = 5;174    const dataShards = deriveShardCounts(total).dataShards;175    for (let i = 0; i < dataShards; i++) {176      expect(buf.isReady()).toBe(false);177      const env = await makeSingleEnvelope({178        index: i,179        total,180        bytes: new Uint8Array([i, i, i]),181      });182      await buf.add(env);183    }184    expect(buf.isReady()).toBe(true);185  });186187  it("toChunkMap returns raw chunk bytes by index", async () => {188    const buf = new ChunkBuffer();189    const env = await makeSingleEnvelope({190      index: 0,191      total: 5,192      bytes: new Uint8Array([1, 2, 3]),193    });194    await buf.add(env);195    const map = buf.toChunkMap();196    expect(map.size).toBe(1);197    expect(Array.from(map.get(0)!)).toEqual([1, 2, 3]);198  });199});200201describe("parseRawScan", () => {202  it("parses a well-formed envelope", async () => {203    const env = await makeSingleEnvelope({ index: 0, total: 5 });204    const parsed = parseRawScan(JSON.stringify(env));205    expect(parsed).toEqual(env);206  });207208  it("returns parse_error for non-JSON input", () => {209    expect(parseRawScan("not json")).toEqual({ error: "parse_error" });210  });211212  it("returns wrong_version for v=1 envelopes", () => {213    const v1 = {214      v: 1,215      c: VERIFY_CODE,216      i: 0,217      t: 5,218      p: false,219      d: "AAAA",220      h: "ABCDEF12",221    };222    expect(parseRawScan(JSON.stringify(v1))).toEqual({223      error: "wrong_version",224    });225  });226227  it("returns parse_error for a v2 object with missing fields", () => {228    expect(parseRawScan(JSON.stringify({ v: 2, c: "PX" }))).toEqual({229      error: "parse_error",230    });231  });232});233234describe("reconstructAndVerify", () => {235  it("round-trips a full envelope set back to the original payload", async () => {236    const payload = await makePayload();237    const envelopes = await buildEnvelopesForPayload(payload);238239    const buf = new ChunkBuffer();240    for (const env of envelopes) {241      const res = await buf.add(env);242      expect(res.accepted).toBe(true);243    }244    expect(buf.isReady()).toBe(true);245246    const restored = await reconstructAndVerify(buf);247    expect(restored.verifyCode).toBe(payload.verifyCode);248    expect(restored.vaultBlob).toBe(payload.vaultBlob);249    expect(restored.vaultFingerprint).toBe(payload.vaultFingerprint);250    expect(restored.entryCount).toBe(payload.entryCount);251  });252253  it("reconstructs after losing 2 chunks", async () => {254    const payload = await makePayload();255    const envelopes = await buildEnvelopesForPayload(payload);256    // Need at least dataShards + 2 chunks total so dropping 2 still leaves257    // enough for reconstruction (small vaults pin to the 2-shard parity floor).258    expect(envelopes.length).toBeGreaterThan(2);259260    // Drop two indices that exercise the parity-decode path: one data shard261    // and one parity shard. The parity floor of 2 means we can lose any two.262    const dropA = 1;263    const dropB = envelopes.length - 1;264265    const buf = new ChunkBuffer();266    for (let i = 0; i < envelopes.length; i++) {267      if (i === dropA || i === dropB) continue;268      const res = await buf.add(envelopes[i]!);269      expect(res.accepted).toBe(true);270    }271    expect(buf.isReady()).toBe(true);272273    const restored = await reconstructAndVerify(buf);274    expect(restored.vaultFingerprint).toBe(payload.vaultFingerprint);275    expect(restored.vaultBlob).toBe(payload.vaultBlob);276  });277278  it("throws RECONSTRUCTION_FINGERPRINT_MISMATCH if the reconstructed blob does not match", async () => {279    // Build a payload whose stored fingerprint deliberately does NOT match280    // its vaultBlob — simulates a successful chunk decode that nevertheless281    // pulled the wrong bytes into a parity-reconstructed slot.282    const payload = await makePayload();283    const tampered: CodexPayload = {284      ...payload,285      vaultFingerprint: "DEADBEEFDEADBEEF",286    };287    const envelopes = await buildEnvelopesForPayload(tampered);288289    const buf = new ChunkBuffer();290    for (const env of envelopes) await buf.add(env);291292    await expect(reconstructAndVerify(buf)).rejects.toMatchObject({293      name: "CodexError",294      code: "RECONSTRUCTION_FINGERPRINT_MISMATCH",295    });296  });297298  it("throws CodexError when called on an empty buffer", async () => {299    const buf = new ChunkBuffer();300    await expect(reconstructAndVerify(buf)).rejects.toBeInstanceOf(CodexError);301  });302});303304describe("ChunkBuffer + parseRawScan integration", () => {305  it("a full scan loop accepts every raw envelope string", async () => {306    const payload = await makePayload(2_500);307    const envelopes = await buildEnvelopesForPayload(payload);308    const buf = new ChunkBuffer();309310    for (const env of envelopes) {311      const parsed = parseRawScan(JSON.stringify(env));312      expect("error" in parsed).toBe(false);313      const res = await buf.add(parsed as ChunkEnvelope);314      expect(res.accepted).toBe(true);315    }316317    const restored = await reconstructAndVerify(buf);318    expect(restored.vaultFingerprint).toBe(payload.vaultFingerprint);319    // sanity-check that base64UrlToBytes is reachable from this file320    expect(base64UrlToBytes(restored.vaultBlob).length).toBeGreaterThan(0);321  });322});