Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# Copyright (c) ONNX Project Contributors
2#
3# SPDX-License-Identifier: Apache-2.0
4from __future__ import annotations
6import os
7import re
8import uuid
9import warnings
10from itertools import chain
11from typing import IO, TYPE_CHECKING
13import onnx.checker as onnx_checker
14import onnx.onnx_cpp2py_export.checker as c_checker
15from onnx.onnx_pb import (
16 AttributeProto,
17 FunctionProto,
18 GraphProto,
19 ModelProto,
20 TensorProto,
21)
23if TYPE_CHECKING:
24 from collections.abc import Callable, Iterable
27def _open_external_data_fd(
28 base_dir: str, location: str, tensor_name: str, read_only: bool
29) -> int:
30 """Open external data via C++ and return a CRT file descriptor."""
31 return c_checker._open_external_data(base_dir, location, tensor_name, read_only)
34# Security: 3-layer defense against malicious external_data entries (GHSA-538c-55jv-c5g9)
35#
36# Layer 1 (here) — Attribute whitelist: Only spec-defined keys are accepted.
37# Unknown keys are warned and ignored, preventing arbitrary attribute injection (CWE-915).
38#
39# Layer 2 (ExternalDataInfo.__init__) — Bounds validation: offset and length must be
40# non-negative integers. Catches invalid values at parse time (CWE-400).
41#
42# Layer 3 (load_external_data_for_tensor) — File-size validation: offset and length are
43# checked against actual file size before reading. This is the critical safety net that
44# prevents memory exhaustion regardless of how the model was constructed (CWE-400).
45#
46# 'basepath' is included because set_external_data() and model_container
47# write it to protobuf entries; it must survive save/load round-trips.
48_ALLOWED_EXTERNAL_DATA_KEYS = frozenset(
49 {"location", "offset", "length", "checksum", "basepath"}
50)
51_SORTED_ALLOWED_KEYS = sorted(_ALLOWED_EXTERNAL_DATA_KEYS)
52_MAX_UNKNOWN_KEYS_IN_WARNING = 10
53_MAX_KEY_DISPLAY_LENGTH = 100
55# Largest alignment boundary offsets may use; see `ExternalData.md`
56_MAX_EXTERNAL_DATA_PADDING = 64 * 1024
59class ExternalDataInfo:
60 def __init__(self, tensor: TensorProto) -> None:
61 self.location = ""
62 self.offset = None
63 self.length = None
64 self.checksum = None
65 self.basepath = ""
67 unknown_keys: set[str] = set()
68 unknown_key_count = 0
69 for entry in tensor.external_data:
70 # Layer 1: reject unknown keys (CWE-915 defense-in-depth)
71 if entry.key in _ALLOWED_EXTERNAL_DATA_KEYS:
72 setattr(self, entry.key, entry.value)
73 else:
74 unknown_key_count += 1
75 if len(unknown_keys) < _MAX_UNKNOWN_KEYS_IN_WARNING:
76 truncated = entry.key[:_MAX_KEY_DISPLAY_LENGTH]
77 if len(entry.key) > _MAX_KEY_DISPLAY_LENGTH:
78 truncated += "..."
79 unknown_keys.add(truncated)
81 if unknown_keys:
82 shown = sorted(unknown_keys)
83 extra = unknown_key_count - len(shown)
84 key_list = repr(shown)
85 if extra > 0:
86 key_list += f" and {extra} more"
87 warnings.warn(
88 f"Ignoring unknown external data key(s) {key_list} "
89 f"for tensor {tensor.name!r}. "
90 f"Allowed keys: {_SORTED_ALLOWED_KEYS}",
91 stacklevel=2,
92 )
94 if self.offset is not None:
95 self.offset = int(self.offset)
96 if self.offset < 0:
97 raise ValueError(
98 f"External data offset must be non-negative, got {self.offset} "
99 f"for tensor {tensor.name!r}"
100 )
102 if self.length is not None:
103 self.length = int(self.length)
104 if self.length < 0:
105 raise ValueError(
106 f"External data length must be non-negative, got {self.length} "
107 f"for tensor {tensor.name!r}"
108 )
111def _validate_external_data_file_bounds(
112 data_file: IO[bytes],
113 info: ExternalDataInfo,
114 tensor_name: str,
115) -> bytes:
116 """Validate offset/length against actual file size and read data.
118 Layer 3 defense-in-depth (CWE-400): prevents memory exhaustion even if the
119 model was crafted via direct protobuf APIs that bypass Python parsing.
121 Returns the raw bytes read from the file.
122 """
123 file_size = os.fstat(data_file.fileno()).st_size
125 if info.offset is not None:
126 if info.offset > file_size:
127 raise ValueError(
128 f"External data offset ({info.offset}) exceeds file size "
129 f"({file_size}) for tensor {tensor_name!r}"
130 )
131 data_file.seek(info.offset)
133 if info.length is not None:
134 read_start = info.offset if info.offset is not None else 0
135 available = file_size - read_start
136 if info.length > available:
137 raise ValueError(
138 f"External data length ({info.length}) exceeds available data "
139 f"({available} bytes from offset {read_start}) "
140 f"for tensor {tensor_name!r}"
141 )
142 return data_file.read(info.length)
143 return data_file.read()
146def _read_external_data_bytes(tensor: TensorProto, base_dir: str) -> bytes:
147 """Reads and returns the raw bytes stored in the external file for tensor.
149 Unlike load_external_data_for_tensor, this does not mutate tensor: the
150 tensor's data_location and external_data fields are left untouched, so
151 the tensor still faithfully describes where its data lives on disk.
153 Arguments:
154 tensor: a TensorProto object.
155 base_dir: directory that contains the external data.
157 Returns:
158 The raw bytes read from the external data file.
159 """
160 info = ExternalDataInfo(tensor)
161 fd = _open_external_data_fd(base_dir, info.location, tensor.name, True)
162 with os.fdopen(fd, "rb") as data_file:
163 return _validate_external_data_file_bounds(data_file, info, tensor.name)
166def load_external_data_for_tensor(tensor: TensorProto, base_dir: str) -> None:
167 """Loads data from an external file for tensor and converts it to an in-memory tensor.
168 Ideally TensorProto should not hold any raw data but if it does it will be ignored.
170 After this call, tensor.raw_data holds the loaded bytes and tensor no longer
171 describes external data: data_location is reset to DEFAULT and external_data is
172 cleared.
174 Arguments:
175 tensor: a TensorProto object.
176 base_dir: directory that contains the external data.
177 """
178 tensor.raw_data = _read_external_data_bytes(tensor, base_dir)
179 tensor.data_location = TensorProto.DEFAULT
180 del tensor.external_data[:]
183def load_external_data_for_model(model: ModelProto, base_dir: str) -> None:
184 """Loads external tensors into model
186 Arguments:
187 model: ModelProto to load external data to
188 base_dir: directory that contains external data
189 """
190 for tensor in _get_all_tensors(model):
191 if uses_external_data(tensor):
192 load_external_data_for_tensor(tensor, base_dir)
195def set_external_data(
196 tensor: TensorProto,
197 location: str,
198 offset: int | None = None,
199 length: int | None = None,
200 checksum: str | None = None,
201 basepath: str | None = None,
202) -> None:
203 if not tensor.HasField("raw_data"):
204 raise ValueError(
205 f"Tensor {tensor.name} does not have raw_data field. Cannot set external data for this tensor."
206 )
208 del tensor.external_data[:]
209 tensor.data_location = TensorProto.EXTERNAL
210 for k, v in {
211 "location": location,
212 "offset": int(offset) if offset is not None else None,
213 "length": int(length) if length is not None else None,
214 "checksum": checksum,
215 "basepath": basepath,
216 }.items():
217 if v is not None:
218 entry = tensor.external_data.add()
219 entry.key = k
220 entry.value = str(v)
223def convert_model_to_external_data(
224 model: ModelProto,
225 all_tensors_to_one_file: bool = True,
226 location: str | None = None,
227 size_threshold: int = 1024,
228 convert_attribute: bool = False,
229) -> None:
230 """Call to set all tensors with raw data as external data. This call should precede 'save_model'.
231 'save_model' saves all the tensors data as external data after calling this function.
233 Arguments:
234 model (ModelProto): Model to be converted.
235 all_tensors_to_one_file (bool): If true, save all tensors to one external file specified by location.
236 If false, save each tensor to a file named with the tensor name.
237 location: specify the external file relative to the model that all tensors to save to.
238 Path is relative to the model path.
239 If not specified, will use the model name.
240 size_threshold: Threshold for size of data. Only when tensor's data is >= the size_threshold
241 it will be converted to external data. To convert every tensor with raw data to external data set size_threshold=0.
242 convert_attribute (bool): If true, convert all tensors to external data
243 If false, convert only non-attribute tensors to external data
245 Raise:
246 ValueError: If location is not a relative path.
247 FileExistsError: If a file already exists in location.
248 """
249 tensors = _get_initializer_tensors(model)
250 if convert_attribute:
251 tensors = _get_all_tensors(model)
253 if all_tensors_to_one_file:
254 file_name = str(uuid.uuid1()) + ".data"
255 if location:
256 if os.path.isabs(location):
257 raise ValueError(
258 "location must be a relative path that is relative to the model path."
259 )
260 if os.path.exists(location):
261 raise FileExistsError(f"External data file exists in {location}.")
262 file_name = location
263 for tensor in tensors:
264 if tensor.HasField("raw_data") and len(tensor.raw_data) >= size_threshold:
265 set_external_data(tensor, file_name)
266 else:
267 for tensor in tensors:
268 if tensor.HasField("raw_data") and len(tensor.raw_data) >= size_threshold:
269 tensor_location = tensor.name
270 if not _is_valid_filename(tensor_location):
271 tensor_location = str(uuid.uuid1())
272 set_external_data(tensor, tensor_location)
275def convert_model_from_external_data(model: ModelProto) -> None:
276 """Call to set all tensors which use external data as embedded data.
277 save_model saves all the tensors data as embedded data after
278 calling this function.
280 Arguments:
281 model (ModelProto): Model to be converted.
282 """
283 for tensor in _get_all_tensors(model):
284 if uses_external_data(tensor):
285 if not tensor.HasField("raw_data"):
286 raise ValueError("raw_data field doesn't exist.")
287 del tensor.external_data[:]
288 tensor.data_location = TensorProto.DEFAULT
291def save_external_data(tensor: TensorProto, base_path: str) -> None:
292 """Writes tensor data to an external file according to information in the `external_data` field.
293 The function checks the external is a valid name and located in folder `base_path`.
295 Arguments:
296 tensor (TensorProto): Tensor object to be serialized
297 base_path: System path of a folder where tensor data is to be stored
299 Raises:
300 ValueError: If the external file is invalid.
301 """
302 info = ExternalDataInfo(tensor)
304 if not tensor.HasField("raw_data"):
305 raise onnx_checker.ValidationError("raw_data field doesn't exist.")
307 fd = _open_external_data_fd(base_path, info.location, tensor.name, False)
308 with os.fdopen(fd, "r+b") as data_file:
309 data_file.seek(0, 2)
310 if info.offset is not None:
311 # A tensor is written at the end of the file, optionally bumped forward by
312 # an alignment gap. The gap only aligns tensors that share a file, so it is
313 # always under one boundary. An offset before the end would overwrite data
314 # already in the file; one far past it would pad with an unbounded, untrusted
315 # amount of zeros. Reject both instead of writing.
316 file_size = data_file.tell()
317 padding = info.offset - file_size
318 if not 0 <= padding <= _MAX_EXTERNAL_DATA_PADDING:
319 raise onnx_checker.ValidationError(
320 f"External data offset ({info.offset}) for tensor {tensor.name!r} "
321 f"must be between the current file size ({file_size}) and "
322 f"{file_size + _MAX_EXTERNAL_DATA_PADDING}."
323 )
324 if padding > 0:
325 data_file.write(b"\0" * padding)
327 data_file.seek(info.offset)
328 offset = data_file.tell()
329 data_file.write(tensor.raw_data)
330 set_external_data(tensor, info.location, offset, data_file.tell() - offset)
333def _get_all_tensors(onnx_model_proto: ModelProto) -> Iterable[TensorProto]:
334 """Scan an ONNX model for all tensors and return as an iterator."""
335 return chain(
336 _get_initializer_tensors(onnx_model_proto),
337 _get_attribute_tensors(onnx_model_proto),
338 )
341def _recursive_attribute_processor(
342 attribute: AttributeProto, func: Callable[[GraphProto], Iterable[TensorProto]]
343) -> Iterable[TensorProto]:
344 """Create an iterator through processing ONNX model attributes with functor."""
345 if attribute.type == AttributeProto.GRAPH:
346 yield from func(attribute.g)
347 if attribute.type == AttributeProto.GRAPHS:
348 for graph in attribute.graphs:
349 yield from func(graph)
352def _get_initializer_tensors_from_graph(graph: GraphProto, /) -> Iterable[TensorProto]:
353 """Create an iterator of initializer tensors from ONNX model graph."""
354 yield from graph.initializer
355 for node in graph.node:
356 for attribute in node.attribute:
357 yield from _recursive_attribute_processor(
358 attribute, _get_initializer_tensors_from_graph
359 )
362def _get_initializer_tensors(onnx_model_proto: ModelProto) -> Iterable[TensorProto]:
363 """Create an iterator of initializer tensors from ONNX model."""
364 yield from _get_initializer_tensors_from_graph(onnx_model_proto.graph)
367def _get_attribute_tensors_from_graph(
368 graph_or_function: GraphProto | FunctionProto, /
369) -> Iterable[TensorProto]:
370 """Create an iterator of tensors from node attributes of an ONNX model graph/function."""
371 for node in graph_or_function.node:
372 for attribute in node.attribute:
373 if attribute.HasField("t"):
374 yield attribute.t
375 yield from attribute.tensors
376 yield from _recursive_attribute_processor(
377 attribute, _get_attribute_tensors_from_graph
378 )
381def _get_attribute_tensors(onnx_model_proto: ModelProto) -> Iterable[TensorProto]:
382 """Create an iterator of tensors from node attributes of an ONNX model."""
383 yield from _get_attribute_tensors_from_graph(onnx_model_proto.graph)
384 for function in onnx_model_proto.functions:
385 yield from _get_attribute_tensors_from_graph(function)
388def _is_valid_filename(filename: str) -> bool:
389 """Utility to check whether the provided filename is valid."""
390 exp = re.compile('^[^<>:;,?"*|/]+$')
391 match = exp.match(filename)
392 return bool(match)
395def uses_external_data(tensor: TensorProto) -> bool:
396 """Returns true if the tensor stores data in an external location."""
397 return (
398 tensor.HasField("data_location")
399 and tensor.data_location == TensorProto.EXTERNAL
400 )
403def remove_external_data_field(tensor: TensorProto, field_key: str) -> None:
404 """Removes a field from a Tensor's external_data key-value store.
406 Modifies tensor object in place.
408 Arguments:
409 tensor (TensorProto): Tensor object from which value will be removed
410 field_key (string): The key of the field to be removed
411 """
412 for i, field in enumerate(tensor.external_data):
413 if field.key == field_key:
414 del tensor.external_data[i]
417def write_external_data_tensors(model: ModelProto, filepath: str) -> ModelProto:
418 """Serializes data for all the tensors which have data location set to TensorProto.External.
420 Note: This function also strips basepath information from all tensors' external_data fields.
422 Arguments:
423 model (ModelProto): Model object which is the source of tensors to serialize.
424 filepath: System path to the directory which should be treated as base path for external data.
426 Returns:
427 ModelProto: The modified model object.
428 """
429 # Writing to external data happens in 2 passes:
430 # 1. Tensors with raw data which pass the necessary conditions (size threshold etc) are marked for serialization
431 # 2. The raw data in these tensors is serialized to a file
432 # Thus serialize only if tensor has raw data and it was marked for serialization
433 tensors_to_write = [
434 tensor
435 for tensor in _get_all_tensors(model)
436 if uses_external_data(tensor) and tensor.HasField("raw_data")
437 ]
438 # save_external_data() appends each tensor to the end of its external data
439 # file and validates that any pre-assigned offset lands there. Sort by that
440 # offset (per destination file) so tensors are written in the order their
441 # offsets imply, regardless of their order in the graph; tensors without a
442 # pre-assigned offset keep their relative order and are written last.
443 infos = {id(tensor): ExternalDataInfo(tensor) for tensor in tensors_to_write}
444 tensors_to_write.sort(
445 key=lambda tensor: (
446 infos[id(tensor)].location,
447 infos[id(tensor)].offset is None,
448 infos[id(tensor)].offset or 0,
449 )
450 )
451 for tensor in tensors_to_write:
452 save_external_data(tensor, filepath)
453 tensor.ClearField("raw_data")
455 return model