From e879e0a36e5f3f02ae1301adc3f4405b117758ee Mon Sep 17 00:00:00 2001 From: Gregory Comer Date: Mon, 8 Dec 2025 16:49:45 -0800 Subject: [PATCH] Prevent signed integer overflow in pixel_shuffle size calculation (#16138) Summary: The size calculation in `get_pixel_shuffle_out_target_size` can trigger signed integer overflow (UB) with very large upscale_factors. This manifested as an integer division by zero fault during fuzzing. Specifically, the calculation of `upscale_factor * upscale_factor` ([source](https://github.com/pytorch/executorch/blob/04f1e4d22383ffcbc770acf5002348e3f95082a2/kernels/portable/cpu/util/copy_ops_util.cpp#L364)) can overflow. In the motivating case, SizesType is a signed 32-bit integer and upscale_factor = 2^17. Since this is an impractically large upscale factor, I'm just adding a constraint that upscale_factor < 32768 (2^15). In theory, SizesType could be defined as less than 32 bits, but this seems unlikely in practice. Differential Revision: D88693324 --- kernels/portable/cpu/op_pixel_shuffle.cpp | 8 ++++++-- kernels/portable/cpu/util/copy_ops_util.cpp | 8 +++++++- kernels/portable/cpu/util/copy_ops_util.h | 2 +- 3 files changed, 14 insertions(+), 4 deletions(-) diff --git a/kernels/portable/cpu/op_pixel_shuffle.cpp b/kernels/portable/cpu/op_pixel_shuffle.cpp index fd9f1739e57..c550b789ed9 100644 --- a/kernels/portable/cpu/op_pixel_shuffle.cpp +++ b/kernels/portable/cpu/op_pixel_shuffle.cpp @@ -84,8 +84,12 @@ Tensor& pixel_shuffle_out( ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out); Tensor::SizesType expected_out_size[kTensorDimensionLimit]; size_t expected_out_dim = 0; - get_pixel_shuffle_out_target_size( - in, upscale_factor, expected_out_size, &expected_out_dim); + ET_KERNEL_CHECK( + ctx, + get_pixel_shuffle_out_target_size( + in, upscale_factor, expected_out_size, &expected_out_dim), + InvalidArgument, + out); // Make sure the output tensor is the right size. ET_KERNEL_CHECK( diff --git a/kernels/portable/cpu/util/copy_ops_util.cpp b/kernels/portable/cpu/util/copy_ops_util.cpp index 1527e6d9e35..5b4398daa09 100644 --- a/kernels/portable/cpu/util/copy_ops_util.cpp +++ b/kernels/portable/cpu/util/copy_ops_util.cpp @@ -346,11 +346,15 @@ bool check_pixel_unshuffle_args( return true; } -void get_pixel_shuffle_out_target_size( +bool get_pixel_shuffle_out_target_size( const Tensor& in, int64_t upscale_factor, executorch::aten::SizesType* out_sizes, size_t* out_ndim) { + // Prevent signed integer overflow when computing upscale_factor ^ 2. + ET_CHECK_OR_RETURN_FALSE( + upscale_factor < 32768, "Upscale factor must be less than 32768."); + *out_ndim = in.dim(); const executorch::aten::SizesType casted_upscale_factor = upscale_factor; @@ -366,6 +370,8 @@ void get_pixel_shuffle_out_target_size( out_sizes[i] = in.size(i) * casted_upscale_factor; i++; out_sizes[i] = in.size(i) * casted_upscale_factor; + + return true; } void get_pixel_unshuffle_out_target_size( diff --git a/kernels/portable/cpu/util/copy_ops_util.h b/kernels/portable/cpu/util/copy_ops_util.h index 15a7916e0e8..af2e8c83cad 100644 --- a/kernels/portable/cpu/util/copy_ops_util.h +++ b/kernels/portable/cpu/util/copy_ops_util.h @@ -139,7 +139,7 @@ bool check_pixel_shuffle_args( int64_t upscale_factor, Tensor& out); -void get_pixel_shuffle_out_target_size( +bool get_pixel_shuffle_out_target_size( const Tensor& in, int64_t upscale_factor, executorch::aten::SizesType* out_sizes,