CVE-2024-27318

@@ -13,7 +13,6 @@
 #include <vector>
 
 #include "onnx/common/file_utils.h"
-#include "onnx/common/path.h"
 #include "onnx/defs/schema.h"
 #include "onnx/defs/tensor_proto_util.h"
 #include "onnx/proto_utils.h"
@@ -128,85 +127,7 @@ void check_tensor(const TensorProto& tensor, const CheckerContext& ctx) {
     for (const StringStringEntryProto& entry : tensor.external_data()) {
       if (entry.has_key() && entry.has_value() && entry.key() == "location") {
         has_location = true;
-#ifdef _WIN32
-        auto file_path = std::filesystem::path(utf8str_to_wstring(entry.value()));
-        if (file_path.is_absolute()) {
-          fail_check(
-              "Location of external TensorProto ( tensor name: ",
-              tensor.name(),
-              ") should be a relative path, but it is an absolute path: ",
-              entry.value());
-        }
-        auto relative_path = file_path.lexically_normal().make_preferred().wstring();
-        // Check that normalized relative path contains ".." on Windows.
-        if (relative_path.find(L"..", 0) != std::string::npos) {
-          fail_check(
-              "Data of TensorProto ( tensor name: ",
-              tensor.name(),
-              ") should be file inside the ",
-              ctx.get_model_dir(),
-              ", but the '",
-              entry.value(),
-              "' points outside the directory");
-        }
-        std::wstring data_path = path_join(utf8str_to_wstring(ctx.get_model_dir()), relative_path);
-        struct _stat64 buff;
-        if (data_path.empty() || (data_path[0] != '#' && _wstat64(data_path.c_str(), &buff) != 0)) {
-          fail_check(
-              "Data of TensorProto ( tensor name: ",
-              tensor.name(),
-              ") should be stored in ",
-              entry.value(),
-              ", but it doesn't exist or is not accessible.");
-        }
-#else // POSIX
-        if (entry.value().empty()) {
-          fail_check("Location of external TensorProto ( tensor name: ", tensor.name(), ") should not be empty.");
-        } else if (entry.value()[0] == '/') {
-          fail_check(
-              "Location of external TensorProto ( tensor name: ",
-              tensor.name(),
-              ") should be a relative path, but it is an absolute path: ",
-              entry.value());
-        }
-        std::string relative_path = clean_relative_path(entry.value());
-        // Check that normalized relative path contains ".." on POSIX
-        if (relative_path.find("..", 0) != std::string::npos) {
-          fail_check(
-              "Data of TensorProto ( tensor name: ",
-              tensor.name(),
-              ") should be file inside the ",
-              ctx.get_model_dir(),
-              ", but the '",
-              entry.value(),
-              "' points outside the directory");
-        }
-        std::string data_path = path_join(ctx.get_model_dir(), relative_path);
-        // use stat64 to check whether the file exists
-#if defined(__APPLE__) || defined(__wasm__) || !defined(__GLIBC__)
-        struct stat buffer; // APPLE, wasm and non-glic stdlibs do not have stat64
-        if (data_path.empty() || (data_path[0] != '#' && stat((data_path).c_str(), &buffer) != 0)) {
-#else
-        struct stat64 buffer; // All POSIX under glibc except APPLE and wasm have stat64
-        if (data_path.empty() || (data_path[0] != '#' && stat64((data_path).c_str(), &buffer) != 0)) {
-#endif
-          fail_check(
-              "Data of TensorProto ( tensor name: ",
-              tensor.name(),
-              ") should be stored in ",
-              data_path,
-              ", but it doesn't exist or is not accessible.");
-        }
-        // Do not allow symlinks or directories.
-        if (data_path.empty() || (data_path[0] != '#' && !S_ISREG(buffer.st_mode))) {
-          fail_check(
-              "Data of TensorProto ( tensor name: ",
-              tensor.name(),
-              ") should be stored in ",
-              data_path,
-              ", but it is not regular file.");
-        }
-#endif
+        resolve_external_data_location(ctx.get_model_dir(), entry.value(), tensor.name());
       }
     }
     if (!has_location) {
@@ -1049,6 +970,93 @@ void check_model(const ModelProto& model, bool full_check, bool skip_opset_compa
   }
 }
 
+std::string resolve_external_data_location(
+    const std::string& base_dir,
+    const std::string& location,
+    const std::string& tensor_name) {
+#ifdef _WIN32
+  auto file_path = std::filesystem::path(utf8str_to_wstring(location));
+  if (file_path.is_absolute()) {
+    fail_check(
+        "Location of external TensorProto ( tensor name: ",
+        tensor_name,
+        ") should be a relative path, but it is an absolute path: ",
+        location);
+  }
+  auto relative_path = file_path.lexically_normal().make_preferred().wstring();
+  // Check that normalized relative path contains ".." on Windows.
+  if (relative_path.find(L"..", 0) != std::string::npos) {
+    fail_check(
+        "Data of TensorProto ( tensor name: ",
+        tensor_name,
+        ") should be file inside the ",
+        base_dir,
+        ", but the '",
+        location,
+        "' points outside the directory");
+  }
+  std::wstring data_path = path_join(utf8str_to_wstring(base_dir), relative_path);
+  struct _stat64 buff;
+  if (data_path.empty() || (data_path[0] != '#' && _wstat64(data_path.c_str(), &buff) != 0)) {
+    fail_check(
+        "Data of TensorProto ( tensor name: ",
+        tensor_name,
+        ") should be stored in ",
+        location,
+        ", but it doesn't exist or is not accessible.");
+  }
+  return wstring_to_utf8str(data_path);
+#else // POSIX
+  if (location.empty()) {
+    fail_check("Location of external TensorProto ( tensor name: ", tensor_name, ") should not be empty.");
+  } else if (location[0] == '/') {
+    fail_check(
+        "Location of external TensorProto ( tensor name: ",
+        tensor_name,
+        ") should be a relative path, but it is an absolute path: ",
+        location);
+  }
+  std::string relative_path = clean_relative_path(location);
+  // Check that normalized relative path contains ".." on POSIX
+  if (relative_path.find("..", 0) != std::string::npos) {
+    fail_check(
+        "Data of TensorProto ( tensor name: ",
+        tensor_name,
+        ") should be file inside the ",
+        base_dir,
+        ", but the '",
+        location,
+        "' points outside the directory");
+  }
+  std::string data_path = path_join(base_dir, relative_path);
+  // use stat64 to check whether the file exists
+#if defined(__APPLE__) || defined(__wasm__) || !defined(__GLIBC__)
+  struct stat buffer; // APPLE, wasm and non-glic stdlibs do not have stat64
+  if (data_path.empty() || (data_path[0] != '#' && stat((data_path).c_str(), &buffer) != 0)) {
+#else
+  struct stat64 buffer; // All POSIX under glibc except APPLE and wasm have stat64
+  if (data_path.empty() || (data_path[0] != '#' && stat64((data_path).c_str(), &buffer) != 0)) {
+#endif
+    fail_check(
+        "Data of TensorProto ( tensor name: ",
+        tensor_name,
+        ") should be stored in ",
+        data_path,
+        ", but it doesn't exist or is not accessible.");
+  }
+  // Do not allow symlinks or directories.
+  if (data_path.empty() || (data_path[0] != '#' && !S_ISREG(buffer.st_mode))) {
+    fail_check(
+        "Data of TensorProto ( tensor name: ",
+        tensor_name,
+        ") should be stored in ",
+        data_path,
+        ", but it is not regular file.");
+  }
+  return data_path;
+#endif
+}
+
 std::set<std::string> experimental_ops = {
     "ATen",
     "Affine",

@@ -160,7 +160,10 @@ void check_model_local_functions(
 
 void check_model(const ModelProto& model, bool full_check = false, bool skip_opset_compatibility_check = false);
 void check_model(const std::string& model_path, bool full_check = false, bool skip_opset_compatibility_check = false);
-
+std::string resolve_external_data_location(
+    const std::string& base_dir,
+    const std::string& location,
+    const std::string& tensor_name);
 bool check_is_experimental_op(const NodeProto& node);
 
 } // namespace checker

@@ -36,11 +36,22 @@ inline std::wstring utf8str_to_wstring(const std::string& utf8str) {
   if (utf8str.size() > INT_MAX) {
     fail_check("utf8str_to_wstring: string is too long for converting to wstring.");
   }
-  int size_required = MultiByteToWideChar(CP_UTF8, 0, utf8str.c_str(), (int)utf8str.size(), NULL, 0);
+  int size_required = MultiByteToWideChar(CP_UTF8, 0, utf8str.c_str(), static_cast<int>(utf8str.size()), NULL, 0);
   std::wstring ws_str(size_required, 0);
-  MultiByteToWideChar(CP_UTF8, 0, utf8str.c_str(), (int)utf8str.size(), &ws_str[0], size_required);
+  MultiByteToWideChar(CP_UTF8, 0, utf8str.c_str(), static_cast<int>(utf8str.size()), &ws_str[0], size_required);
   return ws_str;
 }
+inline std::string wstring_to_utf8str(const std::wstring& ws_str) {
+  if (ws_str.size() > INT_MAX) {
+    fail_check("wstring_to_utf8str: string is too long for converting to UTF-8.");
+  }
+  int size_required =
+      WideCharToMultiByte(CP_UTF8, 0, ws_str.c_str(), static_cast<int>(ws_str.size()), NULL, 0, NULL, NULL);
+  std::string utf8str(size_required, 0);
+  WideCharToMultiByte(
+      CP_UTF8, 0, ws_str.c_str(), static_cast<int>(ws_str.size()), &utf8str[0], size_required, NULL, NULL);
+  return utf8str;
+}
 
 #else
 std::string path_join(const std::string& origin, const std::string& append);

@@ -535,6 +535,8 @@ PYBIND11_MODULE(onnx_cpp2py_export, onnx_cpp2py_export) {
       "full_check"_a = false,
       "skip_opset_compatibility_check"_a = false);
 
+  checker.def("_resolve_external_data_location", &checker::resolve_external_data_location);
+
   // Submodule `version_converter`
   auto version_converter = onnx_cpp2py_export.def_submodule("version_converter");
   version_converter.doc() = "VersionConverter submodule";

@@ -8,6 +8,7 @@
 from itertools import chain
 from typing import Callable, Iterable, Optional
 
+import onnx.onnx_cpp2py_export.checker as c_checker
 from onnx.onnx_pb import AttributeProto, GraphProto, ModelProto, TensorProto
 
 
@@ -38,9 +39,9 @@ def load_external_data_for_tensor(tensor: TensorProto, base_dir: str) -> None:
         base_dir: directory that contains the external data.
     """
     info = ExternalDataInfo(tensor)
-    file_location = _sanitize_path(info.location)
-    external_data_file_path = os.path.join(base_dir, file_location)
-
+    external_data_file_path = c_checker._resolve_external_data_location(  # type: ignore[attr-defined]
+        base_dir, info.location, tensor.name
+    )
     with open(external_data_file_path, "rb") as data_file:
         if info.offset:
             data_file.seek(info.offset)
@@ -254,14 +255,6 @@ def _get_attribute_tensors(onnx_model_proto: ModelProto) -> Iterable[TensorProto
     yield from _get_attribute_tensors_from_graph(onnx_model_proto.graph)
 
 
-def _sanitize_path(path: str) -> str:
-    """Remove path components which would allow traversing up a directory tree from a base path.
-
-    Note: This method is currently very basic and should be expanded.
-    """
-    return path.lstrip("/.")
-
-
 def _is_valid_filename(filename: str) -> bool:
     """Utility to check whether the provided filename is valid."""
     exp = re.compile('^[^<>:;,?"*|/]+$')

@@ -15,6 +15,7 @@
 import onnx
 import onnx.external_data_helper as ext_data
 import onnx.helper
+import onnx.onnx_cpp2py_export.checker as c_checker
 
 
 def _set_external_data(
@@ -288,8 +289,9 @@ def _load_large_initializers(self, file_path):
                 continue
 
             info = ext_data.ExternalDataInfo(tensor)
-            file_location = ext_data._sanitize_path(info.location)
-            external_data_file_path = os.path.join(base_dir, file_location)
+            external_data_file_path = c_checker._resolve_external_data_location(  # type: ignore[attr-defined]
+                base_dir, info.location, tensor.name
+            )
             key = f"#t{i}"
             _set_external_data(tensor, location=key)
 

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
+import itertools
 import os
 import pathlib
 import tempfile
@@ -204,6 +205,52 @@ def test_save_external_single_file_data(self) -> None:
         attribute_tensor = new_model.graph.node[0].attribute[0].t
         np.testing.assert_allclose(to_array(attribute_tensor), self.attribute_value)
 
+    @parameterized.parameterized.expand(itertools.product((True, False), (True, False)))
+    def test_save_external_invalid_single_file_data_and_check(
+        self, use_absolute_path: bool, use_model_path: bool
+    ) -> None:
+        model = onnx.load_model(self.model_filename, self.serialization_format)
+
+        model_dir = os.path.join(self.temp_dir, "save_copy")
+        os.mkdir(model_dir)
+
+        traversal_external_data_dir = os.path.join(
+            self.temp_dir, "invlid_external_data"
+        )
+        os.mkdir(traversal_external_data_dir)
+
+        if use_absolute_path:
+            traversal_external_data_location = os.path.join(
+                traversal_external_data_dir, "tensors.bin"
+            )
+        else:
+            traversal_external_data_location = "../invlid_external_data/tensors.bin"
+
+        external_data_dir = os.path.join(self.temp_dir, "external_data")
+        os.mkdir(external_data_dir)
+        new_model_filepath = os.path.join(model_dir, "model.onnx")
+
+        def convert_model_to_external_data_no_check(model: ModelProto, location: str):
+            for tensor in model.graph.initializer:
+                if tensor.HasField("raw_data"):
+                    set_external_data(tensor, location)
+
+        convert_model_to_external_data_no_check(
+            model,
+            location=traversal_external_data_location,
+        )
+
+        onnx.save_model(model, new_model_filepath, self.serialization_format)
+        if use_model_path:
+            with self.assertRaises(onnx.checker.ValidationError):
+                _ = onnx.load_model(new_model_filepath, self.serialization_format)
+        else:
+            onnx_model = onnx.load_model(
+                new_model_filepath, self.serialization_format, load_external_data=False
+            )
+            with self.assertRaises(onnx.checker.ValidationError):
+                load_external_data_for_model(onnx_model, external_data_dir)
+
 
 @parameterized.parameterized_class(
     [