Skip to content

[Security] Add missing blob-size assertion in ReadAllToBF16 weight loader#924

Open
0xASTRA wants to merge 1 commit into
google:mainfrom
0xASTRA:fix/weights-readallbf16-missing-blob-size-assert
Open

[Security] Add missing blob-size assertion in ReadAllToBF16 weight loader#924
0xASTRA wants to merge 1 commit into
google:mainfrom
0xASTRA:fix/weights-readallbf16-missing-blob-size-assert

Conversation

@0xASTRA
Copy link
Copy Markdown

@0xASTRA 0xASTRA commented May 29, 2026

Summary

ReadAllToBF16 in weights.cc is the only one of three weight-load modes missing the range.bytes == mat.PackedBytes() assertion. The other two modes (MapAll at line ~557, MakeBatches at line ~645) both assert this invariant. Without the check, a crafted .sbs file with correct Rows/Cols metadata but an inflated blob bytes field passes all structural validation and causes a heap OOB write (keep_type branch) or OOB read (DecompressToBF16 branch).

Fix

Add HWY_ASSERT_M(tensor.range.bytes == mat.PackedBytes(), mat.Name()) at the top of the per-tensor lambda in ReadAllToBF16, before either read path executes — mirroring the existing checks in MapAll and MakeBatches.

Test plan

  • Existing weight-loading tests pass
  • A .sbs file with inflated blob bytes for one tensor is rejected with an assertion failure in kReadBF16 mode (--to_bf16 or decode_qbatch_size >= 128)

…ader

MapAll (weights.cc:557) and MakeBatches (weights.cc:645) both assert
that a tensor blob's on-disk size equals the buffer size derived from
the tensor's declared shape (range.bytes == mat.PackedBytes()).

ReadAllToBF16 was the only load mode missing this check. Without it, a
crafted .sbs weight file with correct Rows/Cols metadata but an inflated
blob bytes field passes all structural validation and reaches the
keep_type branch, where tensor.range.bytes bytes are written into a
mat.PackedBytes()-sized buffer — a heap OOB write with attacker-
controlled content. The DecompressToBF16 branch has a sibling OOB read
when the blob is undersized.

Add the same HWY_ASSERT_M at the top of the per-tensor lambda, before
either read path executes, mirroring the existing checks in MapAll and
MakeBatches.
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.

1 participant