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  PARITY_SHARDS,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    includesRecoveryWords: false,94  };95}9697async function makeSingleEnvelope(opts: {98  index: number;99  total: number;100  isParity?: boolean;101  verifyCode?: string;102  bytes?: Uint8Array;103}): Promise<ChunkEnvelope> {104  return buildChunkEnvelope(105    opts.bytes ?? new Uint8Array([1, 2, 3, 4, 5]),106    opts.index,107    opts.total,108    opts.isParity ?? false,109    opts.verifyCode ?? VERIFY_CODE,110  );111}112113describe("ChunkBuffer", () => {114  it("accepts a valid chunk and exposes context", async () => {115    const buf = new ChunkBuffer();116    const env = await makeSingleEnvelope({ index: 0, total: 5 });117    const res = await buf.add(env);118    expect(res.accepted).toBe(true);119    expect(buf.size()).toBe(1);120    expect(buf.getContext()).toEqual({121      verifyCode: VERIFY_CODE,122      total: 5,123      dataShards: 5 - PARITY_SHARDS,124    });125  });126127  it("rejects a chunk from a different verifyCode", async () => {128    const buf = new ChunkBuffer();129    await buf.add(await makeSingleEnvelope({ index: 0, total: 5 }));130    const foreign = await makeSingleEnvelope({131      index: 1,132      total: 5,133      verifyCode: "PX-OTHR-XXXX",134    });135    const res = await buf.add(foreign);136    expect(res).toEqual({ accepted: false, reason: "foreign_codex" });137    expect(buf.size()).toBe(1);138  });139140  it("rejects a chunk with bad integrity hash", async () => {141    const buf = new ChunkBuffer();142    const env = await makeSingleEnvelope({ index: 0, total: 5 });143    // Corrupt the data field without updating the hash — verifyChunkIntegrity144    // must catch this.145    const corrupted: ChunkEnvelope = {146      ...env,147      d: bytesToBase64Url(new Uint8Array([9, 9, 9, 9])),148    };149    const res = await buf.add(corrupted);150    expect(res).toEqual({ accepted: false, reason: "bad_integrity" });151    expect(buf.size()).toBe(0);152  });153154  it("rejects a chunk with wrong envelope version", async () => {155    const buf = new ChunkBuffer();156    const env = await makeSingleEnvelope({ index: 0, total: 5 });157    const wrongVersion = { ...env, v: 1 } as unknown as ChunkEnvelope;158    const res = await buf.add(wrongVersion);159    expect(res.accepted).toBe(false);160    expect(res.reason).toBe("parse_error");161  });162163  it("deduplicates: adding the same index twice keeps only one", async () => {164    const buf = new ChunkBuffer();165    const env = await makeSingleEnvelope({ index: 2, total: 5 });166    expect((await buf.add(env)).accepted).toBe(true);167    const second = await buf.add(env);168    expect(second).toEqual({ accepted: false, reason: "duplicate" });169    expect(buf.size()).toBe(1);170  });171172  it("isReady() flips true at dataShards count", async () => {173    const buf = new ChunkBuffer();174    const total = 5;175    const dataShards = total - PARITY_SHARDS; // 3176    for (let i = 0; i < dataShards; i++) {177      expect(buf.isReady()).toBe(false);178      const env = await makeSingleEnvelope({179        index: i,180        total,181        bytes: new Uint8Array([i, i, i]),182      });183      await buf.add(env);184    }185    expect(buf.isReady()).toBe(true);186  });187188  it("toChunkMap returns raw chunk bytes by index", async () => {189    const buf = new ChunkBuffer();190    const env = await makeSingleEnvelope({191      index: 0,192      total: 5,193      bytes: new Uint8Array([1, 2, 3]),194    });195    await buf.add(env);196    const map = buf.toChunkMap();197    expect(map.size).toBe(1);198    expect(Array.from(map.get(0)!)).toEqual([1, 2, 3]);199  });200});201202describe("parseRawScan", () => {203  it("parses a well-formed envelope", async () => {204    const env = await makeSingleEnvelope({ index: 0, total: 5 });205    const parsed = parseRawScan(JSON.stringify(env));206    expect(parsed).toEqual(env);207  });208209  it("returns parse_error for non-JSON input", () => {210    expect(parseRawScan("not json")).toEqual({ error: "parse_error" });211  });212213  it("returns wrong_version for v=1 envelopes", () => {214    const v1 = {215      v: 1,216      c: VERIFY_CODE,217      i: 0,218      t: 5,219      p: false,220      d: "AAAA",221      h: "ABCDEF12",222    };223    expect(parseRawScan(JSON.stringify(v1))).toEqual({224      error: "wrong_version",225    });226  });227228  it("returns parse_error for a v2 object with missing fields", () => {229    expect(parseRawScan(JSON.stringify({ v: 2, c: "PX" }))).toEqual({230      error: "parse_error",231    });232  });233});234235describe("reconstructAndVerify", () => {236  it("round-trips a full envelope set back to the original payload", async () => {237    const payload = await makePayload();238    const envelopes = await buildEnvelopesForPayload(payload);239240    const buf = new ChunkBuffer();241    for (const env of envelopes) {242      const res = await buf.add(env);243      expect(res.accepted).toBe(true);244    }245    expect(buf.isReady()).toBe(true);246247    const restored = await reconstructAndVerify(buf);248    expect(restored.verifyCode).toBe(payload.verifyCode);249    expect(restored.vaultBlob).toBe(payload.vaultBlob);250    expect(restored.vaultFingerprint).toBe(payload.vaultFingerprint);251    expect(restored.entryCount).toBe(payload.entryCount);252  });253254  it("reconstructs after losing 2 chunks", async () => {255    const payload = await makePayload();256    const envelopes = await buildEnvelopesForPayload(payload);257    // Need at least dataShards + PARITY_SHARDS total to lose exactly PARITY_SHARDS.258    expect(envelopes.length).toBeGreaterThan(PARITY_SHARDS);259260    // Drop two indices that exercise the parity-decode path: one data shard261    // and one parity shard. PARITY_SHARDS=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});