From 0f8df23a8994de91fe56970c281714184413a9a5 Mon Sep 17 00:00:00 2001 From: Test User Date: Fri, 10 Apr 2026 09:46:45 +0800 Subject: [PATCH] fix: prevent integer overflow and OOB reads in WASM safety checks 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 --- src/format.zig | 4 ++- src/turboquant.zig | 64 +++++++++++++++++++++++++++++++++++++++++--- src/wasm_exports.zig | 5 +++- 3 files changed, 67 insertions(+), 6 deletions(-) diff --git a/src/format.zig b/src/format.zig index 0b96b98..b416934 100644 --- a/src/format.zig +++ b/src/format.zig @@ -65,7 +65,9 @@ pub fn readHeader(data: []const u8) FormatError!Header { pub fn slicePayload(data: []const u8, header: Header) FormatError!struct { polar: []const u8, qjl: []const u8 } { const payload_start = HEADER_SIZE; - const payload_end = payload_start + header.polar_bytes + header.qjl_bytes; + // Guard against u32 overflow when summing polar_bytes + qjl_bytes + const payload_len = std.math.add(u32, header.polar_bytes, header.qjl_bytes) catch return FormatError.InvalidPayload; + const payload_end = payload_start + payload_len; if (data.len < payload_end) return FormatError.InvalidPayload; return .{ diff --git a/src/turboquant.zig b/src/turboquant.zig index 21f73b6..a53b7c4 100644 --- a/src/turboquant.zig +++ b/src/turboquant.zig @@ -10,6 +10,26 @@ pub const qjl = @import("qjl.zig"); pub const EncodeError = error{ InvalidDimension, OutOfMemory }; pub const DecodeError = error{ InvalidHeader, InvalidPayload, OutOfMemory }; +/// Validate that payload field sizes in the header are consistent with the +/// declared dimension. Prevents out-of-bounds reads in polar/qjl when a +/// crafted header claims a large dim but provides undersized payloads. +fn validatePayloadSizes(header: format.Header) bool { + const dim: u32 = header.dim; + if (dim == 0 or dim % 2 != 0) return false; + + // polar: (dim/2) pairs * 7 bits, rounded up to bytes, plus 1 padding byte + const num_pairs = dim / 2; + const polar_bits = num_pairs * 7; + const expected_polar = (polar_bits + 7) / 8 + 1; + if (header.polar_bytes != expected_polar) return false; + + // qjl: one sign bit per dimension, rounded up to bytes + const expected_qjl = (dim + 7) / 8; + if (header.qjl_bytes != expected_qjl) return false; + + return true; +} + pub const EngineConfig = struct { dim: usize, seed: u32, @@ -114,6 +134,7 @@ pub const Engine = struct { error.InvalidPayload => return DecodeError.InvalidPayload, }; if (header.dim != e.dim) return DecodeError.InvalidPayload; + if (!validatePayloadSizes(header)) return DecodeError.InvalidPayload; const payload = format.slicePayload(compressed, header) catch |err| switch (err) { error.InvalidHeader => return DecodeError.InvalidHeader, @@ -149,6 +170,7 @@ pub const Engine = struct { pub fn dot(e: *Engine, q: []const f32, compressed: []const u8) f32 { const header = format.readHeader(compressed) catch return 0; if (q.len != e.dim or header.dim != e.dim) return 0; + if (!validatePayloadSizes(header)) return 0; const payload = format.slicePayload(compressed, header) catch return 0; @@ -188,7 +210,7 @@ pub const Engine = struct { // Precompute q_sum for QJL fast path (SIMD reduction) var q_sum: f32 = 0; const d = rotated_q.len; - const lane: usize = comptime if (is_aarch64) 4 else 4; + const lane = std.simd.suggestVectorLength(f32) orelse 4; { var sum_vec: @Vector(lane, f32) = @splat(0); var si: usize = 0; @@ -223,9 +245,6 @@ pub const Engine = struct { out_scores[i] = polar_sum + qjl_sum; } } - - const builtin = @import("builtin"); - const is_aarch64 = builtin.cpu.arch == .aarch64; }; fn computeResidualFromPolar(polar_encoded: []const u8, rotated: []const f32, max_r: f32, residual: []f32) void { @@ -435,6 +454,25 @@ test "decode rejects truncated payload" { try std.testing.expectError(DecodeError.InvalidPayload, result); } +test "decode rejects payload with inconsistent field sizes" { + const allocator = std.testing.allocator; + var engine = try Engine.init(allocator, .{ .dim = 8, .seed = 12345 }); + defer engine.deinit(allocator); + + // Encode a valid vector, then corrupt the polar_bytes field in the header + const x: [8]f32 = .{ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0 }; + const compressed = try engine.encode(allocator, &x); + defer allocator.free(compressed); + + // Corrupt polar_bytes (offset 6..10) to an inconsistent value + var corrupted = try allocator.dupe(u8, compressed); + defer allocator.free(corrupted); + std.mem.writeInt(u32, corrupted[6..10], 0, .little); + + const result = engine.decode(allocator, corrupted); + try std.testing.expectError(DecodeError.InvalidPayload, result); +} + test "dot returns zero on dimension mismatch" { const allocator = std.testing.allocator; var engine = try Engine.init(allocator, .{ .dim = 8, .seed = 12345 }); @@ -449,6 +487,24 @@ test "dot returns zero on dimension mismatch" { try std.testing.expectEqual(0.0, result); } +test "dot returns zero on payload with inconsistent field sizes" { + const allocator = std.testing.allocator; + var engine = try Engine.init(allocator, .{ .dim = 8, .seed = 12345 }); + defer engine.deinit(allocator); + + const x: [8]f32 = .{ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0 }; + const compressed = try engine.encode(allocator, &x); + defer allocator.free(compressed); + + var corrupted = try allocator.dupe(u8, compressed); + defer allocator.free(corrupted); + std.mem.writeInt(u32, corrupted[6..10], 0, .little); + + const q: [8]f32 = .{ 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8 }; + const result = engine.dot(&q, corrupted); + try std.testing.expectEqual(0.0, result); +} + test "roundtrip correct length and finite" { const allocator = std.testing.allocator; var engine = try Engine.init(allocator, .{ .dim = 64, .seed = 9999 }); diff --git a/src/wasm_exports.zig b/src/wasm_exports.zig index ebc3ffc..7b898cf 100644 --- a/src/wasm_exports.zig +++ b/src/wasm_exports.zig @@ -118,8 +118,11 @@ export fn tq_dot_batch( const idx = resolveHandle(handle) orelse return; const engine_ptr = engine_slots[idx] orelse return; + // Guard against u32 overflow: num_vectors * bytes_per_vector must not wrap. + const total_bytes = std.math.mul(u32, num_vectors, bytes_per_vector) catch return; + const q = query_ptr[0..dim]; - const compressed = compressed_ptr[0 .. num_vectors * bytes_per_vector]; + const compressed = compressed_ptr[0..total_bytes]; const scores = out_scores[0..num_vectors]; engine_ptr.dotBatch(q, compressed, bytes_per_vector, num_vectors, scores);