Skip to content

fix: prevent integer overflow and OOB reads in WASM safety checks - #8

Merged
teamchong merged 2 commits into
teamchong:mainfrom
hobostay:fix/wasm-safety-bugs
Apr 19, 2026
Merged

fix: prevent integer overflow and OOB reads in WASM safety checks#8
teamchong merged 2 commits into
teamchong:mainfrom
hobostay:fix/wasm-safety-bugs

Conversation

@hobostay

Copy link
Copy Markdown
Contributor

Summary

Four safety bugs found and fixed in the WASM compression pipeline:

  1. Integer overflow in tq_dot_batch (wasm_exports.zig): num_vectors * bytes_per_vector could wrap u32, creating an undersized slice and causing out-of-bounds reads in the dotBatch loop. Fixed with std.math.mul overflow check.

  2. Integer overflow in format.slicePayload (format.zig): polar_bytes + qjl_bytes could wrap u32, making payload_end smaller than actual payload, bypassing the bounds check. Fixed with std.math.add overflow check.

  3. Missing payload size validation (turboquant.zig): decode() and dot() validated header dimension but not that polar_bytes/qjl_bytes were consistent with that dimension. A crafted compressed buffer with correct dim but undersized payload fields would cause out-of-bounds reads in polar/qjl functions. Added validatePayloadSizes() to check expected sizes before use.

  4. Dead code in dotBatch (turboquant.zig): comptime if (is_aarch64) 4 else 4 always evaluates to 4 — both branches are identical. Replaced with std.simd.suggestVectorLength for consistency with other modules, and removed the unused builtin import.

Test plan

  • All existing tests pass (zig build test)
  • New test: decode rejects payload with inconsistent field sizes — verifies decode rejects corrupted polar_bytes
  • New test: dot returns zero on payload with inconsistent field sizes — verifies dot returns 0 on corrupted payload
  • Manual test: tq_dot_batch with large num_vectors * bytes_per_vector that overflows u32 should return silently

🤖 Generated with Claude Code

@teamchong

Copy link
Copy Markdown
Owner

Thanks for catching these, the four fixes are all real and still applicable on current main.

Three small asks:

  1. Rebase on main — PR feat: TQ-compressed KV cache for browser LLM (WebGPU) #7 just landed and touched the same lines in decodeInto / dot:
  2. One-line conflict in src/turboquant.zig around decodeInto (same pattern in dot)
  3. Format check — there's an extra blank line before the closing }; of Engine at line 246. zig fmt src/ will fix it. CI is failing on format check right now, and this is the only thing it catches.

If it'd be easier, toggle "Allow edits from maintainers" on the PR and I'll push the rebase for you. Otherwise happy to re-review as soon as CI flips green.

Four bugs fixed:

1. Integer overflow in tq_dot_batch: num_vectors * bytes_per_vector
   could wrap u32, creating an undersized buffer and causing
   out-of-bounds reads in the dotBatch loop. Use std.math.mul to
   detect overflow and return early.

2. Integer overflow in format.slicePayload: polar_bytes + qjl_bytes
   could wrap u32, making payload_end smaller than the actual payload
   and bypassing the bounds check. Use std.math.add to detect overflow.

3. Payload size validation in decode/dot: the engine now validates
   that polar_bytes and qjl_bytes in the header are consistent with
   the declared dimension before passing them to polar/qjl functions.
   Prevents out-of-bounds reads from crafted compressed data with
   correct dim but undersized payload fields.

4. Dead code removal in dotBatch: the comptime conditional
   `if (is_aarch64) 4 else 4` always evaluated to 4. Replaced with
   std.simd.suggestVectorLength for consistency with other modules.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@hobostay

Copy link
Copy Markdown
Contributor Author

Hi @teamchong, thanks for the quick review!

I've rebased on main (after #7), resolved the one-line conflict in decodeInto, and fixed the extra blank line before }; in Engine. All three items addressed:

  1. Rebased on latest main — conflict in src/turboquant.zig resolved (kept the validatePayloadSizes check on top of the updated decodeInto/dot from feat: TQ-compressed KV cache for browser LLM (WebGPU) #7)
  2. Format fix applied — removed the extra blank line before };
  3. CI should pass on format check now

Ready for re-review whenever you have a moment.

@teamchong
teamchong merged commit 6756491 into teamchong:main Apr 19, 2026
5 checks passed
teamchong added a commit that referenced this pull request Apr 19, 2026
…Needed

Follow-up to #6756491 (PR #8). The original validatePayloadSizes
duplicated the polar/qjl encoded-size formulas inline:

  const polar_bits = num_pairs * 7;            // BITS_PER_PAIR
  const expected_polar = (polar_bits + 7) / 8 + 1;
  const expected_qjl = (dim + 7) / 8;

Both modules already export the same math as pub helpers
(polar.polarBytesNeeded, qjl.qjlBytesNeeded). Calling them directly
keeps the validator in sync automatically if either format ever
changes its internal widths (e.g. R_BITS/THETA_BITS retuning,
QJL widening from 1-bit). Removes the silent drift risk.

Also adds symmetric test coverage for the qjl_bytes validation
path — the original tests only corrupted polar_bytes (offset 6..10),
leaving the qjl_bytes check at offset 10..14 untested. Two new
tests (decode + dot) corrupt qjl_bytes and assert the same
InvalidPayload / 0.0 behaviour. 63 tests total, all pass.
teamchong added a commit that referenced this pull request Apr 19, 2026
- Rewrite design doc: removes 'not implemented' status, fixes 2-mode
  claim (now 9 branches), corrects u64→u32 offsets in the format
  spec, updates sizes to match the actual 88.3 MB cache, moves done
  items into the shipped checklist, removes the 'Names' padding
  section.

- Bump to 0.4.1 for PR #8 (security fixes: u32 overflow guards in
  tq_dot_batch + format.slicePayload, validatePayloadSizes for
  decode/dot) + follow-up refactor that delegates to
  polarBytesNeeded/qjlBytesNeeded.
teamchong added a commit that referenced this pull request Apr 19, 2026
Dynamic context window for the Gemma 4 E2B draw demo. Engine boots with a small "router" system prompt; when the model emits setType("..."), a side-channel observer detects it and engine.mountKV swaps the active KV to that type's pre-baked branch — no runtime prefill on the swap.

9 per-type branches under demo/src/draw/prompts/ (router + architecture / sequence / flowchart / state / orgchart / er / class / swimlane), packed into a new TQKC multi-branch container format at public/system-cache.bin (88.3 MB). Design: demo/docs/dynamic-context-window.md.

Core pieces:
- ModeTracker in grammar.ts with 9 SDK_MODE constants + onEnter hooks
- engine.mountKV(name) + registerBranch
- worker mountBranchAndPrefill: swap KV + re-prefill user turn atomically
- DiagramType union expanded to all 8 values in types.ts + sdk-types.ts
- Router skips the thinking phase (its only job is to pick setType); branch runs full think+code under specialised KV
- Retry restores to router snapshot (activeBranch bookkeeping included)

Rolls in demo improvements accumulated on this branch: subgroup-butterfly reductions, window-bound softmax, shared-KV skip in batched prefill, uniform-buffer dedup, mobile-aware WebGPU error card, markdown-stripped thinking cloud, cancel-and-restart, singleton lock, orphan auto-fix before retry, retry-LCP rollback.

Follow-up to #8: validatePayloadSizes now delegates to polarBytesNeeded / qjlBytesNeeded instead of re-deriving the formulas, plus symmetric qjl_bytes-corruption tests. Package bumped to 0.4.1.

22 new TS unit tests + 2 new Zig tests (63 total), all pass. TS clean, zig fmt clean, CI green across Linux/macOS/Windows.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants