@@ -83,7 +83,24 @@ std::istream& operator>>(std::istream& stream, SharedContext::SharedWeights::Met
83
83
stream >> value.size ;
84
84
size_t num_dimensions;
85
85
stream >> num_dimensions;
86
- value.dimensions .resize (num_dimensions);
86
+
87
+ if (stream.fail ()) {
88
+ ORT_THROW (" Error: Failed to read num_dimensions from stream." );
89
+ }
90
+
91
+ constexpr size_t MAX_SAFE_DIMENSIONS = 1024 ;
92
+
93
+ size_t safe_num_dimensions = num_dimensions;
94
+
95
+ if (num_dimensions == 0 || safe_num_dimensions > MAX_SAFE_DIMENSIONS) {
96
+ ORT_THROW (" Invalid number of dimensions provided." );
97
+ }
98
+ try {
99
+ value.dimensions .resize (safe_num_dimensions);
100
+ } catch (const std::bad_alloc&) {
101
+ ORT_THROW (" Error: Memory allocation failed while resizing dimensions." );
102
+ }
103
+
87
104
for (auto & dim : value.dimensions ) {
88
105
stream >> dim;
89
106
}
@@ -235,23 +252,23 @@ int GetFirstAvailableDevice(SessionContext& session_context) {
235
252
void FillOutputsWithConstantData (std::shared_ptr<ov::Node> node, Ort::UnownedValue& out_tensor) {
236
253
switch (node->get_element_type ()) {
237
254
case ov::element::Type_t::f32: {
238
- FillOutputHelper<float >(out_tensor, node);
255
+ FillOutputHelper<float >(out_tensor, std::move ( node) );
239
256
break ;
240
257
}
241
258
case ov::element::Type_t::boolean: {
242
- FillOutputHelper<char >(out_tensor, node);
259
+ FillOutputHelper<char >(out_tensor, std::move ( node) );
243
260
break ;
244
261
}
245
262
case ov::element::Type_t::i32: {
246
- FillOutputHelper<int32_t >(out_tensor, node);
263
+ FillOutputHelper<int32_t >(out_tensor, std::move ( node) );
247
264
break ;
248
265
}
249
266
case ov::element::Type_t::i64: {
250
- FillOutputHelper<int64_t >(out_tensor, node);
267
+ FillOutputHelper<int64_t >(out_tensor, std::move ( node) );
251
268
break ;
252
269
}
253
270
case ov::element::Type_t::f16: {
254
- FillOutputHelper<float >(out_tensor, node);
271
+ FillOutputHelper<float >(out_tensor, std::move ( node) );
255
272
break ;
256
273
}
257
274
default :
0 commit comments