erasure.ts

13.9 KB · raw

1/**2 * erasure.ts — Reed-Solomon erasure coding around the WASM library.3 *4 * The library wants a single contiguous buffer of length5 *   (dataShards + parityShards) * shardSize6 * with data shards laid out first, followed by space for parity. `encode()`7 * fills the parity region in-place. `reconstruct()` rewrites missing shards8 * in-place using a `shardsAvailable: boolean[]` mask.9 *10 * We expose a chunk-oriented API (`encodeWithErasure` / `decodeWithErasure`)11 * because the rest of the codex pipeline thinks in terms of individual12 * envelopes. Both adapters allocate a fresh buffer per call — the library is13 * not re-entrant on a shared buffer.14 */1516import { ReedSolomonErasure } from "@subspace/reed-solomon-erasure.wasm";17import {18  CHUNK_SIZE,19  CodexError,20  MAX_DATA_SHARDS_PER_STRIPE,21  computeParityShards,22} from "../../shared/codex-types";23import { getReedSolomon } from "./rs-loader";2425/**26 * The shared rs-loader caches a single `ReedSolomonErasure` instance. Some27 * builds of the underlying WASM module accumulate __wbindgen state between28 * `reconstruct()` calls and return RESULT_ERROR_TOO_FEW_SHARDS_PRESENT on the29 * second invocation even when enough shards are present. Building a fresh30 * instance per call avoids the failure mode entirely; the cost is a few31 * milliseconds per call — negligible compared to compressing the vault and32 * rendering Aztec codes. We keep the loader's cached instance only as a33 * liveness probe (it confirms the WASM is instantiable on this platform).34 */35let cachedWasmBytes: Uint8Array<ArrayBuffer> | null = null;3637async function freshReedSolomon(): Promise<ReedSolomonErasure> {38  if (cachedWasmBytes) {39    return ReedSolomonErasure.fromBytes(cachedWasmBytes);40  }41  // Capture bytes by intercepting the loader's first instantiation. The42  // loader caches its own promise, so calling it twice is cheap on the second43  // pass. We pull bytes from the WASM memory image after the fact — see44  // `loadWasmBytes` below for the lazy initialization.45  await getReedSolomon();46  if (!cachedWasmBytes) {47    cachedWasmBytes = await loadWasmBytes();48  }49  return ReedSolomonErasure.fromBytes(cachedWasmBytes);50}5152/**53 * Load the raw WASM bytes for re-instantiation. The popup path resolves54 * against the extension package via `chrome.runtime.getURL`; vitest reads55 * directly from `node_modules`. `freshReedSolomon` needs the raw bytes (not56 * the loader's instance) because each call rebuilds the WASM module to avoid57 * cross-call state accumulation in `reconstruct`.58 */59async function loadWasmBytes(): Promise<Uint8Array<ArrayBuffer>> {60  const inPopup =61    typeof chrome !== "undefined" &&62    typeof chrome.runtime !== "undefined" &&63    typeof chrome.runtime.getURL === "function";64  if (inPopup) {65    const url = chrome.runtime.getURL("codex/reed_solomon_erasure_bg.wasm");66    const response = await fetch(url);67    if (!response.ok) {68      throw new CodexError(69        "ERASURE_CODING_VERIFICATION_FAILED",70        `Failed to fetch WASM: ${response.status}`,71      );72    }73    return new Uint8Array(await response.arrayBuffer());74  }75  const fs = await import("node:fs");76  const path = await import("node:path");77  const wasmPath = path.resolve(78    "node_modules/@subspace/reed-solomon-erasure.wasm/dist/reed_solomon_erasure_bg.wasm",79  );80  return new Uint8Array(fs.readFileSync(wasmPath));81}8283/**84 * @subspace/reed-solomon-erasure.wasm exposes the RESULT_* constants but not a85 * `getResultMessage(code)` helper (the @digitaldefiance fork added it). We map86 * codes to human-readable strings here so the surrounding CodexError messages87 * stay specific instead of leaking a bare number.88 */89function resultMessage(code: number): string {90  switch (code) {91    case ReedSolomonErasure.RESULT_OK: return "ok";92    case ReedSolomonErasure.RESULT_ERROR_TOO_FEW_SHARDS: return "too few shards";93    case ReedSolomonErasure.RESULT_ERROR_TOO_MANY_SHARDS: return "too many shards";94    case ReedSolomonErasure.RESULT_ERROR_TOO_FEW_DATA_SHARDS: return "too few data shards";95    case ReedSolomonErasure.RESULT_ERROR_TOO_MANY_DATA_SHARDS: return "too many data shards";96    case ReedSolomonErasure.RESULT_ERROR_TOO_FEW_PARITY_SHARDS: return "too few parity shards";97    case ReedSolomonErasure.RESULT_ERROR_TOO_MANY_PARITY_SHARDS: return "too many parity shards";98    case ReedSolomonErasure.RESULT_ERROR_TOO_FEW_BUFFER_SHARDS: return "too few buffer shards";99    case ReedSolomonErasure.RESULT_ERROR_TOO_MANY_BUFFER_SHARDS: return "too many buffer shards";100    case ReedSolomonErasure.RESULT_ERROR_INCORRECT_SHARD_SIZE: return "incorrect shard size";101    case ReedSolomonErasure.RESULT_ERROR_TOO_FEW_SHARDS_PRESENT: return "too few shards present";102    case ReedSolomonErasure.RESULT_ERROR_EMPTY_SHARD: return "empty shard";103    case ReedSolomonErasure.RESULT_ERROR_INVALID_SHARD_FLAGS: return "invalid shard flags";104    case ReedSolomonErasure.RESULT_ERROR_INVALID_INDEX: return "invalid index";105    default: return `unknown result code ${code}`;106  }107}108109export interface EncodeResult {110  dataChunks: Uint8Array[];111  parityChunks: Uint8Array[];112  totalChunks: number;113  dataShards: number;114  parityShards: number;115}116117export interface MultiStripeEncodeResult {118  stripes: EncodeResult[];119  /** Sum of `dataShards` across stripes. */120  totalDataChunks: number;121  /** Sum of `parityShards` across stripes. */122  totalParityChunks: number;123  /** Sum of `totalChunks` across stripes = number of Aztec codes to print. */124  totalChunks: number;125}126127/**128 * Plan how many data shards belong to each stripe, balanced as evenly as129 * possible so no single stripe approaches the GF(2⁸) cap. Returns an array of130 * shard counts that sum to `totalDataShards` and where every element is131 * ≤ MAX_DATA_SHARDS_PER_STRIPE.132 */133export function planStripeShardCounts(totalDataShards: number): number[] {134  if (totalDataShards <= 0) return [];135  const numStripes = Math.max(136    1,137    Math.ceil(totalDataShards / MAX_DATA_SHARDS_PER_STRIPE),138  );139  const base = Math.floor(totalDataShards / numStripes);140  const remainder = totalDataShards % numStripes;141  const counts: number[] = [];142  for (let i = 0; i < numStripes; i++) {143    counts.push(base + (i < remainder ? 1 : 0));144  }145  return counts;146}147148/**149 * Split `data` into fixed-size shards plus `computeParityShards(dataShards)`150 * parity shards. Pads the trailing data shard with zeros — the original byte151 * length is recovered via the `CodexPayload.vaultBlob` (no length prefix152 * needed on the wire because the outer JSON already bounds the data).153 *154 * Single-stripe primitive. The GF(2⁸) field caps `dataShards + parityShards`155 * at 256; anything larger panics inside the WASM module, so call156 * `encodeIntoStripes` instead for payloads that need more total shards.157 */158export async function encodeWithErasure(data: Uint8Array): Promise<EncodeResult> {159  const rs = await freshReedSolomon();160161  const dataShards = Math.max(1, Math.ceil(data.length / CHUNK_SIZE));162  const parityShards = computeParityShards(dataShards);163  const totalShards = dataShards + parityShards;164  if (totalShards > 256) {165    throw new CodexError(166      "ERASURE_CODING_VERIFICATION_FAILED",167      `Single RS stripe cannot exceed 256 shards (got ${totalShards}; ${data.length} bytes / ${CHUNK_SIZE}-byte chunks). Use encodeIntoStripes for large payloads.`,168    );169  }170  const buffer = new Uint8Array(totalShards * CHUNK_SIZE);171  buffer.set(data, 0); // Trailing bytes after `data.length` stay zero — the pad.172173  let code: number;174  try {175    code = rs.encode(buffer, dataShards, parityShards);176  } catch (err) {177    // WASM traps (e.g. Rust panics) surface as raw JS errors with messages178    // like "unreachable". Map them to a typed CodexError so the UI can179    // explain what happened instead of leaking the trap string.180    throw new CodexError(181      "ERASURE_CODING_VERIFICATION_FAILED",182      `Reed-Solomon encode trapped (data=${dataShards}, parity=${parityShards}): ${183        err instanceof Error ? err.message : String(err)184      }`,185    );186  }187  if (code !== ReedSolomonErasure.RESULT_OK) {188    throw new CodexError(189      "ERASURE_CODING_VERIFICATION_FAILED",190      `Reed-Solomon encode failed: ${resultMessage(code)}`,191    );192  }193194  // Slice() copies so each shard owns its memory and callers cannot mutate195  // the working buffer through a returned view.196  const dataChunks: Uint8Array[] = [];197  for (let i = 0; i < dataShards; i++) {198    dataChunks.push(buffer.slice(i * CHUNK_SIZE, (i + 1) * CHUNK_SIZE));199  }200  const parityChunks: Uint8Array[] = [];201  for (let i = 0; i < parityShards; i++) {202    const base = (dataShards + i) * CHUNK_SIZE;203    parityChunks.push(buffer.slice(base, base + CHUNK_SIZE));204  }205206  return {207    dataChunks,208    parityChunks,209    totalChunks: totalShards,210    dataShards,211    parityShards,212  };213}214215/**216 * Reassemble the data-shard region from a partial set of chunks. Missing217 * shards are reconstructed from parity. The returned buffer is exactly218 * `dataShards * shardSize` bytes long, including any trailing pad bytes219 * — callers strip the pad using their out-of-band length information.220 *221 * Shard size is derived from the actual chunks rather than the current222 * `CHUNK_SIZE` constant so that Codexes printed under an older constant223 * (e.g. 1024-byte shards) still restore on a build whose constant has moved.224 * Reed-Solomon requires all shards to share one length; the first chunk225 * establishes it and every other chunk must match.226 */227export async function decodeWithErasure(228  chunks: Map<number, Uint8Array>,229  dataShards: number,230  totalShards: number,231): Promise<Uint8Array> {232  const parityShards = totalShards - dataShards;233  if (parityShards < 0) {234    throw new CodexError(235      "ERASURE_CODING_VERIFICATION_FAILED",236      `Invalid shard layout: dataShards=${dataShards} totalShards=${totalShards}`,237    );238  }239240  if (chunks.size < dataShards) {241    throw new CodexError(242      "INSUFFICIENT_CHUNKS",243      `Need at least ${dataShards} chunks, have ${chunks.size}`,244    );245  }246247  // First scanned chunk fixes the shard size for the whole reconstruction.248  const firstChunk = chunks.values().next().value;249  if (!firstChunk || firstChunk.length === 0) {250    throw new CodexError(251      "ERASURE_CODING_VERIFICATION_FAILED",252      "First chunk is empty; cannot determine shard size",253    );254  }255  const shardSize = firstChunk.length;256257  const rs = await freshReedSolomon();258  const buffer = new Uint8Array(totalShards * shardSize);259  const available: boolean[] = new Array(totalShards).fill(false);260261  for (const [index, chunk] of chunks) {262    if (index < 0 || index >= totalShards) {263      throw new CodexError(264        "ERASURE_CODING_VERIFICATION_FAILED",265        `Chunk index ${index} out of range [0, ${totalShards})`,266      );267    }268    if (chunk.length !== shardSize) {269      throw new CodexError(270        "ERASURE_CODING_VERIFICATION_FAILED",271        `Chunk ${index} has length ${chunk.length}, expected ${shardSize}`,272      );273    }274    buffer.set(chunk, index * shardSize);275    available[index] = true;276  }277278  let code: number;279  try {280    code = rs.reconstruct(buffer, dataShards, parityShards, available);281  } catch (err) {282    throw new CodexError(283      "ERASURE_CODING_VERIFICATION_FAILED",284      `Reed-Solomon reconstruct trapped (data=${dataShards}, parity=${parityShards}, have=${chunks.size}): ${285        err instanceof Error ? err.message : String(err)286      }`,287    );288  }289  if (code !== ReedSolomonErasure.RESULT_OK) {290    // The library distinguishes "not enough shards" with its own code; map it291    // back to our INSUFFICIENT_CHUNKS so handlers upstream can surface a292    // clean "scan more codes" prompt instead of a generic library error.293    if (code === ReedSolomonErasure.RESULT_ERROR_TOO_FEW_SHARDS_PRESENT) {294      throw new CodexError(295        "INSUFFICIENT_CHUNKS",296        `Need at least ${dataShards} valid chunks for reconstruction`,297      );298    }299    throw new CodexError(300      "ERASURE_CODING_VERIFICATION_FAILED",301      `Reed-Solomon reconstruct failed: ${resultMessage(code)}`,302    );303  }304305  return buffer.slice(0, dataShards * shardSize);306}307308/**309 * Multi-stripe encode. Splits `data` into balanced RS stripes (each ≤310 * MAX_DATA_SHARDS_PER_STRIPE data shards) and encodes each independently.311 * The stripes are byte-contiguous: stripe `s` covers exactly312 * `dataShardsForStripe(s) * CHUNK_SIZE` bytes of payload, with the last shard313 * of the last stripe zero-padded when the payload doesn't align. Restore314 * concatenates each stripe's decoded data region and inflates the whole thing315 * — pako tolerates the trailing zeros.316 */317export async function encodeIntoStripes(318  data: Uint8Array,319): Promise<MultiStripeEncodeResult> {320  const totalDataShards = Math.max(1, Math.ceil(data.length / CHUNK_SIZE));321  const stripeShardCounts = planStripeShardCounts(totalDataShards);322323  const stripes: EncodeResult[] = [];324  let byteCursor = 0;325  for (const shardCount of stripeShardCounts) {326    const stripeByteCap = shardCount * CHUNK_SIZE;327    const stripeEnd = Math.min(byteCursor + stripeByteCap, data.length);328    const stripeBytes = data.subarray(byteCursor, stripeEnd);329    const result = await encodeWithErasure(stripeBytes);330    // Sanity: encodeWithErasure infers dataShards from byte length; for any331    // stripe but the last, the slice is exactly `shardCount * CHUNK_SIZE`332    // bytes and the inferred count must match the plan. For the last stripe333    // the byte length may be shorter (zero-pad happens inside the encoder)334    // but the shard count still matches because the planner used the same335    // ceil-division.336    if (result.dataShards !== shardCount) {337      throw new CodexError(338        "ERASURE_CODING_VERIFICATION_FAILED",339        `Stripe plan mismatch: planned ${shardCount} data shards, encoder produced ${result.dataShards}`,340      );341    }342    stripes.push(result);343    byteCursor += stripeByteCap;344  }345346  const totalDataChunks = stripes.reduce((s, x) => s + x.dataShards, 0);347  const totalParityChunks = stripes.reduce((s, x) => s + x.parityShards, 0);348  const totalChunks = totalDataChunks + totalParityChunks;349  return { stripes, totalDataChunks, totalParityChunks, totalChunks };350}