Skip to content

fix: validate tensor indices against subgraph bounds in MarkAllocationLifetimes #3591

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
arunmanni-ai wants to merge 1 commit into
base: main
Choose a base branch
from
+31 −0
Open

fix: validate tensor indices against subgraph bounds in MarkAllocationLifetimes #3591

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
31 changes: 31 additions & 0 deletions tensorflow/lite/micro/micro_allocation_info.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -243,10 +243,19 @@ TfLiteStatus AllocationInfoBuilder::MarkAllocationLifetimes(
&allocation_info[info_.subgraph_offsets[subgraph_idx]];

uint32_t operators_size = NumSubgraphOperators(subgraph);
const int tensors_size =
static_cast<int>(subgraph->tensors() ? subgraph->tensors()->size() : 0);
// Mark all inputs as created at the start of the subgraph invocation.
for (size_t i = 0;
subgraph->inputs() != nullptr && i < subgraph->inputs()->size(); ++i) {
const int tensor_index = subgraph->inputs()->Get(i);
if (tensor_index < 0 || tensor_index >= tensors_size) {
MicroPrintf(
"Subgraph input tensor index %d is out of bounds "
"(subgraph has %d tensors).",
tensor_index, tensors_size);
return kTfLiteError;
}
AllocationInfo* current = &subgraph_allocation_info[tensor_index];
UpdateFirstCreated(current, allocation_scope_count_);
// This will ensure that the tensors that are inputs to the subgraphs
Expand All @@ -262,6 +271,11 @@ TfLiteStatus AllocationInfoBuilder::MarkAllocationLifetimes(
for (size_t n = 0; op->outputs() != nullptr && n < op->outputs()->size();
++n) {
const int tensor_index = op->outputs()->Get(n);
if (tensor_index < 0 || tensor_index >= tensors_size) {
MicroPrintf("Op#%u output tensor index %d is out of bounds "
"(subgraph has %d tensors).", i, tensor_index, tensors_size);
return kTfLiteError;
}
AllocationInfo* current = &subgraph_allocation_info[tensor_index];
UpdateFirstCreated(current, allocation_scope_count_);
}
Expand All @@ -281,6 +295,11 @@ TfLiteStatus AllocationInfoBuilder::MarkAllocationLifetimes(
const int tensor_index = op->inputs()->Get(n);
// Optional bias tensors can have an index of -1 when they are omitted.
if (tensor_index >= 0) {
if (tensor_index >= tensors_size) {
MicroPrintf("Op#%u input tensor index %d is out of bounds "
"(subgraph has %d tensors).", i, tensor_index, tensors_size);
return kTfLiteError;
}
AllocationInfo* current = &subgraph_allocation_info[tensor_index];
// No need to update creation since it is either marked by the subgraph
// or producer op, or it is not part of the memory plan (weight, bias
Expand All @@ -291,6 +310,11 @@ TfLiteStatus AllocationInfoBuilder::MarkAllocationLifetimes(
for (size_t n = 0; op->outputs() != nullptr && n < op->outputs()->size();
++n) {
const int tensor_index = op->outputs()->Get(n);
if (tensor_index < 0 || tensor_index >= tensors_size) {
MicroPrintf("Op#%u output tensor index %d is out of bounds "
"(subgraph has %d tensors).", i, tensor_index, tensors_size);
return kTfLiteError;
}
AllocationInfo* current = &subgraph_allocation_info[tensor_index];
UpdateLastUsed(current, allocation_scope_count_);
}
Expand Down Expand Up @@ -322,6 +346,13 @@ TfLiteStatus AllocationInfoBuilder::MarkAllocationLifetimes(
for (size_t i = 0;
subgraph->outputs() != nullptr && i < subgraph->outputs()->size(); ++i) {
const int tensor_index = subgraph->outputs()->Get(i);
if (tensor_index < 0 || tensor_index >= tensors_size) {
MicroPrintf(
"Subgraph output tensor index %d is out of bounds "
"(subgraph has %d tensors).",
tensor_index, tensors_size);
return kTfLiteError;
}
AllocationInfo* current = &subgraph_allocation_info[tensor_index];
// Make sure to assign the First created value of the subgraph output
// This will handle the case where the subgraph is empty. This helps
Expand Down