Fix heap buffer overflow from reusing a packed sub-byte buffer for a full-byte tensor

onnxruntime/core/framework/allocation_planner.cc

@@ -496,14 +496,27 @@ class PlannerImpl {
     return true;
   }
 
-  /*! \brief Given a tensor-type, return the size of an element of the tensor.
+  /*! \brief Given a tensor-type, return the primitive element type of the tensor.
    */
-  static size_t GetElementSize(const DataType& tensor_type) {
+  static MLDataType GetPrimitiveElementType(const DataType& tensor_type) {
     MLDataType ml_data_type = DataTypeImpl::GetDataType(*tensor_type);
     const TensorTypeBase* tensor_type_base = ml_data_type->AsTensorType();
     ORT_ENFORCE(nullptr != tensor_type_base);
-    MLDataType elt_type = tensor_type_base->GetElementType();
-    return elt_type->Size();
+    return tensor_type_base->GetElementType();
+  }
+
+  /*! \brief Given a tensor-type, return the size in bytes of the C++ carrier used for an element.
+   */
+  static size_t GetElementSize(const DataType& tensor_type) {
+    return GetPrimitiveElementType(tensor_type)->Size();
+  }
+
+  /*! \brief Given a tensor-type, return how many logical (sub-byte) elements are packed into one
+   *  carrier element. Returns 1 for regular types and >1 for packed sub-byte types (e.g. 2 for int4/uint4).
+   */
+  static int32_t GetSubElemCount(const DataType& tensor_type) {
+    const auto* prim_type = GetPrimitiveElementType(tensor_type)->AsPrimitiveDataType();
+    return prim_type != nullptr ? prim_type->GetNumSubElems() : 1;
   }
 
   static bool SameSize(const TensorShapeProto& shape1, const onnxruntime::NodeArg& arg1,
@@ -515,6 +528,16 @@ class PlannerImpl {
     bool is_type1_string = arg1.TypeAsProto()->tensor_type().elem_type() == ONNX_NAMESPACE::TensorProto_DataType_STRING;
     bool is_type2_string = arg2.TypeAsProto()->tensor_type().elem_type() == ONNX_NAMESPACE::TensorProto_DataType_STRING;
 
+    // Packed sub-byte types (e.g. int4/uint4) share the same one-byte C++ carrier size as int8/uint8, but a
+    // carrier stores GetNumSubElems() logical elements, so the physical storage is ceil(N / sub_elems) bytes.
+    // Two tensors with equal logical shape and equal carrier size can therefore have different storage sizes
+    // (e.g. uint4[1024] needs 512 bytes while uint8[1024] needs 1024 bytes). Reusing the smaller buffer for the
+    // larger tensor produces a heap buffer overflow when the tensor is later written. Only treat the tensors as
+    // the same size when the sub-element packing density also matches, which guarantees identical storage bytes.
+    if (GetSubElemCount(ptype1) != GetSubElemCount(ptype2)) {
+      return false;
+    }
+
     // sizeof(std::string) = sizeof(double) on gcc 4.8.x on CentOS. This causes the allocation planner to reuse
     // a tensor of type double. This won't work for string tensors since they need to be placement new'ed.
     // If either of the tensors is a string, don't treat them the same. Moreover, reusing a string tensor for a string

onnxruntime/core/framework/execution_frame.cc

@@ -689,6 +689,22 @@ Status ExecutionFrame::AllocateMLValueTensorPreAllocateBuffer(OrtValue& ort_valu
         return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, message);
       }
     }
+
+    // Defense in depth: equal logical element counts do not guarantee equal physical storage. Packed sub-byte
+    // types (e.g. uint4[N] needs ceil(N/2) bytes) share the same carrier size as full-byte types (uint8[N] needs
+    // N bytes), so a reused buffer sized for the packed type is too small for the full-byte tensor and a later
+    // write would overflow it. Reject the reuse whenever the buffer cannot physically hold the requested tensor.
+    size_t required_storage_bytes = 0;
+    ORT_RETURN_IF_ERROR(
+        Tensor::CalculateTensorStorageSize(element_type, shape, /*alignment*/ 0, required_storage_bytes));
+    const size_t buffer_storage_bytes = reuse_tensor->SizeInBytes();
+    if (required_storage_bytes > buffer_storage_bytes) {
+      return ORT_MAKE_STATUS(
+          ONNXRUNTIME, FAIL, "Cannot re-use buffer: requested tensor needs ", required_storage_bytes,
+          " bytes of storage but the buffer being reused only has ", buffer_storage_bytes,
+          " bytes (buffer shape ", reuse_tensor->Shape(), ", requested shape ", shape,
+          "). This can happen when a packed sub-byte tensor is reused for a full-byte tensor of the same shape.");
+    }
   }
 
   void* reuse_buffer = reuse_tensor->MutableDataRaw();

onnxruntime/test/framework/allocation_planner_test.cc

@@ -24,6 +24,7 @@ using json = nlohmann::json;
 #include "core/util/thread_utils.h"
 
 #include "test/test_environment.h"
+#include "test/unittest_util/framework_test_utils.h"
 #include "test/util/include/asserts.h"
 #include "test/util/include/default_providers.h"
 #ifdef USE_CUDA
@@ -2142,5 +2143,105 @@ TEST(AllocationPlannerTest, AvoidReuseOfBufferForNodeOutputWithNoConsumers) {
 }
 #endif
 
+// Regression test for a heap buffer overflow caused by reusing a packed sub-byte buffer for a full-byte tensor.
+//
+// A packed sub-byte tensor (uint4[N]) needs ceil(N/2) storage bytes, while a full-byte tensor (uint8[N]) of the
+// same logical shape needs N bytes. Both types have a one-byte C++ carrier size, so the allocation planner's
+// SameSize() check used to treat them as the same size (carrier size 1 == 1 and identical shape) and let the
+// uint8 output reuse the smaller uint4 buffer. Writing the uint8 tensor into that half-sized buffer then
+// overflows it (CWE-131 -> CWE-787). The uint8 output must not reuse the uint4 buffer.
+//
+// Graph (opset 21):  X(float) -Cast-> A(uint4) -Cast-> B(float) -Cast-> C(uint8) -Cast-> Y(float)
+// A is fully consumed by the second Cast and freed, so before the fix the planner reused A's 512-byte buffer for
+// the 1023-byte uint8 tensor C.
+//
+// The dimension is deliberately odd so the ceiling division in the packed storage size (ceil(1023 / 2) = 512, not
+// 1023 / 2 = 511) is covered as well.
+TEST(AllocationPlannerTest, AvoidReuseOfPackedSubByteBufferForFullByteTensor) {
+  constexpr int64_t kDim = 1023;
+
+  auto make_tensor_type = [](TensorProto_DataType elem_type) {
+    TypeProto t;
+    t.mutable_tensor_type()->set_elem_type(elem_type);
+    t.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(kDim);
+    return t;
+  };
+
+  auto create_model = [&]() -> Model {
+    Model model("packed_subbyte_reuse", false, ModelMetaData(), PathString(),
+                IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, 21}}, {},
+                DefaultLoggingManager().DefaultLogger());
+    Graph& graph = model.MainGraph();
+
+    TypeProto float_type = make_tensor_type(TensorProto_DataType_FLOAT);
+    TypeProto uint4_type = make_tensor_type(TensorProto_DataType_UINT4);
+    TypeProto uint8_type = make_tensor_type(TensorProto_DataType_UINT8);
+
+    auto& X = graph.GetOrCreateNodeArg("X", &float_type);
+    auto& A = graph.GetOrCreateNodeArg("A", &uint4_type);  // packed sub-byte buffer: ceil(1023/2) = 512 bytes
+    auto& B = graph.GetOrCreateNodeArg("B", &float_type);  // consumes A -> A becomes dead here
+    auto& C = graph.GetOrCreateNodeArg("C", &uint8_type);  // full-byte buffer: 1023 bytes
+    auto& Y = graph.GetOrCreateNodeArg("Y", &float_type);
+
+    auto add_cast = [&graph](const std::string& name, NodeArg& in, NodeArg& out, TensorProto_DataType to) {
+      auto& node = graph.AddNode(name, "Cast", name, {&in}, {&out});
+      node.AddAttribute("to", static_cast<int64_t>(to));
+    };
+    add_cast("cast_to_uint4", X, A, TensorProto_DataType_UINT4);
+    add_cast("cast_a_to_float", A, B, TensorProto_DataType_FLOAT);
+    add_cast("cast_to_uint8", B, C, TensorProto_DataType_UINT8);
+    add_cast("cast_c_to_float", C, Y, TensorProto_DataType_FLOAT);
+
+    graph.SetInputs({&X});
+    graph.SetOutputs({&Y});
+    EXPECT_STATUS_OK(graph.Resolve());
+    return model;
+  };
+
+  SessionOptions so;
+  // Keep memory reuse on (default) and avoid graph optimizations that could rewrite the Cast chain.
+  so.graph_optimization_level = TransformerLevel::Default;
+  InferenceSession sess{so, GetEnvironment()};
+
+  std::string serialized;
+  ASSERT_TRUE(create_model().ToProto().SerializeToString(&serialized));
+  std::stringstream sstr(serialized);
+  ASSERT_STATUS_OK(sess.Load(sstr));
+  ASSERT_STATUS_OK(sess.Initialize());
+
+  const auto& session_state = sess.GetSessionState();
+  const auto& ort_value_index_map = session_state.GetOrtValueNameIdxMap();
+  const SequentialExecutionPlan* plan = session_state.GetExecutionPlan();
+
+  OrtValueIndex a_index, c_index;
+  ASSERT_STATUS_OK(ort_value_index_map.GetIdx("A", a_index));
+  ASSERT_STATUS_OK(ort_value_index_map.GetIdx("C", c_index));
+
+  // The uint8 tensor C must never be planned to reuse the (smaller) uint4 tensor A's buffer.
+  const auto& c_plan = plan->allocation_plan[c_index];
+  const bool reuses_a = (c_plan.alloc_kind == AllocKind::kReuse) && (c_plan.reused_buffer == a_index);
+  EXPECT_FALSE(reuses_a) << "uint8[" << kDim << "] output must not reuse the uint4[" << kDim << "] buffer";
+
+  // The model must still run and produce correct results after round-tripping through uint4 and uint8.
+  constexpr float kValue = 3.0f;
+  std::vector<int64_t> dims{kDim};
+  std::vector<float> input_data(static_cast<size_t>(kDim), kValue);
+  OrtValue input_value;
+  CreateMLValue<float>(TestCPUExecutionProvider()->CreatePreferredAllocators()[0], dims, input_data, &input_value);
+
+  NameMLValMap feeds{{"X", input_value}};
+  std::vector<std::string> output_names{"Y"};
+  std::vector<OrtValue> fetches;
+  ASSERT_STATUS_OK(sess.Run(feeds, output_names, &fetches));
+
+  ASSERT_EQ(fetches.size(), 1u);
+  const Tensor& out = fetches[0].Get<Tensor>();
+  ASSERT_EQ(out.Shape().Size(), kDim);
+  const float* out_data = out.Data<float>();
+  for (int64_t i = 0; i < kDim; ++i) {
+    ASSERT_EQ(out_data[i], kValue) << "mismatch at index " << i;
+  }
+}
+
 }  // namespace test
 }  // namespace onnxruntime

onnxruntime/test/framework/execution_frame_test.cc

@@ -3,6 +3,7 @@
 
 #include "core/common/span_utils.h"
 #include "core/framework/execution_frame.h"
+#include "core/framework/int4.h"
 #include "core/framework/op_kernel.h"
 #include "core/framework/session_state.h"
 #include "core/graph/model.h"
@@ -188,6 +189,104 @@ TEST_F(ExecutionFrameTest, OutputShapeValidationTest) {
   ASSERT_STATUS_OK(frame.GetOrCreateNodeOutputMLValue(int(node->Index()), 1, &actual_shape_diff_from_input, p_ml_value, *node));
 }
 
+// Directly exercises the runtime capacity guard in ExecutionFrame::AllocateMLValueTensorPreAllocateBuffer.
+//
+// A packed sub-byte tensor (uint4[N]) has the same logical element count and the same one-byte C++ carrier size as
+// a full-byte tensor (uint8[N]), but only needs ceil(N/2) storage bytes instead of N. Reusing the smaller uint4
+// buffer for the uint8 tensor would overflow it on the first write, so the reuse must be rejected. The opposite
+// direction (a uint8 buffer reused by a uint4 tensor) is safe and must still be allowed.
+//
+// An odd dimension is used on purpose so the ceiling division in the storage size calculation is covered.
+TEST_F(ExecutionFrameTest, PreAllocatedBufferTooSmallForSubByteTypeTest) {
+  constexpr int64_t kDim = 1023;
+
+  onnxruntime::Model model("test", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(),
+                           {{kOnnxDomain, 12}}, {}, DefaultLoggingManager().DefaultLogger());
+  onnxruntime::Graph& graph = model.MainGraph();
+  TypeProto tensor_float;
+  tensor_float.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT);
+  onnxruntime::NodeArg input_def("X", &tensor_float), output_def("Y", &tensor_float);
+
+  onnxruntime::Node* node = &graph.AddNode("node1", "Relu", "Relu operator", ArgMap{&input_def}, ArgMap{&output_def});
+  node->SetExecutionProviderType(kCpuExecutionProvider);
+  ASSERT_STATUS_OK(graph.Resolve());
+
+  auto cpu_xp = CreateCPUExecutionProvider();
+  auto xp_typ = cpu_xp->Type();
+  ExecutionProviders execution_providers;
+  ASSERT_STATUS_OK(execution_providers.Add(xp_typ, std::move(cpu_xp)));
+  KernelRegistryManager kernel_registry_manager;
+  ASSERT_STATUS_OK(kernel_registry_manager.RegisterKernels(execution_providers));
+
+  DataTransferManager dtm;
+  ExternalDataLoaderManager edlm;
+  profiling::Profiler profiler;
+
+  SessionOptions sess_options;
+  sess_options.enable_mem_pattern = true;
+  sess_options.execution_mode = ExecutionMode::ORT_SEQUENTIAL;
+  sess_options.use_deterministic_compute = false;
+  sess_options.enable_mem_reuse = true;
+
+  SessionState state(graph, execution_providers, &tp_, nullptr, dtm, edlm,
+                     DefaultLoggingManager().DefaultLogger(), profiler, sess_options);
+
+  node->SetExecutionProviderType(xp_typ);
+
+  ASSERT_STATUS_OK(state.FinalizeSessionState(ORT_TSTR(""), kernel_registry_manager));
+
+  const auto& memory_info = execution_providers.Get(xp_typ)->GetOrtDeviceByMemType(OrtMemTypeDefault);
+  const TensorShape shape(std::vector<int64_t>{kDim});
+  MLDataType uint4_type = DataTypeImpl::GetType<UInt4x2>();
+  MLDataType uint8_type = DataTypeImpl::GetType<uint8_t>();
+
+  ASSERT_EQ(Tensor::CalculateTensorStorageSize(uint4_type, shape), static_cast<size_t>((kDim + 1) / 2));
+  ASSERT_EQ(Tensor::CalculateTensorStorageSize(uint8_type, shape), static_cast<size_t>(kDim));
+
+  // A uint8 tensor must not re-use a uint4 buffer of the same logical shape: the buffer is only half the size.
+  {
+    vector<OrtValue> outputs;
+    ExecutionFrame frame({}, {}, {}, outputs, {},
+#ifdef ORT_ENABLE_STREAM
+                         {},
+#endif
+                         state);
+
+    int start_index = frame.GetNodeOffset(node->Index());
+    ASSERT_EQ(start_index, 0);
+
+    OrtValue& uint4_value = *frame.GetMutableNodeInputOrOutputMLValue(start_index);
+    ASSERT_STATUS_OK(frame.AllocateMLValueTensorSelfOwnBuffer(uint4_value, start_index, uint4_type, memory_info, shape));
+
+    OrtValue& uint8_value = *frame.GetMutableNodeInputOrOutputMLValue(start_index + 1);
+    const Status status = frame.AllocateMLValueTensorPreAllocateBuffer(uint8_value, start_index, uint8_type,
+                                                                       memory_info, shape);
+    ASSERT_FALSE(status.IsOK());
+    EXPECT_THAT(status.ErrorMessage(), ::testing::HasSubstr("Cannot re-use buffer"));
+    EXPECT_FALSE(uint8_value.IsAllocated());
+  }
+
+  // The reverse direction is safe: a uint4 tensor fits in a uint8 buffer of the same logical shape.
+  {
+    vector<OrtValue> outputs;
+    ExecutionFrame frame({}, {}, {}, outputs, {},
+#ifdef ORT_ENABLE_STREAM
+                         {},
+#endif
+                         state);
+
+    int start_index = frame.GetNodeOffset(node->Index());
+
+    OrtValue& uint8_value = *frame.GetMutableNodeInputOrOutputMLValue(start_index);
+    ASSERT_STATUS_OK(frame.AllocateMLValueTensorSelfOwnBuffer(uint8_value, start_index, uint8_type, memory_info, shape));
+
+    OrtValue& uint4_value = *frame.GetMutableNodeInputOrOutputMLValue(start_index + 1);
+    ASSERT_STATUS_OK(frame.AllocateMLValueTensorPreAllocateBuffer(uint4_value, start_index, uint4_type,
+                                                                  memory_info, shape));
+    EXPECT_EQ(uint4_value.Get<Tensor>().DataRaw(), uint8_value.Get<Tensor>().DataRaw());
+  }
+}
+
 TEST_F(ExecutionFrameTest, FeedInDataTest) {
   onnxruntime::Model model("test", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(),
                            std::unordered_map<std::string, int>{{"", 10}}, {},