Fix MobileNetV3 Preprocess - #46105

…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