Fix/stickify allocation overflow chain - #3618
Conversation
getsize_ztensor computed dim4*dim3*ceil(dim2/sticks)*ceil(dim1/cells)*4096 as a plain uint64_t chain of multiplications with no overflow check. Each dimension is a uint32_t from an attacker-controlled zdnn_tensor_desc. With the per-dim and total-size validation in verify_transformed_descriptor disabled (#if 0), crafted dims can wrap the product to a tiny value. malloc_aligned_4k then succeeds with a tiny buffer, and the subsequent transform_ztensor call writes data sized to the full un-wrapped logical dimensions past the end of that buffer. Hunk 1: wrap each factor in __builtin_mul_overflow and return 0 on any overflow. malloc_aligned_4k already has an existing !size early-out (Stickify.cpp:272) that fails closed on this 0 return. Hunk 2: in stickify()'s pre-flight buffer-size check, cache the getsize_ztensor result and explicitly reject !required_size alongside the existing size-comparison, so a 0 from an overflowing product is caught before any buffer access rather than passing the size comparison by coincidence. Signed-off-by: Jonathan Huang <jonathanhuang@ibm.com>
…k wrap This commit completes the three-part fix for the Stickify allocation overflow chain (f006 + f018), building on the getsize_ztensor overflow guard already committed for f013. --- f006: Re-enable verify_transformed_descriptor dim/size checks --- The entire validation body of verify_transformed_descriptor() was wrapped in #if 0 since 2025-01-08 (commit 86dbaf0, issue onnx#3034) because the real zDNN runtime calls zdnn_get_max_for_dim() / zdnn_get_nnpa_max_tensor_size(), which read nnpa_query_result state that is not initialized in onnx-mlir's compile-time stickify build. The #if 0 block is removed and replaced with equivalent checks using compile-time constants from NNPALimit.hpp (NNPAGetMaxForDim / NNPA_MAXIMUM_TENSOR_SIZE), which are already used elsewhere in the NNPA path and do not require nnpa_query_result to be initialized. The zero-dim check and the total-size check were the only guards anywhere in the pipeline for those two conditions. Without them, attacker-controlled dims from a crafted .onnx Constant flow directly into getsize_ztensor / allochelper_ztensor_alloc with no bound at all, enabling the f013 overflow and the f018 wrap described below. --- f018: Fix dead SIZE_MAX check in malloc_aligned_4k --- The guard 'if (!size || size > SIZE_MAX)' has a dead second condition: size_t can never exceed SIZE_MAX by definition. The actual hazard is 'malloc(size + extra_allocation)' on the next line: if size lands within extra_allocation (4103 bytes) of SIZE_MAX, the addition wraps size_t to a tiny value, malloc succeeds with a tiny block, but the caller records the original near-SIZE_MAX value as buffer_size. Any subsequent write trusting buffer_size (memset / transform_ztensor) then writes far past the real allocation. Fix: replace 'size > SIZE_MAX' with 'size > SIZE_MAX - extra_allocation' so the check correctly catches the wrap before malloc is called. Signed-off-by: Jonathan Huang <jonathanhuang@ibm.com>
|
Can one of the admins verify this patch? |
SummaryThis PR fixes three related security defects in src/Accelerators/NNPA/Support/Stickify/Stickify.cpp that form a chain: the first disables the gate that should reject bad inputs, the second is an integer overflow that fires when bad inputs reach it, and the third is a dead safety check that was supposed to catch the overflow but never could. All three must be fixed together for the allocation to be safe. When onnx-mlir compiles a model targeting the NNPA hardware accelerator, weight constants are converted to the NNPA hardware layout ("stickified") at compile time. This process allocates a buffer sized from the tensor's declared dimensions, then writes the converted data into it. The three bugs below all affect that allocation sequence in the same file. f006 — The validation gate was entirely disabledverify_transformed_descriptor() is supposed to check that a tensor's dimensions are non-zero and within NNPA's hardware limits before any memory is allocated. The entire body of this function was wrapped in #if 0 since January 2025 (commit 86dbaf0, issue #3034) because the real zDNN library's limit-query functions require hardware state that isn't available at compile time. The function unconditionally returned OK for every input. f013 — The buffer size calculation could overflowgetsize_ztensor() computes how many bytes to allocate as a chain of five multiplications: dim4 × dim3 × ceil(dim2/32) × ceil(dim1/cells) × 4096. Each dimension is an attacker-controlled 32-bit integer. With f006's gate disabled, there was nothing to bound the inputs, so a crafted tensor could cause this product to wrap around uint64_t to a tiny value — for example, choosing dims whose product would be 2^64 + 4096 causes the result to be just 4096. That tiny value drives malloc, giving a 4 KB buffer, while transform_ztensor then writes data sized for the full logical dimensions into it — a heap buffer overflow. f018 — A safety check in the allocator was dead codemalloc_aligned_4k() needs a small amount of extra memory beyond the requested size (4103 bytes) for page-alignment bookkeeping, so it calls malloc(size + 4103). Before doing so it checks if (!size || size > SIZE_MAX). The second condition — size > SIZE_MAX — can never be true: size is of type size_t, and SIZE_MAX is by definition the largest value a size_t can hold. The check was therefore dead code. The real danger is the addition: if size is within 4103 of SIZE_MAX, then size + 4103 wraps around to a tiny number, malloc succeeds with a tiny block, but the caller records the original near-SIZE_MAX value as the buffer's size. Any subsequent write that trusts that recorded size walks far past the real allocation. |
| if (__builtin_mul_overflow( | ||
| size, static_cast<uint64_t>(tfrmd_desc->dim3), &size) || | ||
| __builtin_mul_overflow( | ||
| size, static_cast<uint64_t>(CEIL(number_of_sticks, |
There was a problem hiding this comment.
nit: maybe the ceil could be computed before the test to improve readability
There was a problem hiding this comment.
I was looking at using these, but code assistant indicated that there is already support in LLVM that is more portable:
LLVM — two candidates in [MathExtras.h]
(llvm-project/llvm/include/llvm/Support/MathExtras.h):
llvm::SaturatingMultiply in line :638
llvm::MulOverflow in line 753
3 participants
Analysis files for these three bugs: