From e611bbcedb936e03fc0d6e527c48eec9fd841f65 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Fri, 1 May 2026 16:52:14 -0700 Subject: [PATCH 1/2] Add overflow checks to getLeadingDims and getTrailingDims Summary: Add `c10::mul_overflows()` checks to the dimension-product loops in `getLeadingDims()` and `getTrailingDims()`. Both functions multiply tensor dimension sizes in a loop with no overflow protection. On 32-bit targets where `size_t` is 32 bits, malicious tensor dimensions from a crafted `.pte` file can cause the product to wrap silently, producing a small value that is then used for buffer offset calculations in 40+ kernels via `coordinateToIndex()`. This enables heap buffer overflows during operator execution. MACA-2026-001 (T267380210). Differential Revision: D103467782 --- runtime/core/exec_aten/util/tensor_util.h | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/runtime/core/exec_aten/util/tensor_util.h b/runtime/core/exec_aten/util/tensor_util.h index 26b97e5a7a2..4d742bda0fa 100644 --- a/runtime/core/exec_aten/util/tensor_util.h +++ b/runtime/core/exec_aten/util/tensor_util.h @@ -9,6 +9,7 @@ #pragma once #include +#include #include #include // std::array #include // PRId64 @@ -932,7 +933,12 @@ inline size_t getLeadingDims( ssize_t(tensor.dim())); size_t dims = 1; for (const auto i : c10::irange(dim)) { - dims *= static_cast(tensor.size(i)); + size_t next_dims; + ET_CHECK_MSG( + !c10::mul_overflows(dims, static_cast(tensor.size(i)), &next_dims), + "Overflow computing leading dims at dimension %zd", + (ssize_t)i); + dims = next_dims; } return dims; } @@ -949,7 +955,12 @@ inline size_t getTrailingDims( ssize_t(tensor.dim())); size_t dims = 1; for (size_t i = dim + 1; i < static_cast(tensor.dim()); ++i) { - dims *= static_cast(tensor.size(i)); + size_t next_dims; + ET_CHECK_MSG( + !c10::mul_overflows(dims, static_cast(tensor.size(i)), &next_dims), + "Overflow computing trailing dims at dimension %zu", + i); + dims = next_dims; } return dims; } From 34ef14dc89dfad608bd605ff321bf059341b3236 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Fri, 1 May 2026 17:30:04 -0700 Subject: [PATCH 2/2] Fix formatting --- runtime/core/exec_aten/util/tensor_util.h | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/runtime/core/exec_aten/util/tensor_util.h b/runtime/core/exec_aten/util/tensor_util.h index 4d742bda0fa..b9cbab4c1ef 100644 --- a/runtime/core/exec_aten/util/tensor_util.h +++ b/runtime/core/exec_aten/util/tensor_util.h @@ -935,7 +935,8 @@ inline size_t getLeadingDims( for (const auto i : c10::irange(dim)) { size_t next_dims; ET_CHECK_MSG( - !c10::mul_overflows(dims, static_cast(tensor.size(i)), &next_dims), + !c10::mul_overflows( + dims, static_cast(tensor.size(i)), &next_dims), "Overflow computing leading dims at dimension %zd", (ssize_t)i); dims = next_dims; @@ -957,7 +958,8 @@ inline size_t getTrailingDims( for (size_t i = dim + 1; i < static_cast(tensor.dim()); ++i) { size_t next_dims; ET_CHECK_MSG( - !c10::mul_overflows(dims, static_cast(tensor.size(i)), &next_dims), + !c10::mul_overflows( + dims, static_cast(tensor.size(i)), &next_dims), "Overflow computing trailing dims at dimension %zu", i); dims = next_dims;