Skip to content

Add bounds checking on buffer_idx for constant_buffer and constant_data in XNNPACK #18820

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
+53 −14
Open

Add bounds checking on buffer_idx for constant_buffer and constant_data in XNNPACK #18820

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
67 changes: 53 additions & 14 deletions backends/xnnpack/runtime/XNNCompiler.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -173,7 +173,7 @@ Obtaining the constant data pointer can either be from within the flatbuffer
payload (deprecated) or via offsets to the constant_data_ptr. If no constant
data associated with the tensor value, then returns nullptr.
*/
const uint8_t* getConstantDataPtr(
Result<const uint8_t*> getConstantDataPtr(
uint32_t buffer_idx,
GraphPtr flatbuffer_graph,
const uint8_t* constant_data_ptr,
Expand All @@ -184,13 +184,39 @@ const uint8_t* getConstantDataPtr(
if (!constant_data_ptr) {
// TODO(T172265611): Remove constant_buffer in flatbuffer path after BC
// window
const auto& constant_buffer = *flatbuffer_graph->constant_buffer();
return constant_buffer[buffer_idx]->storage()->data();
auto* cb = flatbuffer_graph->constant_buffer();
ET_CHECK_OR_RETURN_ERROR(
cb != nullptr, InvalidProgram, "constant_buffer is null");
ET_CHECK_OR_RETURN_ERROR(
buffer_idx < cb->size(),
InvalidProgram,
"buffer_idx %u out of bounds for constant_buffer of size %u",
buffer_idx,
cb->size());
Comment on lines +191 to +195
Copy link

Copilot AI Apr 10, 2026

Choose a reason for hiding this comment

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

The format string uses %u for cb->size(), but flatbuffers::Vector::size() is size_t. This is undefined behavior on some platforms; use %zu (or cast to an unsigned type and keep %u) so the log formatting matches the argument type.

Copilot uses AI. Check for mistakes.
auto* buffer_entry = (*cb)[buffer_idx];
ET_CHECK_OR_RETURN_ERROR(
buffer_entry != nullptr && buffer_entry->storage() != nullptr,
InvalidProgram,
"Null constant_buffer entry at buffer_idx %u",
buffer_idx);
return buffer_entry->storage()->data();
} else {
ConstantDataOffsetPtr constant_data_offset =
flatbuffer_graph->constant_data()->Get(buffer_idx);
auto* cd = flatbuffer_graph->constant_data();
ET_CHECK_OR_RETURN_ERROR(
cd != nullptr, InvalidProgram, "constant_data is null");
ET_CHECK_OR_RETURN_ERROR(
buffer_idx < cd->size(),
InvalidProgram,
"buffer_idx %u out of bounds for constant_data of size %u",
buffer_idx,
cd->size());
Comment on lines +208 to +212
Copy link

Copilot AI Apr 10, 2026

Choose a reason for hiding this comment

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

Same issue here: the log uses %u for cd->size() (a size_t). Please switch to %zu or cast appropriately to avoid varargs type mismatches/UB in the error path.

Copilot uses AI. Check for mistakes.
ConstantDataOffsetPtr constant_data_offset = cd->Get(buffer_idx);
Copy link

Copilot AI Apr 10, 2026

Choose a reason for hiding this comment

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

cd->Get(buffer_idx) can still return nullptr if the vector contains a null offset (possible with corrupted input). constant_data_offset is dereferenced immediately afterward; add a null check for constant_data_offset and handle it as an invalid constant_data entry (log + propagate error).

Suggested change
ConstantDataOffsetPtr constant_data_offset = cd->Get(buffer_idx);
ConstantDataOffsetPtr constant_data_offset = cd->Get(buffer_idx);
if (constant_data_offset == nullptr) {
ET_LOG(
Error,
"Invalid null constant_data entry at buffer_idx %u",
buffer_idx);
return nullptr;
}

Copilot uses AI. Check for mistakes.
ET_CHECK_OR_RETURN_ERROR(
constant_data_offset != nullptr,
InvalidProgram,
"Null constant_data entry at buffer_idx %u",
buffer_idx);
uint64_t offset = constant_data_offset->offset();

bool has_named_key = flatbuffers::IsFieldPresent(
constant_data_offset, fb_xnnpack::ConstantDataOffset::VT_NAMED_KEY);
// If there is no tensor name
Expand All @@ -203,7 +229,7 @@ const uint8_t* getConstantDataPtr(
weights_cache->load_unpacked_data(data_name);
if (!data_ptr.ok()) {
ET_LOG(Error, "Failed to load weights from cache");
return nullptr;
return Error::InvalidProgram;
}
Comment on lines 229 to 233
Copy link

Copilot AI Apr 10, 2026

Choose a reason for hiding this comment

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

On weights-cache failures, this now always returns Error::InvalidProgram, even though load_unpacked_data() already returns a more specific error (e.g., InvalidExternalData). Consider returning data_ptr.error() (and optionally logging it) so callers can distinguish invalid models vs missing/invalid external data.

Copilot uses AI. Check for mistakes.
return data_ptr.get();
#else
Expand All @@ -215,7 +241,7 @@ const uint8_t* getConstantDataPtr(
"Failed to get constant data for key %s from named_data_map. Error code: %u",
data_name.c_str(),
static_cast<uint32_t>(buffer.error()));
return nullptr;
return Error::InvalidProgram;
Copy link

Copilot AI Apr 10, 2026

Choose a reason for hiding this comment

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

This branch discards the underlying named_data_map->get_data() error and always returns Error::InvalidProgram. Propagating buffer.error() would preserve the real failure mode (e.g., NotFound/InvalidExternalData) and make error handling/debugging more accurate.

Suggested change
return Error::InvalidProgram;
return buffer.error();

Copilot uses AI. Check for mistakes.
}
const uint8_t* data_ptr =
static_cast<const uint8_t*>(buffer.get().data());
Expand All @@ -229,7 +255,7 @@ const uint8_t* getConstantDataPtr(
return nullptr;
}

const uint8_t* getConstantDataPtr(
Result<const uint8_t*> getConstantDataPtr(
const fb_xnnpack::XNNTensorValue* tensor_value,
GraphPtr flatbuffer_graph,
const uint8_t* constant_data_ptr,
Expand Down Expand Up @@ -298,13 +324,17 @@ Error defineTensor(

// Get Pointer to constant data from flatbuffer, if its non-constant
// it is a nullptr
const uint8_t* buffer_ptr = getConstantDataPtr(
auto buffer_result = getConstantDataPtr(
tensor_value,
flatbuffer_graph,
constant_data_ptr,
named_data_map,
freeable_buffers,
weights_cache);
if (!buffer_result.ok()) {
return buffer_result.error();
}
const uint8_t* buffer_ptr = buffer_result.get();

xnn_status status;
// The type we might have to convert to
Expand Down Expand Up @@ -449,13 +479,17 @@ Error defineTensor(
const float* scale = qparams->scale()->data();

if (qparams->scale_buffer_idx() != 0) {
scale = reinterpret_cast<const float*>(getConstantDataPtr(
auto scale_result = getConstantDataPtr(
qparams->scale_buffer_idx(),
flatbuffer_graph,
constant_data_ptr,
named_data_map,
freeable_buffers,
weights_cache));
weights_cache);
if (!scale_result.ok()) {
return scale_result.error();
}
scale = reinterpret_cast<const float*>(scale_result.get());
ET_CHECK_OR_RETURN_ERROR(
scale != nullptr, Internal, "Failed to load scale data.");
}
Expand Down Expand Up @@ -491,13 +525,18 @@ Error defineTensor(
// Block scales are preferably serialized as bf16 but can also be
// serialized as fp32 for backwards compatability.
if (qparams->scale_buffer_idx() != 0) {
scale_data = reinterpret_cast<const uint16_t*>(getConstantDataPtr(
auto scale_data_result = getConstantDataPtr(
qparams->scale_buffer_idx(),
flatbuffer_graph,
constant_data_ptr,
named_data_map,
freeable_buffers,
weights_cache));
weights_cache);
if (!scale_data_result.ok()) {
return scale_data_result.error();
}
scale_data =
reinterpret_cast<const uint16_t*>(scale_data_result.get());
ET_CHECK_OR_RETURN_ERROR(
scale_data != nullptr, Internal, "Failed to load scale data.");
scale_numel = qparams->num_scales();
Expand Down
Loading