Skip to content

litert/tensor: guard AveragePool2D/MaxPool2D/Conv2D/DepthwiseConv2D against under-rank input in arithmetic.h - #11154

Open
destro4evr-rgb wants to merge 1 commit into
google:masterfrom
destro4evr-rgb:fix/litert-conv-pool-arithmetic-h-rank-checks
Open

litert/tensor: guard AveragePool2D/MaxPool2D/Conv2D/DepthwiseConv2D against under-rank input in arithmetic.h#11154
destro4evr-rgb wants to merge 1 commit into
google:masterfrom
destro4evr-rgb:fix/litert-conv-pool-arithmetic-h-rank-checks

Conversation

@destro4evr-rgb

@destro4evr-rgb destro4evr-rgb commented Sep 5, 2026

Copy link
Copy Markdown
Contributor

Problem

AveragePool2D(), MaxPool2D(), Conv2DImpl(), and DepthwiseConv2DImpl() in litert/tensor/arithmetic.h subscript fixed shape indices (1, 2, 3) without first checking that the input/filter shape vectors are large enough. A crafted TFLite model with a rank-0 or rank-1 tensor on any of these operators causes an out-of-bounds std::vector::operator[] access during graph construction (model-load time), before inference begins, resulting in SIGSEGV.

Fix

Add shape.size() < 4 checks before the first shape subscript in each of the four functions, returning an InvalidArgumentError on violation — consistent with the pattern used in FullyConnected, EmbeddingLookup, and GatherNd.

Affected functions

Function File First crash line
AveragePool2D arithmetic.h:1220 shape[1] on rank-0 input
MaxPool2D arithmetic.h:1260 shape[1] on rank-0 input
Conv2DImpl arithmetic.h:1305 shape[1] on rank-0 input or filter
DepthwiseConv2DImpl arithmetic.h:1379 shape[1] on rank-0 input or filter

…gainst under-rank input in arithmetic.h

AveragePool2D, MaxPool2D, Conv2DImpl, and DepthwiseConv2DImpl all
access fixed shape indices (shape[1], shape[2], shape[3]) without
first checking that the input (and filter) shape vectors are large
enough. A crafted TFLite model with a rank-0 or rank-1 tensor on
any of these operators causes an out-of-bounds std::vector::operator[]
access at graph-construction time, before inference, resulting in
SIGSEGV.

PR google#11057 added equivalent guards to the ToXnnpack() methods in
arithmetic.cc (graph-compilation time), but the earlier crash in
arithmetic.h (graph-construction time) was not covered. Add rank < 4
checks before the first shape subscript in each function.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants