Validate constant tensor byte size in DML OnnxTensorWrapper
OnnxTensorWrapper wraps an onnx::TensorProto and hands out a raw pointer
via GetData() together with dimensions from GetShape(). Consumers size
their reads from the shape and the declared element type, never from the
actual size of the backing buffer. When a TensorProto declares dims and a
data type whose implied byte size exceeds the bytes actually present, the
consumer reads past the end of the heap allocation.
For graph initializers this is already mitigated: ONNX shape inference
calls ParseData, which rejects a raw_data field shorter than
product(dims) * sizeof(element). That covers ops whose inference reads
constant inputs (Slice, Pad, Tile, Expand, Reshape, Resize, Split).
Attribute tensors are not covered by any of that. OpNodeInfoWrapper::
GetTensorAttribute builds an OnnxTensorWrapper directly from
AttributeProto::t(). An attribute tensor is never a graph initializer, so
Graph::ConvertInitializersIntoOrtValues and its embedded-data validation
never see it, and ConstantOfShape's ONNX shape inference only reads the
attribute's data_type - never its data. DmlOperatorConstantOfShape then
checks elementCount == 1 from the dims and memcpy's a byte count derived
from the declared data type out of GetData(). A `value` attribute of
{INT64, dims=[1], raw_data = 1 byte} produces a 7-byte heap over-read
whose contents become the GPU fill pattern, and therefore appear directly
in the model output.
Add VerifyTensorProtoFitsInBuffer, called at the end of the
OnnxTensorWrapper constructor so it covers the external-data, raw_data,
and typed-field branches alike. It computes the required byte size with
GetSizeInBytesFromTensorProto<0> and throws E_INVALIDARG when the backing
buffer is smaller. STRING is exempt because its size query returns
elem_count * sizeof(std::string) rather than a serialized byte size, and
the complex types are sized inline because that helper reports them as
NOT_IMPLEMENTED even though ToMLTensorDataType maps them. Every other
size-query failure - unrepresentable dims, byte-count overflow, an
element type this EP does not admit - is treated as malformed rather than
skipped, so the bounds check is total over the types that can reach a
consumer.
The check is a lower bound only. ValidateEmbeddedTensorProtoDataSizeAndShape
was deliberately not reused because it also enforces exact equality and a
2 GiB cap; RegisterDynamicKernel re-packs arbitrarily large weights into
raw_data protos that reach this constructor, so the cap risks regressing
legitimate models.
Tests: onnxruntime/test/providers/dml_onnx_tensor_wrapper_test.cc builds a
ConstantOfShape model in memory with a truncated `value` attribute and
asserts that session initialization now fails; with the constructor check
removed that test fails, confirming the session previously initialized
while over-reading. A positive control loads, initializes, and runs a
well-formed equivalent and checks the output. ConstantFolding is disabled
in both so the node reaches the DML EP instead of being folded on CPU at
Level1. The constructor's typed-field branch is guarded too but has no
test, because every malformed proto that would reach it is rejected
earlier by the ONNX checker or by DmlOperatorConstantOfShape's own
element-count assertion; the file documents this.
Verified with onnxruntime_test_all (1896 tests, 1879 passed, 0 failed) and
onnxruntime_provider_test --gtest_filter=*Dml* (14 passed) on a machine
with a real DirectML device.
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
