From 8950c56fd556d9406c624e02cfe2773fae9b6717 Mon Sep 17 00:00:00 2001 From: TristanInSec Date: Fri, 17 Apr 2026 12:04:37 -0400 Subject: [PATCH] Add overflow checks to reduce-nd workspace size computation The workspace size for reduce operations is computed by multiplying normalized input shape dimensions as size_t values. On 32-bit platforms, this can overflow to zero or a small value, causing an undersized workspace allocation that the reduction kernel then writes past. Add __builtin_mul_overflow checks to both the innermost and non-innermost reduction paths, returning xnn_status_invalid_parameter when the workspace size would overflow. --- src/operators/reduce-nd.c | 24 ++++++++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/src/operators/reduce-nd.c b/src/operators/reduce-nd.c index cdc737eb0a4..d304fd059c2 100644 --- a/src/operators/reduce-nd.c +++ b/src/operators/reduce-nd.c @@ -241,8 +241,16 @@ static enum xnn_status reshape_reduce_nd( size_t num_reduction_elements; if (normalized_reduction_axes[num_reduction_axes - 1] == num_input_dims - 1) { if (workspace_size != NULL) { - const size_t num_output_elements = normalized_input_shape[0] * normalized_input_shape[2] * normalized_input_shape[4]; - *workspace_size = (num_output_elements << log2_accumulator_element_size) + XNN_EXTRA_BYTES; + size_t num_output_elements; + size_t tmp; + if (__builtin_mul_overflow(normalized_input_shape[0], normalized_input_shape[2], &tmp) || + __builtin_mul_overflow(tmp, normalized_input_shape[4], &num_output_elements) || + __builtin_mul_overflow(num_output_elements, (size_t)1 << log2_accumulator_element_size, &tmp)) { + xnn_log_error("failed to reshape %s operator: workspace size overflow", + xnn_operator_type_to_string_v2(reduce_op)); + return xnn_status_invalid_parameter; + } + *workspace_size = tmp + XNN_EXTRA_BYTES; } num_reduction_elements = normalized_input_shape[1] * normalized_input_shape[3] * normalized_input_shape[5]; const size_t axis_dim = normalized_input_shape[5]; @@ -283,8 +291,16 @@ static enum xnn_status reshape_reduce_nd( // Reduction along the non-innermost dimension const size_t channel_like_dim = normalized_input_shape[XNN_MAX_TENSOR_DIMS - 1]; if (workspace_size != NULL) { - const size_t num_output_elements = normalized_input_shape[1] * normalized_input_shape[3] * normalized_input_shape[5]; - *workspace_size = (num_output_elements << log2_accumulator_element_size) + XNN_EXTRA_BYTES; + size_t num_output_elements; + size_t tmp; + if (__builtin_mul_overflow(normalized_input_shape[1], normalized_input_shape[3], &tmp) || + __builtin_mul_overflow(tmp, normalized_input_shape[5], &num_output_elements) || + __builtin_mul_overflow(num_output_elements, (size_t)1 << log2_accumulator_element_size, &tmp)) { + xnn_log_error("failed to reshape %s operator: workspace size overflow", + xnn_operator_type_to_string_v2(reduce_op)); + return xnn_status_invalid_parameter; + } + *workspace_size = tmp + XNN_EXTRA_BYTES; } num_reduction_elements = normalized_input_shape[0] * normalized_input_shape[2] * normalized_input_shape[4]; const size_t axis_dim = normalized_input_shape[4];