Validate buffer slices during deserialization by M0nd0R · Pull Request #44653 · openxla/xla

M0nd0R · 2026-06-23T02:36:36Z

Summary

validate BufferAllocation::Slice::FromProto rejects negative offsets and sizes
reject serialized slices whose [offset, offset + size) range exceeds the referenced allocation
add regression coverage for malformed BufferAllocationSliceProto values

Why

BufferAllocation::Slice::FromProto reconstructs runtime buffer slices from serialized executable data. It previously validated only buffer_allocation_index, while the compile-time assignment path guarantees that offset >= 0, size >= 0, and offset + size <= allocation.size(). This change enforces that same invariant before constructing the slice.

Testing

git diff --check
clang++ -std=c++17 -fsanitize=address -g -DNDEBUG -DAPPLY_FIX poc_primitive_asan.cc -o /tmp/openxla_poc_fixed && /tmp/openxla_poc_fixed from the local reproduction bundle; the fixed logic rejects the malformed slice
attempted bazel test //xla/service:buffer_assignment_test; local sparse-checkout setup initially missed top-level third_party and tools, and after adding them the invocation stayed in Bazel repository mapping until interrupted. Retrying with --noexperimental_repo_remote_exec fails during repository initialization because rules_ml_toolchain requires that experimental flag.

Imported from GitHub PR #44653 ## Summary - validate `BufferAllocation::Slice::FromProto` rejects negative offsets and sizes - reject serialized slices whose `[offset, offset + size)` range exceeds the referenced allocation - add regression coverage for malformed `BufferAllocationSliceProto` values ## Why `BufferAllocation::Slice::FromProto` reconstructs runtime buffer slices from serialized executable data. It previously validated only `buffer_allocation_index`, while the compile-time assignment path guarantees that `offset >= 0`, `size >= 0`, and `offset + size <= allocation.size()`. This change enforces that same invariant before constructing the slice. ## Testing - `git diff --check` - `clang++ -std=c++17 -fsanitize=address -g -DNDEBUG -DAPPLY_FIX poc_primitive_asan.cc -o /tmp/openxla_poc_fixed && /tmp/openxla_poc_fixed` from the local reproduction bundle; the fixed logic rejects the malformed slice - attempted `bazel test //xla/service:buffer_assignment_test`; local sparse-checkout setup initially missed top-level `third_party` and `tools`, and after adding them the invocation stayed in Bazel repository mapping until interrupted. Retrying with `--noexperimental_repo_remote_exec` fails during repository initialization because `rules_ml_toolchain` requires that experimental flag. Copybara import of the project: -- 67c6ebd by M0nd0R <[email protected]>: Validate buffer slices during deserialization Merging this change closes #44653 FUTURE_COPYBARA_INTEGRATE_REVIEW=#44653 from M0nd0R:codex/validate-buffer-slice-from-proto 67c6ebd PiperOrigin-RevId: 950753890

Imported from GitHub PR openxla/xla#44653 ## Summary - validate `BufferAllocation::Slice::FromProto` rejects negative offsets and sizes - reject serialized slices whose `[offset, offset + size)` range exceeds the referenced allocation - add regression coverage for malformed `BufferAllocationSliceProto` values ## Why `BufferAllocation::Slice::FromProto` reconstructs runtime buffer slices from serialized executable data. It previously validated only `buffer_allocation_index`, while the compile-time assignment path guarantees that `offset >= 0`, `size >= 0`, and `offset + size <= allocation.size()`. This change enforces that same invariant before constructing the slice. ## Testing - `git diff --check` - `clang++ -std=c++17 -fsanitize=address -g -DNDEBUG -DAPPLY_FIX poc_primitive_asan.cc -o /tmp/openxla_poc_fixed && /tmp/openxla_poc_fixed` from the local reproduction bundle; the fixed logic rejects the malformed slice - attempted `bazel test //xla/service:buffer_assignment_test`; local sparse-checkout setup initially missed top-level `third_party` and `tools`, and after adding them the invocation stayed in Bazel repository mapping until interrupted. Retrying with `--noexperimental_repo_remote_exec` fails during repository initialization because `rules_ml_toolchain` requires that experimental flag. Copybara import of the project: -- 67c6ebdc5932ccdf4fb97c98a6e21b6656f0c08a by M0nd0R <[email protected]>: Validate buffer slices during deserialization Merging this change closes #44653 Reverts changelist 949495726 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#44653 from M0nd0R:codex/validate-buffer-slice-from-proto 67c6ebdc5932ccdf4fb97c98a6e21b6656f0c08a PiperOrigin-RevId: 950753890

Imported from GitHub PR #44653 ## Summary - validate `BufferAllocation::Slice::FromProto` rejects negative offsets and sizes - reject serialized slices whose `[offset, offset + size)` range exceeds the referenced allocation - add regression coverage for malformed `BufferAllocationSliceProto` values ## Why `BufferAllocation::Slice::FromProto` reconstructs runtime buffer slices from serialized executable data. It previously validated only `buffer_allocation_index`, while the compile-time assignment path guarantees that `offset >= 0`, `size >= 0`, and `offset + size <= allocation.size()`. This change enforces that same invariant before constructing the slice. ## Testing - `git diff --check` - `clang++ -std=c++17 -fsanitize=address -g -DNDEBUG -DAPPLY_FIX poc_primitive_asan.cc -o /tmp/openxla_poc_fixed && /tmp/openxla_poc_fixed` from the local reproduction bundle; the fixed logic rejects the malformed slice - attempted `bazel test //xla/service:buffer_assignment_test`; local sparse-checkout setup initially missed top-level `third_party` and `tools`, and after adding them the invocation stayed in Bazel repository mapping until interrupted. Retrying with `--noexperimental_repo_remote_exec` fails during repository initialization because `rules_ml_toolchain` requires that experimental flag. Copybara import of the project: -- 67c6ebd by M0nd0R <[email protected]>: Validate buffer slices during deserialization Merging this change closes #44653 FUTURE_COPYBARA_INTEGRATE_REVIEW=#44653 from M0nd0R:codex/validate-buffer-slice-from-proto 67c6ebd PiperOrigin-RevId: 950753890

Imported from GitHub PR openxla/xla#44653 ## Summary - validate `BufferAllocation::Slice::FromProto` rejects negative offsets and sizes - reject serialized slices whose `[offset, offset + size)` range exceeds the referenced allocation - add regression coverage for malformed `BufferAllocationSliceProto` values ## Why `BufferAllocation::Slice::FromProto` reconstructs runtime buffer slices from serialized executable data. It previously validated only `buffer_allocation_index`, while the compile-time assignment path guarantees that `offset >= 0`, `size >= 0`, and `offset + size <= allocation.size()`. This change enforces that same invariant before constructing the slice. ## Testing - `git diff --check` - `clang++ -std=c++17 -fsanitize=address -g -DNDEBUG -DAPPLY_FIX poc_primitive_asan.cc -o /tmp/openxla_poc_fixed && /tmp/openxla_poc_fixed` from the local reproduction bundle; the fixed logic rejects the malformed slice - attempted `bazel test //xla/service:buffer_assignment_test`; local sparse-checkout setup initially missed top-level `third_party` and `tools`, and after adding them the invocation stayed in Bazel repository mapping until interrupted. Retrying with `--noexperimental_repo_remote_exec` fails during repository initialization because `rules_ml_toolchain` requires that experimental flag. Copybara import of the project: -- 67c6ebdc5932ccdf4fb97c98a6e21b6656f0c08a by M0nd0R <[email protected]>: Validate buffer slices during deserialization Merging this change closes #44653 Reverts changelist 949495726 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#44653 from M0nd0R:codex/validate-buffer-slice-from-proto 67c6ebdc5932ccdf4fb97c98a6e21b6656f0c08a PiperOrigin-RevId: 950753890

Imported from GitHub PR #44653 ## Summary - validate `BufferAllocation::Slice::FromProto` rejects negative offsets and sizes - reject serialized slices whose `[offset, offset + size)` range exceeds the referenced allocation - add regression coverage for malformed `BufferAllocationSliceProto` values ## Why `BufferAllocation::Slice::FromProto` reconstructs runtime buffer slices from serialized executable data. It previously validated only `buffer_allocation_index`, while the compile-time assignment path guarantees that `offset >= 0`, `size >= 0`, and `offset + size <= allocation.size()`. This change enforces that same invariant before constructing the slice. ## Testing - `git diff --check` - `clang++ -std=c++17 -fsanitize=address -g -DNDEBUG -DAPPLY_FIX poc_primitive_asan.cc -o /tmp/openxla_poc_fixed && /tmp/openxla_poc_fixed` from the local reproduction bundle; the fixed logic rejects the malformed slice - attempted `bazel test //xla/service:buffer_assignment_test`; local sparse-checkout setup initially missed top-level `third_party` and `tools`, and after adding them the invocation stayed in Bazel repository mapping until interrupted. Retrying with `--noexperimental_repo_remote_exec` fails during repository initialization because `rules_ml_toolchain` requires that experimental flag. Copybara import of the project: -- 67c6ebd by M0nd0R <[email protected]>: Validate buffer slices during deserialization Merging this change closes #44653 FUTURE_COPYBARA_INTEGRATE_REVIEW=#44653 from M0nd0R:codex/validate-buffer-slice-from-proto 67c6ebd PiperOrigin-RevId: 950753890

Imported from GitHub PR openxla/xla#44653 ## Summary - validate `BufferAllocation::Slice::FromProto` rejects negative offsets and sizes - reject serialized slices whose `[offset, offset + size)` range exceeds the referenced allocation - add regression coverage for malformed `BufferAllocationSliceProto` values ## Why `BufferAllocation::Slice::FromProto` reconstructs runtime buffer slices from serialized executable data. It previously validated only `buffer_allocation_index`, while the compile-time assignment path guarantees that `offset >= 0`, `size >= 0`, and `offset + size <= allocation.size()`. This change enforces that same invariant before constructing the slice. ## Testing - `git diff --check` - `clang++ -std=c++17 -fsanitize=address -g -DNDEBUG -DAPPLY_FIX poc_primitive_asan.cc -o /tmp/openxla_poc_fixed && /tmp/openxla_poc_fixed` from the local reproduction bundle; the fixed logic rejects the malformed slice - attempted `bazel test //xla/service:buffer_assignment_test`; local sparse-checkout setup initially missed top-level `third_party` and `tools`, and after adding them the invocation stayed in Bazel repository mapping until interrupted. Retrying with `--noexperimental_repo_remote_exec` fails during repository initialization because `rules_ml_toolchain` requires that experimental flag. Copybara import of the project: -- 67c6ebdc5932ccdf4fb97c98a6e21b6656f0c08a by M0nd0R <[email protected]>: Validate buffer slices during deserialization Merging this change closes #44653 Reverts changelist 949495726 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#44653 from M0nd0R:codex/validate-buffer-slice-from-proto 67c6ebdc5932ccdf4fb97c98a6e21b6656f0c08a PiperOrigin-RevId: 950753890

akuegel · 2026-07-21T11:03:21Z

It looks like this change is causing test failures in xla/backends/gpu/runtime/dynamic_slice_thunk_test.cc and xla/backends/gpu/runtime/gpublas_lt_matmul_thunk_test.cc

Also in //xla/backends/gpu/runtime:thunk_proto_deserialization_test: https://github.com/openxla/xla/actions/runs/29749387329/job/88375566893

@M0nd0R can you please take a look?

M0nd0R · 2026-07-22T02:19:46Z

Thanks for the heads-up — looked into the failures.

Root cause: the new FromProto bounds check is correct; several existing test fixtures had allocation sizes that never matched the slices they deserialize (they only worked before because FromProto checked the index, not offset/size).

What was failing

ThunkProtoDeserializationTest.CublasLtGroupedMatmulThunk / GpuBlasLtMatmulThunkTest.ThunkProtoSerializationGroupedMatmul
Slice for allocation #4 is size=1302400 (f32[2,407,400]), but the fixture allocated #4 as 651200 (the non-grouped f32[407,400] size).
DynamicSliceThunkTest.*ProtoRoundTrip
CheckProtoRoundTrip hard-coded every allocation to 1024, so the 1MiB workspace slices (size=1048576) failed the new check.

Fix (pushed):

bump grouped-matmul allocation #4 to 1302400 in both matmul proto tests
size CheckProtoRoundTrip allocations from the thunk’s argument/offset slices so they always cover what FromProto reconstructs

Please re-run the affected targets when convenient:

//xla/backends/gpu/runtime:thunk_proto_deserialization_test
//xla/backends/gpu/runtime:gpublas_lt_matmul_thunk_test
//xla/backends/gpu/runtime:dynamic_slice_thunk_test

Grouped-matmul fixtures used a B allocation sized for the non-grouped shape, and DynamicSliceThunk round-trips hard-coded 1KiB allocations that cannot hold the 1MiB workspace slices once FromProto validates ranges.

Imported from GitHub PR #44653 ## Summary - validate `BufferAllocation::Slice::FromProto` rejects negative offsets and sizes - reject serialized slices whose `[offset, offset + size)` range exceeds the referenced allocation - add regression coverage for malformed `BufferAllocationSliceProto` values ## Why `BufferAllocation::Slice::FromProto` reconstructs runtime buffer slices from serialized executable data. It previously validated only `buffer_allocation_index`, while the compile-time assignment path guarantees that `offset >= 0`, `size >= 0`, and `offset + size <= allocation.size()`. This change enforces that same invariant before constructing the slice. ## Testing - `git diff --check` - `clang++ -std=c++17 -fsanitize=address -g -DNDEBUG -DAPPLY_FIX poc_primitive_asan.cc -o /tmp/openxla_poc_fixed && /tmp/openxla_poc_fixed` from the local reproduction bundle; the fixed logic rejects the malformed slice - attempted `bazel test //xla/service:buffer_assignment_test`; local sparse-checkout setup initially missed top-level `third_party` and `tools`, and after adding them the invocation stayed in Bazel repository mapping until interrupted. Retrying with `--noexperimental_repo_remote_exec` fails during repository initialization because `rules_ml_toolchain` requires that experimental flag. Copybara import of the project: -- 67c6ebd by M0nd0R <[email protected]>: Validate buffer slices during deserialization -- 5ce256b by mundur <[email protected]>: Fix deserialization tests for slice bounds checks Grouped-matmul fixtures used a B allocation sized for the non-grouped shape, and DynamicSliceThunk round-trips hard-coded 1KiB allocations that cannot hold the 1MiB workspace slices once FromProto validates ranges. Merging this change closes #44653 FUTURE_COPYBARA_INTEGRATE_REVIEW=#44653 from M0nd0R:codex/validate-buffer-slice-from-proto 5ce256b PiperOrigin-RevId: 950753890

Imported from GitHub PR openxla/xla#44653 ## Summary - validate `BufferAllocation::Slice::FromProto` rejects negative offsets and sizes - reject serialized slices whose `[offset, offset + size)` range exceeds the referenced allocation - add regression coverage for malformed `BufferAllocationSliceProto` values ## Why `BufferAllocation::Slice::FromProto` reconstructs runtime buffer slices from serialized executable data. It previously validated only `buffer_allocation_index`, while the compile-time assignment path guarantees that `offset >= 0`, `size >= 0`, and `offset + size <= allocation.size()`. This change enforces that same invariant before constructing the slice. ## Testing - `git diff --check` - `clang++ -std=c++17 -fsanitize=address -g -DNDEBUG -DAPPLY_FIX poc_primitive_asan.cc -o /tmp/openxla_poc_fixed && /tmp/openxla_poc_fixed` from the local reproduction bundle; the fixed logic rejects the malformed slice - attempted `bazel test //xla/service:buffer_assignment_test`; local sparse-checkout setup initially missed top-level `third_party` and `tools`, and after adding them the invocation stayed in Bazel repository mapping until interrupted. Retrying with `--noexperimental_repo_remote_exec` fails during repository initialization because `rules_ml_toolchain` requires that experimental flag. Copybara import of the project: -- 67c6ebdc5932ccdf4fb97c98a6e21b6656f0c08a by M0nd0R <[email protected]>: Validate buffer slices during deserialization -- 5ce256b2ca289ca3bb01e537025a90aea4f9361c by mundur <[email protected]>: Fix deserialization tests for slice bounds checks Grouped-matmul fixtures used a B allocation sized for the non-grouped shape, and DynamicSliceThunk round-trips hard-coded 1KiB allocations that cannot hold the 1MiB workspace slices once FromProto validates ranges. Merging this change closes #44653 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#44653 from M0nd0R:codex/validate-buffer-slice-from-proto 5ce256b2ca289ca3bb01e537025a90aea4f9361c PiperOrigin-RevId: 950753890

Imported from GitHub PR #44653 ## Summary - validate `BufferAllocation::Slice::FromProto` rejects negative offsets and sizes - reject serialized slices whose `[offset, offset + size)` range exceeds the referenced allocation - add regression coverage for malformed `BufferAllocationSliceProto` values ## Why `BufferAllocation::Slice::FromProto` reconstructs runtime buffer slices from serialized executable data. It previously validated only `buffer_allocation_index`, while the compile-time assignment path guarantees that `offset >= 0`, `size >= 0`, and `offset + size <= allocation.size()`. This change enforces that same invariant before constructing the slice. ## Testing - `git diff --check` - `clang++ -std=c++17 -fsanitize=address -g -DNDEBUG -DAPPLY_FIX poc_primitive_asan.cc -o /tmp/openxla_poc_fixed && /tmp/openxla_poc_fixed` from the local reproduction bundle; the fixed logic rejects the malformed slice - attempted `bazel test //xla/service:buffer_assignment_test`; local sparse-checkout setup initially missed top-level `third_party` and `tools`, and after adding them the invocation stayed in Bazel repository mapping until interrupted. Retrying with `--noexperimental_repo_remote_exec` fails during repository initialization because `rules_ml_toolchain` requires that experimental flag. Copybara import of the project: -- 67c6ebd by M0nd0R <[email protected]>: Validate buffer slices during deserialization -- 5ce256b by mundur <[email protected]>: Fix deserialization tests for slice bounds checks Grouped-matmul fixtures used a B allocation sized for the non-grouped shape, and DynamicSliceThunk round-trips hard-coded 1KiB allocations that cannot hold the 1MiB workspace slices once FromProto validates ranges. Merging this change closes #44653 FUTURE_COPYBARA_INTEGRATE_REVIEW=#44653 from M0nd0R:codex/validate-buffer-slice-from-proto 5ce256b PiperOrigin-RevId: 950753890

Imported from GitHub PR openxla/xla#44653 ## Summary - validate `BufferAllocation::Slice::FromProto` rejects negative offsets and sizes - reject serialized slices whose `[offset, offset + size)` range exceeds the referenced allocation - add regression coverage for malformed `BufferAllocationSliceProto` values ## Why `BufferAllocation::Slice::FromProto` reconstructs runtime buffer slices from serialized executable data. It previously validated only `buffer_allocation_index`, while the compile-time assignment path guarantees that `offset >= 0`, `size >= 0`, and `offset + size <= allocation.size()`. This change enforces that same invariant before constructing the slice. ## Testing - `git diff --check` - `clang++ -std=c++17 -fsanitize=address -g -DNDEBUG -DAPPLY_FIX poc_primitive_asan.cc -o /tmp/openxla_poc_fixed && /tmp/openxla_poc_fixed` from the local reproduction bundle; the fixed logic rejects the malformed slice - attempted `bazel test //xla/service:buffer_assignment_test`; local sparse-checkout setup initially missed top-level `third_party` and `tools`, and after adding them the invocation stayed in Bazel repository mapping until interrupted. Retrying with `--noexperimental_repo_remote_exec` fails during repository initialization because `rules_ml_toolchain` requires that experimental flag. Copybara import of the project: -- 67c6ebdc5932ccdf4fb97c98a6e21b6656f0c08a by M0nd0R <[email protected]>: Validate buffer slices during deserialization -- 5ce256b2ca289ca3bb01e537025a90aea4f9361c by mundur <[email protected]>: Fix deserialization tests for slice bounds checks Grouped-matmul fixtures used a B allocation sized for the non-grouped shape, and DynamicSliceThunk round-trips hard-coded 1KiB allocations that cannot hold the 1MiB workspace slices once FromProto validates ranges. Merging this change closes #44653 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#44653 from M0nd0R:codex/validate-buffer-slice-from-proto 5ce256b2ca289ca3bb01e537025a90aea4f9361c PiperOrigin-RevId: 950753890

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#44653 from M0nd0R:codex/validate-buffer-slice-from-proto 5ce256b2ca289ca3bb01e537025a90aea4f9361c PiperOrigin-RevId: 951172456

Imported from GitHub PR openxla/xla#44653 ## Summary - validate `BufferAllocation::Slice::FromProto` rejects negative offsets and sizes - reject serialized slices whose `[offset, offset + size)` range exceeds the referenced allocation - add regression coverage for malformed `BufferAllocationSliceProto` values ## Why `BufferAllocation::Slice::FromProto` reconstructs runtime buffer slices from serialized executable data. It previously validated only `buffer_allocation_index`, while the compile-time assignment path guarantees that `offset >= 0`, `size >= 0`, and `offset + size <= allocation.size()`. This change enforces that same invariant before constructing the slice. ## Testing - `git diff --check` - `clang++ -std=c++17 -fsanitize=address -g -DNDEBUG -DAPPLY_FIX poc_primitive_asan.cc -o /tmp/openxla_poc_fixed && /tmp/openxla_poc_fixed` from the local reproduction bundle; the fixed logic rejects the malformed slice - attempted `bazel test //xla/service:buffer_assignment_test`; local sparse-checkout setup initially missed top-level `third_party` and `tools`, and after adding them the invocation stayed in Bazel repository mapping until interrupted. Retrying with `--noexperimental_repo_remote_exec` fails during repository initialization because `rules_ml_toolchain` requires that experimental flag. Copybara import of the project: -- 67c6ebdc5932ccdf4fb97c98a6e21b6656f0c08a by M0nd0R <[email protected]>: Validate buffer slices during deserialization -- 5ce256b2ca289ca3bb01e537025a90aea4f9361c by mundur <[email protected]>: Fix deserialization tests for slice bounds checks Grouped-matmul fixtures used a B allocation sized for the non-grouped shape, and DynamicSliceThunk round-trips hard-coded 1KiB allocations that cannot hold the 1MiB workspace slices once FromProto validates ranges. Merging this change closes #44653 PiperOrigin-RevId: 952025970

…perform direct `std::copy` (`memcpy`) when `shape().layout()` is monotonic with dim 0 major. FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#44653 from M0nd0R:codex/validate-buffer-slice-from-proto 5ce256b2ca289ca3bb01e537025a90aea4f9361c PiperOrigin-RevId: 933959965

…roto BufferAssignment::FromProto passes assignee.offset() and assignee.size() from a deserialized BufferAllocationProto.Assigned proto directly to BufferAllocation::AddAssignment without checking for negative values. The CHECK_LE guards in AddAssignment are bypassed by negative int64 values: CHECK_LE(-1, 1024) = true (for any real allocation) CHECK_LE(-1 + 4, 1024) = true So a negative offset is stored in assigned_buffers_ and later used in LLVM IR codegen (ir_emitter.cc) as: tempbuf_address_base + b()->getInt64(slice.offset()) With offset=-1 this is an out-of-bounds address. This mirrors the fix applied to Slice::FromProto in PR openxla#44653 (2840cd3), which fixed the thunk deserialization path but missed this sibling path in BufferAssignment::FromProto. Fix: add an explicit negative check before calling AddAssignment, returning absl::OutOfRangeError for negative offset or size values. Add two regression tests confirming the bypass behavior and documenting the required fix behavior.

…ment::FromProto Imported from GitHub PR #46105 ## Problem `BufferAssignment::FromProto` passes `assignee.offset()` and `assignee.size()` from a deserialized `BufferAllocationProto.Assigned` proto directly to `BufferAllocation::AddAssignment` without checking for negative values. The guards inside `AddAssignment` are: ```cpp CHECK_LE(offset, size_) // checks offset <= allocation.size() CHECK_LE(offset + size, size_) // checks offset + size <= allocation.size() ``` These `CHECK_LE` guards are **bypassed by negative int64 values**: `-1 <= 1024` is true for any real allocation, so a negative offset passes both checks and gets stored in `assigned_buffers_`. The stored negative offset is later used in LLVM IR codegen (`ir_emitter.cc`): ```cpp tempbuf_address_base + b()->getInt64(slice.offset()) ``` With `offset = -1`, this produces a pointer one byte before the allocation — an out-of-bounds address. ## Attack Surface Reachable from `CpuAotLoader::LoadAotCompilationResult(serialized_aot_result)` where `serialized_aot_result` is a user-supplied serialized AOT executable. The proto fields `offset` and `size` are both `int64` in `hlo.proto`, so negative values are fully valid at the wire level. ## Relationship to PR #44653 PR #44653 (`2840cd3`) fixed `Slice::FromProto` with: ```cpp if (proto.offset() < 0 || proto.size() < 0) { return absl::OutOfRangeError(...); } ``` This PR applies the **same fix** to `BufferAssignment::FromProto` → `AddAssignment` path, which handles `Assigned` buffers in `BufferAllocationProto` — the sibling that the previous fix missed. ## Fix Added a pre-validation check in `BufferAssignment::FromProto` before calling `AddAssignment`: ```cpp if (assignee.offset() < 0 || assignee.size() < 0) { return absl::OutOfRangeError(absl::StrCat( "BufferAssignment proto contains an assigned buffer with negative ", "offset or size: logical_buffer_id=", logical_buffer_id, " offset=", assignee.offset(), " size=", assignee.size())); } ``` ## Testing Two new test cases in `buffer_assignment_test.cc`: - `BufferAssignmentFromProtoTest.RejectsNegativeOffset`: confirms `offset=-1` bypasses `CHECK_LE` guards (proving the vulnerability) and documents that the fix must reject it before reaching `AddAssignment` - `BufferAssignmentFromProtoTest.RejectsNegativeSize`: same for `size=-1` --- *Disclosure: this contribution was authored with the aid of an AI coding assistant and reviewed before submission.* Copybara import of the project: -- 6fab2ec by prasanna8585 <[email protected]>: [XLA] Validate assigned buffer offset/size in BufferAssignment::FromProto BufferAssignment::FromProto passes assignee.offset() and assignee.size() from a deserialized BufferAllocationProto.Assigned proto directly to BufferAllocation::AddAssignment without checking for negative values. The CHECK_LE guards in AddAssignment are bypassed by negative int64 values: CHECK_LE(-1, 1024) = true (for any real allocation) CHECK_LE(-1 + 4, 1024) = true So a negative offset is stored in assigned_buffers_ and later used in LLVM IR codegen (ir_emitter.cc) as: tempbuf_address_base + b()->getInt64(slice.offset()) With offset=-1 this is an out-of-bounds address. This mirrors the fix applied to Slice::FromProto in PR #44653 (2840cd3), which fixed the thunk deserialization path but missed this sibling path in BufferAssignment::FromProto. Fix: add an explicit negative check before calling AddAssignment, returning absl::OutOfRangeError for negative offset or size values. Add two regression tests confirming the bypass behavior and documenting the required fix behavior. Merging this change closes #46105 FUTURE_COPYBARA_INTEGRATE_REVIEW=#46105 from prasanna8585:fix/buffer-assignment-from-proto-negative-offset 6fab2ec PiperOrigin-RevId: 952638612

@beckerhe

…:AddAssignment Imported from GitHub PR #46105 ## Problem `BufferAssignment::FromProto` passes `assignee.offset()` and `assignee.size()` from a deserialized `BufferAllocationProto.Assigned` proto directly to `BufferAllocation::AddAssignment` without checking for negative values. The guards inside `AddAssignment` are: ```cpp CHECK_LE(offset, size_) // checks offset <= allocation.size() CHECK_LE(offset + size, size_) // checks offset + size <= allocation.size() ``` These `CHECK_LE` guards are **bypassed by negative int64 values**: `-1 <= 1024` is true for any real allocation, so a negative offset passes both checks and gets stored in `assigned_buffers_`. The stored negative offset is later used in LLVM IR codegen (`ir_emitter.cc`): ```cpp tempbuf_address_base + b()->getInt64(slice.offset()) ``` With `offset = -1`, this produces a pointer one byte before the allocation — an out-of-bounds address. ## Attack Surface Reachable from `CpuAotLoader::LoadAotCompilationResult(serialized_aot_result)` where `serialized_aot_result` is a user-supplied serialized AOT executable. The proto fields `offset` and `size` are both `int64` in `hlo.proto`, so negative values are fully valid at the wire level. ## Relationship to PR #44653 PR #44653 (`2840cd3`) fixed `Slice::FromProto` with: ```cpp if (proto.offset() < 0 || proto.size() < 0) { return absl::OutOfRangeError(...); } ``` This PR applies the **same fix** to `BufferAssignment::FromProto` → `AddAssignment` path, which handles `Assigned` buffers in `BufferAllocationProto` — the sibling that the previous fix missed. ## Fix Added a pre-validation check in `BufferAssignment::FromProto` before calling `AddAssignment`: ```cpp if (assignee.offset() < 0 || assignee.size() < 0) { return absl::OutOfRangeError(absl::StrCat( "BufferAssignment proto contains an assigned buffer with negative ", "offset or size: logical_buffer_id=", logical_buffer_id, " offset=", assignee.offset(), " size=", assignee.size())); } ``` ## Testing Two new test cases in `buffer_assignment_test.cc`: - `BufferAssignmentFromProtoTest.RejectsNegativeOffset`: confirms `offset=-1` bypasses `CHECK_LE` guards (proving the vulnerability) and documents that the fix must reject it before reaching `AddAssignment` - `BufferAssignmentFromProtoTest.RejectsNegativeSize`: same for `size=-1` --- *Disclosure: this contribution was authored with the aid of an AI coding assistant and reviewed before submission.* Copybara import of the project: -- 6fab2ec by prasanna8585 <[email protected]>: [XLA] Validate assigned buffer offset/size in BufferAssignment::FromProto BufferAssignment::FromProto passes assignee.offset() and assignee.size() from a deserialized BufferAllocationProto.Assigned proto directly to BufferAllocation::AddAssignment without checking for negative values. The CHECK_LE guards in AddAssignment are bypassed by negative int64 values: CHECK_LE(-1, 1024) = true (for any real allocation) CHECK_LE(-1 + 4, 1024) = true So a negative offset is stored in assigned_buffers_ and later used in LLVM IR codegen (ir_emitter.cc) as: tempbuf_address_base + b()->getInt64(slice.offset()) With offset=-1 this is an out-of-bounds address. This mirrors the fix applied to Slice::FromProto in PR #44653 (2840cd3), which fixed the thunk deserialization path but missed this sibling path in BufferAssignment::FromProto. Fix: add an explicit negative check before calling AddAssignment, returning absl::OutOfRangeError for negative offset or size values. Add two regression tests confirming the bypass behavior and documenting the required fix behavior. -- 5fdce18 by prasanna8585 <[email protected]>: Address review: move fix from FromProto to BufferAllocation::AddAssignment Per review feedback from @beckerhe: move the negative offset/size validation from a precondition check in BufferAssignment::FromProto into BufferAllocation::AddAssignment itself, converting the existing CHECK_LE statements to TF_RET_CHECK. This protects every caller of AddAssignment (there are ~9 call sites across buffer_assignment.cc), not just the FromProto deserialization path, and returns an absl::Status error instead of crashing the process when the precondition is violated. Also adds explicit offset >= 0 / size >= 0 checks: converting CHECK_LE(offset, size_) to TF_RET_CHECK alone would not catch a negative offset, since offset <= size_ is still true for any negative offset and non-negative allocation size. All 9 existing call sites already propagate the returned Status via RETURN_IF_ERROR, so no call-site changes were needed beyond removing the now-redundant check in FromProto. Merging this change closes #46105 FUTURE_COPYBARA_INTEGRATE_REVIEW=#46105 from prasanna8585:fix/buffer-assignment-from-proto-negative-offset 5fdce18 PiperOrigin-RevId: 952638612

@beckerhe

…:AddAssignment Imported from GitHub PR #46105 ## Problem `BufferAssignment::FromProto` passes `assignee.offset()` and `assignee.size()` from a deserialized `BufferAllocationProto.Assigned` proto directly to `BufferAllocation::AddAssignment` without checking for negative values. The guards inside `AddAssignment` are: ```cpp CHECK_LE(offset, size_) // checks offset <= allocation.size() CHECK_LE(offset + size, size_) // checks offset + size <= allocation.size() ``` These `CHECK_LE` guards are **bypassed by negative int64 values**: `-1 <= 1024` is true for any real allocation, so a negative offset passes both checks and gets stored in `assigned_buffers_`. The stored negative offset is later used in LLVM IR codegen (`ir_emitter.cc`): ```cpp tempbuf_address_base + b()->getInt64(slice.offset()) ``` With `offset = -1`, this produces a pointer one byte before the allocation — an out-of-bounds address. ## Attack Surface Reachable from `CpuAotLoader::LoadAotCompilationResult(serialized_aot_result)` where `serialized_aot_result` is a user-supplied serialized AOT executable. The proto fields `offset` and `size` are both `int64` in `hlo.proto`, so negative values are fully valid at the wire level. ## Relationship to PR #44653 PR #44653 (`2840cd3`) fixed `Slice::FromProto` with: ```cpp if (proto.offset() < 0 || proto.size() < 0) { return absl::OutOfRangeError(...); } ``` This PR applies the **same fix** to `BufferAssignment::FromProto` → `AddAssignment` path, which handles `Assigned` buffers in `BufferAllocationProto` — the sibling that the previous fix missed. ## Fix Added a pre-validation check in `BufferAssignment::FromProto` before calling `AddAssignment`: ```cpp if (assignee.offset() < 0 || assignee.size() < 0) { return absl::OutOfRangeError(absl::StrCat( "BufferAssignment proto contains an assigned buffer with negative ", "offset or size: logical_buffer_id=", logical_buffer_id, " offset=", assignee.offset(), " size=", assignee.size())); } ``` ## Testing Two new test cases in `buffer_assignment_test.cc`: - `BufferAssignmentFromProtoTest.RejectsNegativeOffset`: confirms `offset=-1` bypasses `CHECK_LE` guards (proving the vulnerability) and documents that the fix must reject it before reaching `AddAssignment` - `BufferAssignmentFromProtoTest.RejectsNegativeSize`: same for `size=-1` --- *Disclosure: this contribution was authored with the aid of an AI coding assistant and reviewed before submission.* Copybara import of the project: -- 6fab2ec by prasanna8585 <[email protected]>: [XLA] Validate assigned buffer offset/size in BufferAssignment::FromProto BufferAssignment::FromProto passes assignee.offset() and assignee.size() from a deserialized BufferAllocationProto.Assigned proto directly to BufferAllocation::AddAssignment without checking for negative values. The CHECK_LE guards in AddAssignment are bypassed by negative int64 values: CHECK_LE(-1, 1024) = true (for any real allocation) CHECK_LE(-1 + 4, 1024) = true So a negative offset is stored in assigned_buffers_ and later used in LLVM IR codegen (ir_emitter.cc) as: tempbuf_address_base + b()->getInt64(slice.offset()) With offset=-1 this is an out-of-bounds address. This mirrors the fix applied to Slice::FromProto in PR #44653 (2840cd3), which fixed the thunk deserialization path but missed this sibling path in BufferAssignment::FromProto. Fix: add an explicit negative check before calling AddAssignment, returning absl::OutOfRangeError for negative offset or size values. Add two regression tests confirming the bypass behavior and documenting the required fix behavior. -- 5fdce18 by prasanna8585 <[email protected]>: Address review: move fix from FromProto to BufferAllocation::AddAssignment Per review feedback from @beckerhe: move the negative offset/size validation from a precondition check in BufferAssignment::FromProto into BufferAllocation::AddAssignment itself, converting the existing CHECK_LE statements to TF_RET_CHECK. This protects every caller of AddAssignment (there are ~9 call sites across buffer_assignment.cc), not just the FromProto deserialization path, and returns an absl::Status error instead of crashing the process when the precondition is violated. Also adds explicit offset >= 0 / size >= 0 checks: converting CHECK_LE(offset, size_) to TF_RET_CHECK alone would not catch a negative offset, since offset <= size_ is still true for any negative offset and non-negative allocation size. All 9 existing call sites already propagate the returned Status via RETURN_IF_ERROR, so no call-site changes were needed beyond removing the now-redundant check in FromProto. Merging this change closes #46105 FUTURE_COPYBARA_INTEGRATE_REVIEW=#46105 from prasanna8585:fix/buffer-assignment-from-proto-negative-offset 5fdce18 PiperOrigin-RevId: 952638612

@beckerhe

…:AddAssignment Imported from GitHub PR #46105 ## Problem `BufferAssignment::FromProto` passes `assignee.offset()` and `assignee.size()` from a deserialized `BufferAllocationProto.Assigned` proto directly to `BufferAllocation::AddAssignment` without checking for negative values. The guards inside `AddAssignment` are: ```cpp CHECK_LE(offset, size_) // checks offset <= allocation.size() CHECK_LE(offset + size, size_) // checks offset + size <= allocation.size() ``` These `CHECK_LE` guards are **bypassed by negative int64 values**: `-1 <= 1024` is true for any real allocation, so a negative offset passes both checks and gets stored in `assigned_buffers_`. The stored negative offset is later used in LLVM IR codegen (`ir_emitter.cc`): ```cpp tempbuf_address_base + b()->getInt64(slice.offset()) ``` With `offset = -1`, this produces a pointer one byte before the allocation — an out-of-bounds address. ## Attack Surface Reachable from `CpuAotLoader::LoadAotCompilationResult(serialized_aot_result)` where `serialized_aot_result` is a user-supplied serialized AOT executable. The proto fields `offset` and `size` are both `int64` in `hlo.proto`, so negative values are fully valid at the wire level. ## Relationship to PR #44653 PR #44653 (`2840cd3`) fixed `Slice::FromProto` with: ```cpp if (proto.offset() < 0 || proto.size() < 0) { return absl::OutOfRangeError(...); } ``` This PR applies the **same fix** to `BufferAssignment::FromProto` → `AddAssignment` path, which handles `Assigned` buffers in `BufferAllocationProto` — the sibling that the previous fix missed. ## Fix Added a pre-validation check in `BufferAssignment::FromProto` before calling `AddAssignment`: ```cpp if (assignee.offset() < 0 || assignee.size() < 0) { return absl::OutOfRangeError(absl::StrCat( "BufferAssignment proto contains an assigned buffer with negative ", "offset or size: logical_buffer_id=", logical_buffer_id, " offset=", assignee.offset(), " size=", assignee.size())); } ``` ## Testing Two new test cases in `buffer_assignment_test.cc`: - `BufferAssignmentFromProtoTest.RejectsNegativeOffset`: confirms `offset=-1` bypasses `CHECK_LE` guards (proving the vulnerability) and documents that the fix must reject it before reaching `AddAssignment` - `BufferAssignmentFromProtoTest.RejectsNegativeSize`: same for `size=-1` --- *Disclosure: this contribution was authored with the aid of an AI coding assistant and reviewed before submission.* Copybara import of the project: -- 6fab2ec by prasanna8585 <[email protected]>: [XLA] Validate assigned buffer offset/size in BufferAssignment::FromProto BufferAssignment::FromProto passes assignee.offset() and assignee.size() from a deserialized BufferAllocationProto.Assigned proto directly to BufferAllocation::AddAssignment without checking for negative values. The CHECK_LE guards in AddAssignment are bypassed by negative int64 values: CHECK_LE(-1, 1024) = true (for any real allocation) CHECK_LE(-1 + 4, 1024) = true So a negative offset is stored in assigned_buffers_ and later used in LLVM IR codegen (ir_emitter.cc) as: tempbuf_address_base + b()->getInt64(slice.offset()) With offset=-1 this is an out-of-bounds address. This mirrors the fix applied to Slice::FromProto in PR #44653 (2840cd3), which fixed the thunk deserialization path but missed this sibling path in BufferAssignment::FromProto. Fix: add an explicit negative check before calling AddAssignment, returning absl::OutOfRangeError for negative offset or size values. Add two regression tests confirming the bypass behavior and documenting the required fix behavior. -- 5fdce18 by prasanna8585 <[email protected]>: Address review: move fix from FromProto to BufferAllocation::AddAssignment Per review feedback from @beckerhe: move the negative offset/size validation from a precondition check in BufferAssignment::FromProto into BufferAllocation::AddAssignment itself, converting the existing CHECK_LE statements to TF_RET_CHECK. This protects every caller of AddAssignment (there are ~9 call sites across buffer_assignment.cc), not just the FromProto deserialization path, and returns an absl::Status error instead of crashing the process when the precondition is violated. Also adds explicit offset >= 0 / size >= 0 checks: converting CHECK_LE(offset, size_) to TF_RET_CHECK alone would not catch a negative offset, since offset <= size_ is still true for any negative offset and non-negative allocation size. All 9 existing call sites already propagate the returned Status via RETURN_IF_ERROR, so no call-site changes were needed beyond removing the now-redundant check in FromProto. Merging this change closes #46105 FUTURE_COPYBARA_INTEGRATE_REVIEW=#46105 from prasanna8585:fix/buffer-assignment-from-proto-negative-offset 5fdce18 PiperOrigin-RevId: 952638612

@beckerhe

…:AddAssignment Imported from GitHub PR #46105 ## Problem `BufferAssignment::FromProto` passes `assignee.offset()` and `assignee.size()` from a deserialized `BufferAllocationProto.Assigned` proto directly to `BufferAllocation::AddAssignment` without checking for negative values. The guards inside `AddAssignment` are: ```cpp CHECK_LE(offset, size_) // checks offset <= allocation.size() CHECK_LE(offset + size, size_) // checks offset + size <= allocation.size() ``` These `CHECK_LE` guards are **bypassed by negative int64 values**: `-1 <= 1024` is true for any real allocation, so a negative offset passes both checks and gets stored in `assigned_buffers_`. The stored negative offset is later used in LLVM IR codegen (`ir_emitter.cc`): ```cpp tempbuf_address_base + b()->getInt64(slice.offset()) ``` With `offset = -1`, this produces a pointer one byte before the allocation — an out-of-bounds address. ## Attack Surface Reachable from `CpuAotLoader::LoadAotCompilationResult(serialized_aot_result)` where `serialized_aot_result` is a user-supplied serialized AOT executable. The proto fields `offset` and `size` are both `int64` in `hlo.proto`, so negative values are fully valid at the wire level. ## Relationship to PR #44653 PR #44653 (`2840cd3`) fixed `Slice::FromProto` with: ```cpp if (proto.offset() < 0 || proto.size() < 0) { return absl::OutOfRangeError(...); } ``` This PR applies the **same fix** to `BufferAssignment::FromProto` → `AddAssignment` path, which handles `Assigned` buffers in `BufferAllocationProto` — the sibling that the previous fix missed. ## Fix Added a pre-validation check in `BufferAssignment::FromProto` before calling `AddAssignment`: ```cpp if (assignee.offset() < 0 || assignee.size() < 0) { return absl::OutOfRangeError(absl::StrCat( "BufferAssignment proto contains an assigned buffer with negative ", "offset or size: logical_buffer_id=", logical_buffer_id, " offset=", assignee.offset(), " size=", assignee.size())); } ``` ## Testing Two new test cases in `buffer_assignment_test.cc`: - `BufferAssignmentFromProtoTest.RejectsNegativeOffset`: confirms `offset=-1` bypasses `CHECK_LE` guards (proving the vulnerability) and documents that the fix must reject it before reaching `AddAssignment` - `BufferAssignmentFromProtoTest.RejectsNegativeSize`: same for `size=-1` --- *Disclosure: this contribution was authored with the aid of an AI coding assistant and reviewed before submission.* Copybara import of the project: -- 6fab2ec by prasanna8585 <[email protected]>: [XLA] Validate assigned buffer offset/size in BufferAssignment::FromProto BufferAssignment::FromProto passes assignee.offset() and assignee.size() from a deserialized BufferAllocationProto.Assigned proto directly to BufferAllocation::AddAssignment without checking for negative values. The CHECK_LE guards in AddAssignment are bypassed by negative int64 values: CHECK_LE(-1, 1024) = true (for any real allocation) CHECK_LE(-1 + 4, 1024) = true So a negative offset is stored in assigned_buffers_ and later used in LLVM IR codegen (ir_emitter.cc) as: tempbuf_address_base + b()->getInt64(slice.offset()) With offset=-1 this is an out-of-bounds address. This mirrors the fix applied to Slice::FromProto in PR #44653 (2840cd3), which fixed the thunk deserialization path but missed this sibling path in BufferAssignment::FromProto. Fix: add an explicit negative check before calling AddAssignment, returning absl::OutOfRangeError for negative offset or size values. Add two regression tests confirming the bypass behavior and documenting the required fix behavior. -- 5fdce18 by prasanna8585 <[email protected]>: Address review: move fix from FromProto to BufferAllocation::AddAssignment Per review feedback from @beckerhe: move the negative offset/size validation from a precondition check in BufferAssignment::FromProto into BufferAllocation::AddAssignment itself, converting the existing CHECK_LE statements to TF_RET_CHECK. This protects every caller of AddAssignment (there are ~9 call sites across buffer_assignment.cc), not just the FromProto deserialization path, and returns an absl::Status error instead of crashing the process when the precondition is violated. Also adds explicit offset >= 0 / size >= 0 checks: converting CHECK_LE(offset, size_) to TF_RET_CHECK alone would not catch a negative offset, since offset <= size_ is still true for any negative offset and non-negative allocation size. All 9 existing call sites already propagate the returned Status via RETURN_IF_ERROR, so no call-site changes were needed beyond removing the now-redundant check in FromProto. Merging this change closes #46105 FUTURE_COPYBARA_INTEGRATE_REVIEW=#46105 from prasanna8585:fix/buffer-assignment-from-proto-negative-offset 5fdce18 PiperOrigin-RevId: 952638612

@beckerhe

…:AddAssignment Imported from GitHub PR #46105 ## Problem `BufferAssignment::FromProto` passes `assignee.offset()` and `assignee.size()` from a deserialized `BufferAllocationProto.Assigned` proto directly to `BufferAllocation::AddAssignment` without checking for negative values. The guards inside `AddAssignment` are: ```cpp CHECK_LE(offset, size_) // checks offset <= allocation.size() CHECK_LE(offset + size, size_) // checks offset + size <= allocation.size() ``` These `CHECK_LE` guards are **bypassed by negative int64 values**: `-1 <= 1024` is true for any real allocation, so a negative offset passes both checks and gets stored in `assigned_buffers_`. The stored negative offset is later used in LLVM IR codegen (`ir_emitter.cc`): ```cpp tempbuf_address_base + b()->getInt64(slice.offset()) ``` With `offset = -1`, this produces a pointer one byte before the allocation — an out-of-bounds address. ## Attack Surface Reachable from `CpuAotLoader::LoadAotCompilationResult(serialized_aot_result)` where `serialized_aot_result` is a user-supplied serialized AOT executable. The proto fields `offset` and `size` are both `int64` in `hlo.proto`, so negative values are fully valid at the wire level. ## Relationship to PR #44653 PR #44653 (`2840cd3`) fixed `Slice::FromProto` with: ```cpp if (proto.offset() < 0 || proto.size() < 0) { return absl::OutOfRangeError(...); } ``` This PR applies the **same fix** to `BufferAssignment::FromProto` → `AddAssignment` path, which handles `Assigned` buffers in `BufferAllocationProto` — the sibling that the previous fix missed. ## Fix Added a pre-validation check in `BufferAssignment::FromProto` before calling `AddAssignment`: ```cpp if (assignee.offset() < 0 || assignee.size() < 0) { return absl::OutOfRangeError(absl::StrCat( "BufferAssignment proto contains an assigned buffer with negative ", "offset or size: logical_buffer_id=", logical_buffer_id, " offset=", assignee.offset(), " size=", assignee.size())); } ``` ## Testing Two new test cases in `buffer_assignment_test.cc`: - `BufferAssignmentFromProtoTest.RejectsNegativeOffset`: confirms `offset=-1` bypasses `CHECK_LE` guards (proving the vulnerability) and documents that the fix must reject it before reaching `AddAssignment` - `BufferAssignmentFromProtoTest.RejectsNegativeSize`: same for `size=-1` --- *Disclosure: this contribution was authored with the aid of an AI coding assistant and reviewed before submission.* Copybara import of the project: -- 6fab2ec by prasanna8585 <[email protected]>: [XLA] Validate assigned buffer offset/size in BufferAssignment::FromProto BufferAssignment::FromProto passes assignee.offset() and assignee.size() from a deserialized BufferAllocationProto.Assigned proto directly to BufferAllocation::AddAssignment without checking for negative values. The CHECK_LE guards in AddAssignment are bypassed by negative int64 values: CHECK_LE(-1, 1024) = true (for any real allocation) CHECK_LE(-1 + 4, 1024) = true So a negative offset is stored in assigned_buffers_ and later used in LLVM IR codegen (ir_emitter.cc) as: tempbuf_address_base + b()->getInt64(slice.offset()) With offset=-1 this is an out-of-bounds address. This mirrors the fix applied to Slice::FromProto in PR #44653 (2840cd3), which fixed the thunk deserialization path but missed this sibling path in BufferAssignment::FromProto. Fix: add an explicit negative check before calling AddAssignment, returning absl::OutOfRangeError for negative offset or size values. Add two regression tests confirming the bypass behavior and documenting the required fix behavior. -- 5fdce18 by prasanna8585 <[email protected]>: Address review: move fix from FromProto to BufferAllocation::AddAssignment Per review feedback from @beckerhe: move the negative offset/size validation from a precondition check in BufferAssignment::FromProto into BufferAllocation::AddAssignment itself, converting the existing CHECK_LE statements to TF_RET_CHECK. This protects every caller of AddAssignment (there are ~9 call sites across buffer_assignment.cc), not just the FromProto deserialization path, and returns an absl::Status error instead of crashing the process when the precondition is violated. Also adds explicit offset >= 0 / size >= 0 checks: converting CHECK_LE(offset, size_) to TF_RET_CHECK alone would not catch a negative offset, since offset <= size_ is still true for any negative offset and non-negative allocation size. All 9 existing call sites already propagate the returned Status via RETURN_IF_ERROR, so no call-site changes were needed beyond removing the now-redundant check in FromProto. Merging this change closes #46105 FUTURE_COPYBARA_INTEGRATE_REVIEW=#46105 from prasanna8585:fix/buffer-assignment-from-proto-negative-offset 5fdce18 PiperOrigin-RevId: 952638612

@beckerhe

…:AddAssignment Imported from GitHub PR #46105 ## Problem `BufferAssignment::FromProto` passes `assignee.offset()` and `assignee.size()` from a deserialized `BufferAllocationProto.Assigned` proto directly to `BufferAllocation::AddAssignment` without checking for negative values. The guards inside `AddAssignment` are: ```cpp CHECK_LE(offset, size_) // checks offset <= allocation.size() CHECK_LE(offset + size, size_) // checks offset + size <= allocation.size() ``` These `CHECK_LE` guards are **bypassed by negative int64 values**: `-1 <= 1024` is true for any real allocation, so a negative offset passes both checks and gets stored in `assigned_buffers_`. The stored negative offset is later used in LLVM IR codegen (`ir_emitter.cc`): ```cpp tempbuf_address_base + b()->getInt64(slice.offset()) ``` With `offset = -1`, this produces a pointer one byte before the allocation — an out-of-bounds address. ## Attack Surface Reachable from `CpuAotLoader::LoadAotCompilationResult(serialized_aot_result)` where `serialized_aot_result` is a user-supplied serialized AOT executable. The proto fields `offset` and `size` are both `int64` in `hlo.proto`, so negative values are fully valid at the wire level. ## Relationship to PR #44653 PR #44653 (`2840cd3`) fixed `Slice::FromProto` with: ```cpp if (proto.offset() < 0 || proto.size() < 0) { return absl::OutOfRangeError(...); } ``` This PR applies the **same fix** to `BufferAssignment::FromProto` → `AddAssignment` path, which handles `Assigned` buffers in `BufferAllocationProto` — the sibling that the previous fix missed. ## Fix Added a pre-validation check in `BufferAssignment::FromProto` before calling `AddAssignment`: ```cpp if (assignee.offset() < 0 || assignee.size() < 0) { return absl::OutOfRangeError(absl::StrCat( "BufferAssignment proto contains an assigned buffer with negative ", "offset or size: logical_buffer_id=", logical_buffer_id, " offset=", assignee.offset(), " size=", assignee.size())); } ``` ## Testing Two new test cases in `buffer_assignment_test.cc`: - `BufferAssignmentFromProtoTest.RejectsNegativeOffset`: confirms `offset=-1` bypasses `CHECK_LE` guards (proving the vulnerability) and documents that the fix must reject it before reaching `AddAssignment` - `BufferAssignmentFromProtoTest.RejectsNegativeSize`: same for `size=-1` --- *Disclosure: this contribution was authored with the aid of an AI coding assistant and reviewed before submission.* Copybara import of the project: -- 6fab2ec by prasanna8585 <[email protected]>: [XLA] Validate assigned buffer offset/size in BufferAssignment::FromProto BufferAssignment::FromProto passes assignee.offset() and assignee.size() from a deserialized BufferAllocationProto.Assigned proto directly to BufferAllocation::AddAssignment without checking for negative values. The CHECK_LE guards in AddAssignment are bypassed by negative int64 values: CHECK_LE(-1, 1024) = true (for any real allocation) CHECK_LE(-1 + 4, 1024) = true So a negative offset is stored in assigned_buffers_ and later used in LLVM IR codegen (ir_emitter.cc) as: tempbuf_address_base + b()->getInt64(slice.offset()) With offset=-1 this is an out-of-bounds address. This mirrors the fix applied to Slice::FromProto in PR #44653 (2840cd3), which fixed the thunk deserialization path but missed this sibling path in BufferAssignment::FromProto. Fix: add an explicit negative check before calling AddAssignment, returning absl::OutOfRangeError for negative offset or size values. Add two regression tests confirming the bypass behavior and documenting the required fix behavior. -- 5fdce18 by prasanna8585 <[email protected]>: Address review: move fix from FromProto to BufferAllocation::AddAssignment Per review feedback from @beckerhe: move the negative offset/size validation from a precondition check in BufferAssignment::FromProto into BufferAllocation::AddAssignment itself, converting the existing CHECK_LE statements to TF_RET_CHECK. This protects every caller of AddAssignment (there are ~9 call sites across buffer_assignment.cc), not just the FromProto deserialization path, and returns an absl::Status error instead of crashing the process when the precondition is violated. Also adds explicit offset >= 0 / size >= 0 checks: converting CHECK_LE(offset, size_) to TF_RET_CHECK alone would not catch a negative offset, since offset <= size_ is still true for any negative offset and non-negative allocation size. All 9 existing call sites already propagate the returned Status via RETURN_IF_ERROR, so no call-site changes were needed beyond removing the now-redundant check in FromProto. Merging this change closes #46105 COPYBARA_INTEGRATE_REVIEW=#46105 from prasanna8585:fix/buffer-assignment-from-proto-negative-offset 5fdce18 PiperOrigin-RevId: 955212759

Validate buffer slices during deserialization

67c6ebd

M0nd0R changed the title ~~[codex] Validate buffer slices during deserialization~~ Validate buffer slices during deserialization Jun 23, 2026

M0nd0R marked this pull request as ready for review July 17, 2026 14:59

akuegel approved these changes Jul 20, 2026

View reviewed changes

copybara-service Bot mentioned this pull request Jul 20, 2026

PR #44653: Validate buffer slices during deserialization #45952

Merged

copybara-service Bot mentioned this pull request Jul 20, 2026

PR #44653: Validate buffer slices during deserialization tensorflow/tensorflow#123627

Merged

M0nd0R force-pushed the codex/validate-buffer-slice-from-proto branch from e4e4698 to 791e4ff Compare July 22, 2026 02:24

Fix deserialization tests for slice bounds checks

5ce256b

Grouped-matmul fixtures used a B allocation sized for the non-grouped shape, and DynamicSliceThunk round-trips hard-coded 1KiB allocations that cannot hold the 1MiB workspace slices once FromProto validates ranges.

M0nd0R force-pushed the codex/validate-buffer-slice-from-proto branch from 791e4ff to 5ce256b Compare July 22, 2026 02:25

akuegel approved these changes Jul 22, 2026

View reviewed changes

golechwierowicz approved these changes Jul 22, 2026

View reviewed changes

copybara-service Bot closed this in 2840cd3 Jul 22, 2026

copybara-service Bot mentioned this pull request Jul 22, 2026

Automated Code Change tensorflow/tensorflow#123659

Merged

copybara-service Bot mentioned this pull request Jul 22, 2026

* Optimize MutableLiteralBase::PopulateFromArray in literal.h to perform direct std::copy (memcpy) when shape().layout() is monotonic with dim 0 major. tensorflow/tensorflow#121544

Draft

prasanna8585 mentioned this pull request Jul 23, 2026

[XLA] Validate buffer offset and size in BufferAllocation::AddAssignment #46105

Closed

copybara-service Bot mentioned this pull request Jul 23, 2026

PR #46105: [XLA] Validate buffer offset and size in BufferAllocation::AddAssignment #46112

Merged

Navigation Menu

Repository navigation

Validate buffer slices during deserialization - #44653

Validate buffer slices during deserialization#44653
M0nd0R wants to merge 2 commits into
openxla:mainfrom
M0nd0R:codex/validate-buffer-slice-from-proto

M0nd0R commented Jun 23, 2026

Uh oh!

akuegel commented Jul 21, 2026

Uh oh!

M0nd0R commented Jul 22, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

M0nd0R commented Jun 23, 2026

Summary

Why

Testing

Uh oh!

akuegel commented Jul 21, 2026

Uh oh!

M0nd0R commented Jul 22, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants