Skip to content

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

Merged
copybara-service[bot] merged 2 commits into
google:devfrom
0xASTRA:fix/weights-readallbf16-missing-blob-size-assert
Jun 2, 2026
Merged

[Security] Add missing blob-size assertion in ReadAllToBF16 weight loader#924
copybara-service[bot] merged 2 commits into
google:devfrom
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.
@jan-wassenberg
Copy link
Copy Markdown
Member

Thanks for making the fix :) We're only able to merge on the dev branch, would you mind retargeting to that?

@0xASTRA 0xASTRA changed the base branch from main to dev June 2, 2026 13:38
@0xASTRA
Copy link
Copy Markdown
Author

0xASTRA commented Jun 2, 2026

Thanks for making the fix :) We're only able to merge on the dev branch, would you mind retargeting to that?

Done!

@jan-wassenberg jan-wassenberg added the copybara-import Trigger Copybara for merging pull requests label Jun 2, 2026
@0xASTRA
Copy link
Copy Markdown
Author

0xASTRA commented Jun 2, 2026

Thanks! Can it be closed?

@copybara-service copybara-service Bot merged commit 4013e52 into google:dev Jun 2, 2026
2 of 3 checks passed
@jan-wassenberg
Copy link
Copy Markdown
Member

FYI the logic here is a bit off. We only want to check this if keep_type == true, i.e. move the code down by one line. I'll fix that soon.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

copybara-import Trigger Copybara for merging pull requests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants