Go: restore the possibility of go getting tensorflow and fix tests - #44653

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 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 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 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 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 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

…ment::FromProto

Imported from GitHub PR openxla/xla#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:

--
6fab2ecd47988e30cfe195becfaecdb641479e8c 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=openxla/xla#46105 from prasanna8585:fix/buffer-assignment-from-proto-negative-offset 6fab2ecd47988e30cfe195becfaecdb641479e8c
PiperOrigin-RevId: 952638612

…ment::FromProto

Imported from GitHub PR openxla/xla#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:

--
6fab2ecd47988e30cfe195becfaecdb641479e8c 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=openxla/xla#46105 from prasanna8585:fix/buffer-assignment-from-proto-negative-offset 6fab2ecd47988e30cfe195becfaecdb641479e8c
PiperOrigin-RevId: 952638612

…ment::FromProto

Imported from GitHub PR openxla/xla#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:

--
6fab2ecd47988e30cfe195becfaecdb641479e8c 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=openxla/xla#46105 from prasanna8585:fix/buffer-assignment-from-proto-negative-offset 6fab2ecd47988e30cfe195becfaecdb641479e8c
PiperOrigin-RevId: 952638612

…ment::FromProto

Imported from GitHub PR openxla/xla#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:

--
6fab2ecd47988e30cfe195becfaecdb641479e8c 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=openxla/xla#46105 from prasanna8585:fix/buffer-assignment-from-proto-negative-offset 6fab2ecd47988e30cfe195becfaecdb641479e8c
PiperOrigin-RevId: 952638612

…:AddAssignment

Imported from GitHub PR openxla/xla#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:

--
6fab2ecd47988e30cfe195becfaecdb641479e8c 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.

--
5fdce18c4bdf2e1e53731611bc08938cf3669a05 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=openxla/xla#46105 from prasanna8585:fix/buffer-assignment-from-proto-negative-offset 5fdce18c4bdf2e1e53731611bc08938cf3669a05
PiperOrigin-RevId: 952638612

…:AddAssignment

Imported from GitHub PR openxla/xla#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:

--
6fab2ecd47988e30cfe195becfaecdb641479e8c 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.

--
5fdce18c4bdf2e1e53731611bc08938cf3669a05 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=openxla/xla#46105 from prasanna8585:fix/buffer-assignment-from-proto-negative-offset 5fdce18c4bdf2e1e53731611bc08938cf3669a05
PiperOrigin-RevId: 952638612

…:AddAssignment

Imported from GitHub PR openxla/xla#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:

--
6fab2ecd47988e30cfe195becfaecdb641479e8c 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.

--
5fdce18c4bdf2e1e53731611bc08938cf3669a05 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=openxla/xla#46105 from prasanna8585:fix/buffer-assignment-from-proto-negative-offset 5fdce18c4bdf2e1e53731611bc08938cf3669a05
PiperOrigin-RevId: 952638612

…:AddAssignment

Imported from GitHub PR openxla/xla#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:

--
6fab2ecd47988e30cfe195becfaecdb641479e8c 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.

--
5fdce18c4bdf2e1e53731611bc08938cf3669a05 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=openxla/xla#46105 from prasanna8585:fix/buffer-assignment-from-proto-negative-offset 5fdce18c4bdf2e1e53731611bc08938cf3669a05
PiperOrigin-RevId: 952638612

…:AddAssignment

Imported from GitHub PR openxla/xla#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:

--
6fab2ecd47988e30cfe195becfaecdb641479e8c 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.

--
5fdce18c4bdf2e1e53731611bc08938cf3669a05 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=openxla/xla#46105 from prasanna8585:fix/buffer-assignment-from-proto-negative-offset 5fdce18c4bdf2e1e53731611bc08938cf3669a05
PiperOrigin-RevId: 952638612

…:AddAssignment

Imported from GitHub PR openxla/xla#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:

--
6fab2ecd47988e30cfe195becfaecdb641479e8c 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.

--
5fdce18c4bdf2e1e53731611bc08938cf3669a05 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

PiperOrigin-RevId: 955212759