google / google/XNNPACK

litert/tensor: derive loop counts from the buffer element count, not shape[0] (shape/buffer guard-series extension)

Open
#11,199 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
C
Stars
2.5k
Forks
560
Avg merge
1d 6h
Merged PRs (30d)
163

Description

````markdown
### One more invariant for the shape/buffer guard series: derive loop counts
### from the buffer, not from `shape[0]`

The `TensorInformation` pair (`shape` int32 span, `buffer` shared_ptr)
is filled in independently by the caller, and nothing in litert/tensor enforces
`shape_product * element_size == buffer->size()` (`LockedBufferSpan::As()`,
buffer.h:70-77, reinterprets with `bytes_ / sizeof(U)` — silent truncation,
no consistency failure). The recent guard series has been closing the
shape-validation family op by op (rank-0, empty axis buffer, axis bounds,
negative dims, scale-array size). The variant none of them cover yet: code
that loops `i < shape[0]` while indexing a buffer-sized span. If an embedding
application ever sets a shape whose count exceeds its buffer's element count,
these loops read past the allocation.

**Sites** (litert/tensor @ master `d50a82e`):

1. `arithmetic.h` reduce constructors — Sum :790 (loops :819-827, :833-844),
ReduceMax :867 (:897-905, :911-922), Mean :944 (:976-984, :990-1001):
`for (int i = 0; i < b_info.shape[0]; ++i) { int axis = b_data[i]; ... }`
with `b_data` buffer-sized. The 0D branch reads `b_data[0]` after only
`shape.empty()` is checked — a non-null zero-byte span with `shape = {1}`
overreads the same way.
2. `backends/xnnpack/arithmetic.cc` `MeanOperation::ToXnnpack` :1054-1061:
`num_axes = axes_info.shape[0]` but the `std::vector` is built
buffer-sized from `axes_data` — both go to `xnn_define_static_reduce_v2`,
whose `memcpy` (src/subgraph/static-reduce.c:374-375) reads `num_axes`
entries.
3. Same file `SliceOperation::ToXnnpack` :1102-1124: `num_dims =
begin_info.shape[0]`; loops read `begin_data[i]` / `size_data[i]` with no
length check on the size tensors — Transpose (:1002) and Tile (:1341)
already check; Slice does not. Feeds static-slice.c:265-267.

**Bounded, and why it still merits a guard**: every value read is immediately
axis-bounds-validated (`axis < 0 || axis >= rank`) in the same loops, and the
XNNPACK boundary rejects counts above `XNN_MAX_TENSOR_DIMS` before the memcpy
(static-reduce.c:361-368, static-slice.c:183-188), so the overread is capped
at 5 * 8 = 40 bytes past the vector and cannot steer indexing — worst case is
a fault on an unmapped page. That is exactly the crash-class the merged guards
exist to prevent from reaching a user's build.

**Suggested invariant** (either half closes the class):
- assert at op construction that the buffer's element count >= `shape[0]`
(Sum/ReduceMax/Mean constructors, Mean/Slice ToXnnpack entry), or
- pass the count from `LockedBufferSpan::size() / sizeof(T)` instead of
`shape[0]` at the sites above.

The established pattern to copy is the merged guard series: 89ef795, 9f5e4f9,
7f616c1, 4013867, 78574e5, deaf412, 73ef51e, 42992df, 3a32935 — each adds a
construction-time check with a clear error and stops the op before the glue
reads unvalidated memory.
````

Contributor guide

Open the contributing guide

Research direction

Start with the listed reduce constructors in arithmetic.h, then inspect MeanOperation::ToXnnpack and SliceOperation::ToXnnpack in backends/xnnpack/arithmetic.cc. Compare their shape-based loop counts with LockedBufferSpan element counts and the existing guard commits. Done means every listed buffer read is bounded by the available element count, including the 0D case, with relevant existing tests passing.

Written by the indexing model from the issue text.

Assessment

Tech stack
c, cpp
Domain
backend, security
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
67/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.