Skip to content

Commit a83a158

Browse files
Verify in-memory references when loading a model (#26764)
### Description When loading the model verify that in-memory references point to an existing and valid OrtValue as if created my ModelEditor. ### Motivation and Context This prevents reading from arbitrary memory location when loading the model. --------- Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
1 parent 334bc04 commit a83a158

File tree

4 files changed

+103
-1
lines changed

4 files changed

+103
-1
lines changed

onnxruntime/core/graph/graph.cc

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3730,8 +3730,24 @@ Status Graph::ConvertInitializersIntoOrtValues() {
37303730
auto& graph_proto = *graph.graph_proto_;
37313731
for (int i = 0, lim = graph_proto.initializer_size(); i < lim; ++i) {
37323732
auto& tensor_proto = *graph_proto.mutable_initializer(i);
3733+
37333734
if (utils::HasExternalData(tensor_proto)) {
3734-
continue; // ignore data on disk, that will be loaded either by EP or at session_state finalize
3735+
if (utils::HasExternalDataInMemory(tensor_proto)) {
3736+
// This can happen when the model is created with ModelEditor.
3737+
// We want to guard against malicious models with arbitrary in-memory references.
3738+
if (OrtValue v; GetOrtValueInitializer(tensor_proto.name(), v)) {
3739+
ORT_RETURN_IF_NOT(graph_utils::CheckInMemoryDataMatch(tensor_proto, v.Get<Tensor>()),
3740+
"In-memory data mismatch for initializer: ", tensor_proto.name(),
3741+
" this is an invalid model");
3742+
} else {
3743+
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL,
3744+
"The model contains initializers with arbitrary in-memory references.",
3745+
"This is an invalid model.");
3746+
}
3747+
}
3748+
// ignore data on disk, that will be loaded either by EP or at session_state finalize
3749+
// ignore valid in-memory references
3750+
continue;
37353751
}
37363752

37373753
size_t size_in_bytes = 0;
@@ -3830,6 +3846,10 @@ Status Graph::AddInitializedOrtValue(const ONNX_NAMESPACE::TensorProto& tensor_p
38303846
const auto proto_shape = utils::GetTensorShapeFromTensorProto(tensor_proto);
38313847
ORT_RETURN_IF_NOT(proto_shape == tensor.Shape(), "Shape mismatch with ortvalue_initializer");
38323848

3849+
const bool data_ptr_match = graph_utils::CheckInMemoryDataMatch(tensor_proto, tensor);
3850+
ORT_RETURN_IF_NOT(data_ptr_match,
3851+
"In-memory data pointer mismatch between tensor proto and ortvalue_initializer");
3852+
38333853
ortvalue_initializers_.insert_or_assign(tensor_proto.name(), ortvalue_initializer);
38343854
} else {
38353855
ORT_ENFORCE(ortvalue_initializers_.count(tensor_proto.name()) == 0,

onnxruntime/test/shared_lib/test_inference.cc

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4761,3 +4761,36 @@ TEST(CApiTest, custom_cast) {
47614761
inputs, "output", expected_dims_y, expected_values_y, 0,
47624762
custom_op_domain, nullptr);
47634763
}
4764+
4765+
TEST(CApiTest, ModelWithMaliciousExternalDataShouldFailToLoad) {
4766+
// Attempt to create an ORT session with the malicious model
4767+
// This should fail due to the invalid external data reference
4768+
constexpr const ORTCHAR_T* model_path = TSTR("testdata/test_evil_weights.onnx");
4769+
4770+
Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "test");
4771+
Ort::SessionOptions session_options;
4772+
4773+
bool exception_thrown = false;
4774+
std::string exception_message;
4775+
4776+
try {
4777+
// This should throw an exception due to malicious external data
4778+
Ort::Session session(env, model_path, session_options);
4779+
} catch (const Ort::Exception& e) {
4780+
exception_thrown = true;
4781+
exception_message = e.what();
4782+
} catch (const std::exception& e) {
4783+
exception_thrown = true;
4784+
exception_message = e.what();
4785+
}
4786+
4787+
// Verify that loading the model failed
4788+
EXPECT_TRUE(exception_thrown) << "Expected model loading to fail due to malicious external data";
4789+
4790+
// Verify that the exception message indicates security or external data issues
4791+
EXPECT_TRUE(exception_message.find("in-memory") != std::string::npos ||
4792+
exception_message.find("references") != std::string::npos ||
4793+
exception_message.find("invalid") != std::string::npos ||
4794+
exception_message.find("model") != std::string::npos)
4795+
<< "Exception message should indicate external data or security issue. Got: " << exception_message;
4796+
}
212 Bytes
Binary file not shown.
Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
import onnx
2+
3+
4+
def create_exp_model():
5+
inputs = []
6+
nodes = []
7+
tensors = []
8+
outputs = []
9+
10+
input_ = onnx.helper.make_tensor_value_info("input", onnx.TensorProto.INT64, [None, None])
11+
inputs.append(input_)
12+
13+
evil_tensor = onnx.helper.make_tensor(
14+
name="evil_weights",
15+
data_type=onnx.TensorProto.INT64,
16+
# dims=[100, 100],
17+
dims=[10],
18+
vals=[],
19+
)
20+
tensors.append(evil_tensor)
21+
evil_tensor.data_location = onnx.TensorProto.EXTERNAL
22+
entry1 = evil_tensor.external_data.add()
23+
entry1.key = "location"
24+
entry1.value = "*/_ORT_MEM_ADDR_/*"
25+
entry2 = evil_tensor.external_data.add()
26+
entry2.key = "offset"
27+
entry2.value = "12230656"
28+
entry3 = evil_tensor.external_data.add()
29+
entry3.key = "length"
30+
entry3.value = "80"
31+
32+
tensors.append(onnx.helper.make_tensor(name="0x1", data_type=onnx.TensorProto.INT64, dims=[1], vals=[0x1]))
33+
nodes.append(onnx.helper.make_node(op_type="Add", inputs=["evil_weights", "0x1"], outputs=["output"]))
34+
35+
outputs.append(onnx.helper.make_tensor_value_info("output", onnx.TensorProto.INT64, [10]))
36+
37+
graph = onnx.helper.make_graph(nodes, "test", inputs, outputs, tensors)
38+
model = onnx.helper.make_model(
39+
graph,
40+
opset_imports=[onnx.helper.make_opsetid("", 18), onnx.helper.make_opsetid("ai.onnx.ml", 3)],
41+
ir_version=11,
42+
)
43+
44+
return model
45+
46+
47+
if __name__ == "__main__":
48+
model = create_exp_model()
49+
onnx.save(model, "test_evil_weights.onnx")

0 commit comments

Comments
 (0)