Skip to content

Fix overflows in et #19057

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lucylq wants to merge 1 commit into
base: main
Choose a base branch
from
+34 −10
Open

Fix overflows in et #19057

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion backends/apple/metal/runtime/metal_backend.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
* LICENSE file in the root directory of this source tree.
*/

#include <c10/util/safe_numerics.h>
#include <dlfcn.h>
#include <executorch/runtime/backend/interface.h>
#include <executorch/runtime/core/error.h>
Expand Down Expand Up @@ -459,8 +460,10 @@ class ET_EXPERIMENTAL MetalBackend final

ET_LOG(Debug, "MetalBackend n_outputs %zd generated", n_outputs);

size_t n_io_sum = 0;
ET_CHECK_OR_RETURN_ERROR(
n_inputs + n_outputs == args.size(),
!c10::add_overflows(n_inputs, n_outputs, &n_io_sum) &&
n_io_sum == args.size(),
InvalidArgument,
"number of user input %zd and output %zd generated from AOT Inductor does not match ET runner's %zd. Exit.",
n_inputs,
Expand Down
5 changes: 4 additions & 1 deletion backends/cuda/runtime/cuda_backend.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
* LICENSE file in the root directory of this source tree.
*/

#include <c10/util/safe_numerics.h>
#include <cuda_runtime.h>
#include <executorch/runtime/backend/interface.h>
#include <executorch/runtime/backend/options.h>
Expand Down Expand Up @@ -550,8 +551,10 @@ class ET_EXPERIMENTAL CudaBackend final

setCurrentCUDAStream(handle->get_cuda_stream(), 0);

size_t n_io_sum = 0;
ET_CHECK_OR_RETURN_ERROR(
n_inputs + n_outputs == args.size(),
!c10::add_overflows(n_inputs, n_outputs, &n_io_sum) &&
n_io_sum == args.size(),
InvalidArgument,
"number of user input %zd and output %zd generated from AOT Inductor does not match ET runner's %zd. Exit.",
n_inputs,
Expand Down
8 changes: 6 additions & 2 deletions examples/qualcomm/oss_scripts/llama/runner/prompt_processor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
* LICENSE file in the root directory of this source tree.
*/

#include <c10/util/safe_numerics.h>
#include <executorch/examples/qualcomm/oss_scripts/llama/runner/prompt_processor.h>
#include <numeric>
using executorch::aten::TensorImpl;
Expand Down Expand Up @@ -248,9 +249,12 @@ Result<uint64_t> PromptProcessor<T>::prefill(
ET_CHECK_MSG(
start_pos == 0, "Bert model doesn't support multi-turn conversation.");
} else if (!enable_attention_sink) {
int64_t end_pos = 0;
ET_CHECK_MSG(
(start_pos + num_prompt_tokens) <=
(metadata_.context_len - metadata_.ar_len),
!c10::add_overflows(
start_pos, static_cast<int64_t>(num_prompt_tokens), &end_pos) &&
end_pos <= static_cast<int64_t>(metadata_.context_len) -
static_cast<int64_t>(metadata_.ar_len),
"The sequence length exceeds the maximum limit that the prompt processor can handle.");
Comment on lines +252 to 258
Copy link

Copilot AI Apr 22, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This uses c10::add_overflows with signed int64_t operands. In this repo's MSVC fallback implementation, add_overflows computes a + b for signed types, which can trigger signed-overflow UB before detection. To keep this check safe cross-platform, prefer unsigned arithmetic (after asserting start_pos >= 0) or rewrite as a non-overflowing bounds check using wider/unsigned types.

Copilot uses AI. Check for mistakes.
}

Expand Down
6 changes: 5 additions & 1 deletion examples/qualcomm/oss_scripts/llama/runner/runner.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
// A llama 3.2 runner that includes preprocessing and post processing
// logic. The module takes in a string as input and emits a string as output.

#include <c10/util/safe_numerics.h>
#include <executorch/examples/models/llama/runner/runner.h>
#include <executorch/examples/models/llama/tokenizer/llama_tiktoken.h>
#include <executorch/examples/qualcomm/oss_scripts/llama/runner/client_mem.h>
Expand Down Expand Up @@ -433,8 +434,11 @@ Error Runner<T>::generate_from_prompt_or_file(
}
int num_prompt_tokens = prompt_tokens.size();
ET_CHECK_MSG(num_prompt_tokens >= 1, "Expected at least 1 prompt token");
int64_t end_pos = 0;
ET_CHECK_MSG(
cur_pos_ + num_prompt_tokens < seq_len,
!c10::add_overflows(
cur_pos_, static_cast<int64_t>(num_prompt_tokens), &end_pos) &&
end_pos < static_cast<int64_t>(seq_len),
Comment on lines +437 to +441
Copy link

Copilot AI Apr 22, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

On MSVC, c10::add_overflows' fallback implementation performs a + b for signed types, which can itself invoke signed-overflow UB before the overflow check triggers. Since cur_pos_/end_pos are int64_t, this reintroduces a platform-specific UB risk. Consider switching this check to unsigned arithmetic (e.g., validate cur_pos_ >= 0, cast operands to uint64_t/size_t, then use add_overflows) or use an explicit bounds check like cur_pos_ <= seq_len - num_prompt_tokens after normalizing types.

Copilot uses AI. Check for mistakes.
"sequence length exceeded - please increase the seq_len value");

// Prompt Processor first
Expand Down
10 changes: 8 additions & 2 deletions extension/flat_tensor/flat_tensor_data_map.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@

#include <executorch/extension/flat_tensor/flat_tensor_data_map.h>

#include <c10/util/safe_numerics.h>

#include <executorch/extension/flat_tensor/serialize/flat_tensor_generated.h>
#include <executorch/extension/flat_tensor/serialize/flat_tensor_header.h>

Expand Down Expand Up @@ -73,9 +75,13 @@ Result<const flat_tensor_flatbuffer::NamedData*> get_named_data(
key.data(),
segments->size());
// Validate the segment.
uint64_t seg_end = 0;
ET_CHECK_OR_RETURN_ERROR(
(segments->Get(segment_index)->offset() +
segments->Get(segment_index)->size()) <= segment_end_offset,
!c10::add_overflows(
static_cast<uint64_t>(segments->Get(segment_index)->offset()),
static_cast<uint64_t>(segments->Get(segment_index)->size()),
&seg_end) &&
seg_end <= static_cast<uint64_t>(segment_end_offset),
InvalidExternalData,
Comment on lines +78 to 85
Copy link

Copilot AI Apr 22, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

segment_end_offset is a size_t but is passed as header_.segment_base_offset + header_.segment_data_size (uint64_t). On 32-bit builds this can truncate, making the bounds check ineffective for large offsets even though seg_end is computed in uint64_t. Consider changing segment_end_offset (and related computations) to uint64_t throughout this helper to avoid truncation and remove the need for casts.

Copilot uses AI. Check for mistakes.
"Invalid segment offset %" PRIu64
" is larger than the segment_base_offset + segment_data_size %" PRIu64
Expand Down
4 changes: 3 additions & 1 deletion runtime/core/array_ref.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
#include <cstdint>

#include <c10/util/irange.h>
#include <c10/util/safe_numerics.h>
#include <executorch/runtime/platform/assert.h>

namespace executorch {
Expand Down Expand Up @@ -161,7 +162,8 @@ class ArrayRef final {
/// slice(n, m) - Take M elements of the array starting at element N
ArrayRef<T> slice(size_t N, size_t M) const {
// cant slice longer then the array
ET_CHECK(N + M <= size());
size_t end = 0;
ET_CHECK(!c10::add_overflows(N, M, &end) && end <= size());
return ArrayRef<T>(data() + N, M);
}

Expand Down
6 changes: 4 additions & 2 deletions runtime/core/hierarchical_allocator.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
#pragma once

#include <c10/util/irange.h>
#include <c10/util/safe_numerics.h>

#include <executorch/runtime/core/memory_allocator.h>
#include <executorch/runtime/core/result.h>
Expand Down Expand Up @@ -58,8 +59,9 @@ class HierarchicalAllocator final {
size_t offset_bytes,
size_t size_bytes) {
// Check for integer overflow in offset_bytes + size_bytes.
size_t end_bytes = 0;
ET_CHECK_OR_RETURN_ERROR(
size_bytes <= SIZE_MAX - offset_bytes,
!c10::add_overflows(offset_bytes, size_bytes, &end_bytes),
InvalidArgument,
"Integer overflow in offset_bytes (%" ET_PRIsize_t
") + size_bytes (%" ET_PRIsize_t ")",
Expand All @@ -73,7 +75,7 @@ class HierarchicalAllocator final {
buffers_.size());
Span<uint8_t> buffer = buffers_[memory_id];
ET_CHECK_OR_RETURN_ERROR(
offset_bytes + size_bytes <= buffer.size(),
end_bytes <= buffer.size(),
MemoryAllocationFailed,
"offset_bytes (%" ET_PRIsize_t ") + size_bytes (%" ET_PRIsize_t
") >= allocator size (%" ET_PRIsize_t
Expand Down
Loading