restore.test.ts

10.6 KB · raw

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("toStripeChunks returns raw chunk bytes per stripe", 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 stripes = buf.toStripeChunks();196    expect(stripes.length).toBe(1);197    const map = stripes[0]!.chunks;198    expect(map.size).toBe(1);199    expect(Array.from(map.get(0)!)).toEqual([1, 2, 3]);200  });201});202203describe("parseRawScan", () => {204  it("parses a well-formed envelope", async () => {205    const env = await makeSingleEnvelope({ index: 0, total: 5 });206    const parsed = parseRawScan(JSON.stringify(env));207    expect(parsed).toEqual(env);208  });209210  it("returns parse_error for non-JSON input", () => {211    expect(parseRawScan("not json")).toEqual({ error: "parse_error" });212  });213214  it("returns wrong_version for v=1 envelopes", () => {215    const v1 = {216      v: 1,217      c: VERIFY_CODE,218      i: 0,219      t: 5,220      p: false,221      d: "AAAA",222      h: "ABCDEF12",223    };224    expect(parseRawScan(JSON.stringify(v1))).toEqual({225      error: "wrong_version",226    });227  });228229  it("returns parse_error for a v2 object with missing fields", () => {230    expect(parseRawScan(JSON.stringify({ v: 2, c: "PX" }))).toEqual({231      error: "parse_error",232    });233  });234});235236describe("reconstructAndVerify", () => {237  it("round-trips a full envelope set back to the original payload", async () => {238    const payload = await makePayload();239    const envelopes = await buildEnvelopesForPayload(payload);240241    const buf = new ChunkBuffer();242    for (const env of envelopes) {243      const res = await buf.add(env);244      expect(res.accepted).toBe(true);245    }246    expect(buf.isReady()).toBe(true);247248    const restored = await reconstructAndVerify(buf);249    expect(restored.verifyCode).toBe(payload.verifyCode);250    expect(restored.vaultBlob).toBe(payload.vaultBlob);251    expect(restored.vaultFingerprint).toBe(payload.vaultFingerprint);252    expect(restored.entryCount).toBe(payload.entryCount);253  });254255  it("reconstructs after losing 2 chunks", async () => {256    const payload = await makePayload();257    const envelopes = await buildEnvelopesForPayload(payload);258    // Need at least dataShards + 2 chunks total so dropping 2 still leaves259    // enough for reconstruction (small vaults pin to the 2-shard parity floor).260    expect(envelopes.length).toBeGreaterThan(2);261262    // Drop two indices that exercise the parity-decode path: one data shard263    // and one parity shard. The parity floor of 2 means we can lose any two.264    const dropA = 1;265    const dropB = envelopes.length - 1;266267    const buf = new ChunkBuffer();268    for (let i = 0; i < envelopes.length; i++) {269      if (i === dropA || i === dropB) continue;270      const res = await buf.add(envelopes[i]!);271      expect(res.accepted).toBe(true);272    }273    expect(buf.isReady()).toBe(true);274275    const restored = await reconstructAndVerify(buf);276    expect(restored.vaultFingerprint).toBe(payload.vaultFingerprint);277    expect(restored.vaultBlob).toBe(payload.vaultBlob);278  });279280  it("throws RECONSTRUCTION_FINGERPRINT_MISMATCH if the reconstructed blob does not match", async () => {281    // Build a payload whose stored fingerprint deliberately does NOT match282    // its vaultBlob — simulates a successful chunk decode that nevertheless283    // pulled the wrong bytes into a parity-reconstructed slot.284    const payload = await makePayload();285    const tampered: CodexPayload = {286      ...payload,287      vaultFingerprint: "DEADBEEFDEADBEEF",288    };289    const envelopes = await buildEnvelopesForPayload(tampered);290291    const buf = new ChunkBuffer();292    for (const env of envelopes) await buf.add(env);293294    await expect(reconstructAndVerify(buf)).rejects.toMatchObject({295      name: "CodexError",296      code: "RECONSTRUCTION_FINGERPRINT_MISMATCH",297    });298  });299300  it("throws CodexError when called on an empty buffer", async () => {301    const buf = new ChunkBuffer();302    await expect(reconstructAndVerify(buf)).rejects.toBeInstanceOf(CodexError);303  });304});305306describe("ChunkBuffer + parseRawScan integration", () => {307  it("a full scan loop accepts every raw envelope string", async () => {308    const payload = await makePayload(2_500);309    const envelopes = await buildEnvelopesForPayload(payload);310    const buf = new ChunkBuffer();311312    for (const env of envelopes) {313      const parsed = parseRawScan(JSON.stringify(env));314      expect("error" in parsed).toBe(false);315      const res = await buf.add(parsed as ChunkEnvelope);316      expect(res.accepted).toBe(true);317    }318319    const restored = await reconstructAndVerify(buf);320    expect(restored.vaultFingerprint).toBe(payload.vaultFingerprint);321    // sanity-check that base64UrlToBytes is reachable from this file322    expect(base64UrlToBytes(restored.vaultBlob).length).toBeGreaterThan(0);323  });324});