From 20114ad628f7f7d83a99125bd1b337acd8fa9f9b Mon Sep 17 00:00:00 2001 From: TristanInSec Date: Fri, 17 Apr 2026 12:03:43 -0400 Subject: [PATCH] Add overflow detection to xnn_shape_multiply_* functions The five shape multiplication functions multiply size_t dimensions in a loop without checking for overflow. On platforms where size_t is 32-bit (ARM32, WebAssembly), dimension products exceeding 2^32 silently wrap, causing get_tensor_size() to return undersized allocation sizes. Add a checked_mul helper using __builtin_mul_overflow that returns SIZE_MAX on overflow. This propagates through the multiplication chain so callers see SIZE_MAX instead of a wrapped value. --- src/tensor.c | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/src/tensor.c b/src/tensor.c index d53f0c42c83..4a7d6351f84 100644 --- a/src/tensor.c +++ b/src/tensor.c @@ -637,12 +637,20 @@ enum xnn_status xnn_define_blockwise_quantized_tensor_value( block_size, dims, data, external_id, flags, xnn_datatype_bf16, id_out); } +static inline size_t checked_mul(size_t a, size_t b) { + size_t result; + if (__builtin_mul_overflow(a, b, &result)) { + return SIZE_MAX; + } + return result; +} + size_t xnn_shape_multiply_all_dims( const struct xnn_shape* shape) { size_t batch_size = 1; for (size_t i = 0; i < shape->num_dims; i++) { - batch_size *= shape->dim[i]; + batch_size = checked_mul(batch_size, shape->dim[i]); } return batch_size; } @@ -653,7 +661,7 @@ size_t xnn_shape_multiply_batch_dims( { size_t batch_size = 1; for (size_t i = 0; i + num_nonbatch_dims < shape->num_dims; i++) { - batch_size *= shape->dim[i]; + batch_size = checked_mul(batch_size, shape->dim[i]); } return batch_size; } @@ -663,7 +671,7 @@ size_t xnn_shape_multiply_non_channel_dims( { size_t batch_size = 1; for (size_t i = 0; i + 1 < shape->num_dims; i++) { - batch_size *= shape->dim[i]; + batch_size = checked_mul(batch_size, shape->dim[i]); } return batch_size; } @@ -674,7 +682,7 @@ size_t xnn_shape_multiply_leading_dims( { size_t batch_size = 1; for (size_t i = 0; i < num_leading_dims; i++) { - batch_size *= shape->dim[i]; + batch_size = checked_mul(batch_size, shape->dim[i]); } return batch_size; } @@ -685,7 +693,7 @@ size_t xnn_shape_multiply_trailing_dims( { size_t product = 1; for (size_t i = start_dim; i < shape->num_dims; i++) { - product *= shape->dim[i]; + product = checked_mul(product, shape->dim[i]); } return product; }