Test validate path

Author: yuslepukhin

include/onnxruntime/core/graph/graph.h

@@ -1463,10 +1463,11 @@ class Graph {  // NOLINT(clang-analyzer-optin.performance.Padding): preserve exi
 
   /// <summary>
   /// This function converts all the graph TensorProto initializers into OrtValues
-  /// and creates a in-memory external data reference for each OrtValue.
+  /// and creates a in-memory external data reference for each OrtValue. It validates external paths data references.
   /// </summary>
+  /// <param name="whitelisted_external_paths"></param>
   /// <returns></returns>
-  Status ConvertInitializersIntoOrtValues();
+  Status ConvertInitializersIntoOrtValues(gsl::span<const std::filesystem::path> whitelisted_external_paths);
 
   /**
    * @brief This function examines the specified initializers in the graph and converts them inline

onnxruntime/core/framework/tensorprotoutils.cc

@@ -401,23 +401,49 @@ Status ParseWhiteListedPaths(const PathString& paths_str,
 }
 
 Status ValidateExternalDataPath(const std::filesystem::path& base_dir,
-                                const std::filesystem::path& location) {
+                                const std::filesystem::path& location,
+                                gsl::span<const std::filesystem::path> whitelisted_external_folders) {
   // Reject absolute paths
   ORT_RETURN_IF(location.is_absolute(),
                 "Absolute paths not allowed for external data location");
-  if (!base_dir.empty()) {
-    // Resolve and verify the path stays within model directory
-    auto base_canonical = std::filesystem::weakly_canonical(base_dir);
-    // If the symlink exists, it resolves to the target path;
-    // so if the symlink is outside the directory it would be caught here.
-    auto resolved = std::filesystem::weakly_canonical(base_dir / location);
-    // Check that resolved path starts with base directory
+
+  auto validate_location_under_dir = [&location](const std::filesystem::path& dir) -> bool {
+    if (dir.empty()) {
+      return false;
+    }
+    auto base_canonical = std::filesystem::weakly_canonical(dir);
+    auto resolved = std::filesystem::weakly_canonical(dir / location);
     auto [base_end, resolved_it] = std::mismatch(
         base_canonical.begin(), base_canonical.end(),
         resolved.begin(), resolved.end());
-    ORT_RETURN_IF(base_end != base_canonical.end(),
-                  "External data path: ", location, " escapes model directory: ", base_dir);
+    return base_end == base_canonical.end();
+  };
+
+  if (!base_dir.empty()) {
+    if (validate_location_under_dir(base_dir)) {
+      return Status::OK();
+    }
   }
+
+  // base_dir validation failed or base_dir is empty, try whitelisted folders
+  if (!whitelisted_external_folders.empty()) {
+    for (const auto& folder : whitelisted_external_folders) {
+      if (validate_location_under_dir(folder)) {
+        return Status::OK();
+      }
+    }
+
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                           "External data path: ", location,
+                           " is not under any allowed directory");
+  }
+
+  // No whitelisted folders supplied
+  if (!base_dir.empty()) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                           "External data path: ", location, " escapes model directory: ", base_dir);
+  }
+
   return Status::OK();
 }
 

onnxruntime/core/framework/tensorprotoutils.h

@@ -539,14 +539,17 @@ Status ParseWhiteListedPaths(const PathString& paths_str,
 /// <summary>
 /// The functions will make sure the 'location' specified in the external data is under the 'base_dir'.
 /// If the `base_dir` is empty, the function only ensures that `location` is not an absolute path.
+/// If validation fails for base_dir, the function will check against whitelisted_external_folders.
 /// </summary>
 /// <param name="base_dir">model location directory</param>
 /// <param name="location">location is a string retrieved from TensorProto external data that is not
 ///                        an in-memory tag</param>
+/// <param name="whitelisted_external_folders">additional folders where external data is allowed</param>
 /// <returns>The function will fail if the resolved full path is not under the model directory
-///          or one of the subdirectories</returns>
+///          or one of the whitelisted folders</returns>
 Status ValidateExternalDataPath(const std::filesystem::path& base_dir,
-                                const std::filesystem::path& location);
+                                const std::filesystem::path& location,
+                                gsl::span<const std::filesystem::path> whitelisted_external_folders = {});
 
 #endif  // !defined(SHARED_PROVIDER)
 

onnxruntime/core/graph/graph.cc

@@ -3737,7 +3737,7 @@ Status Graph::Resolve(const ResolveOptions& options) {
   return ForThisAndAllSubgraphs(all_subgraphs, finalize_func);
 }
 
-Status Graph::ConvertInitializersIntoOrtValues() {
+Status Graph::ConvertInitializersIntoOrtValues(gsl::span<const std::filesystem::path> whitelisted_external_paths) {
   std::vector<Graph*> all_subgraphs;
   FindAllSubgraphs(all_subgraphs);
 
@@ -3771,7 +3771,7 @@ Status Graph::ConvertInitializersIntoOrtValues() {
           std::unique_ptr<onnxruntime::ExternalDataInfo> external_data_info;
           ORT_RETURN_IF_ERROR(onnxruntime::ExternalDataInfo::Create(tensor_proto.external_data(), external_data_info));
           const auto& location = external_data_info->GetRelPath();
-          auto st = utils::ValidateExternalDataPath(model_dir, location);
+          auto st = utils::ValidateExternalDataPath(model_dir, location, whitelisted_external_paths);
           if (!st.IsOK()) {
             return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL,
                                    "External data path validation failed for initializer: ", tensor_proto.name(),

onnxruntime/core/providers/shared_library/provider_api.h

@@ -453,11 +453,6 @@ inline bool HasExternalDataInMemory(const ONNX_NAMESPACE::TensorProto& ten_proto
   return g_host->Utils__HasExternalDataInMemory(ten_proto);
 }
 
-inline Status ValidateExternalDataPath(const std::filesystem::path& base_dir,
-                                       const std::filesystem::path& location) {
-  return g_host->Utils__ValidateExternalDataPath(base_dir, location);
-}
-
 }  // namespace utils
 
 namespace graph_utils {

onnxruntime/core/providers/shared_library/provider_interfaces.h

@@ -1004,9 +1004,6 @@ struct ProviderHost {
 
   virtual bool Utils__HasExternalDataInMemory(const ONNX_NAMESPACE::TensorProto& ten_proto) = 0;
 
-  virtual Status Utils__ValidateExternalDataPath(const std::filesystem::path& base_path,
-                                                 const std::filesystem::path& location) = 0;
-
   // Model
   virtual std::unique_ptr<Model> Model__construct(ONNX_NAMESPACE::ModelProto&& model_proto, const PathString& model_path,
                                                   const IOnnxRuntimeOpSchemaRegistryList* local_registries,

onnxruntime/core/session/inference_session.cc

@@ -1390,7 +1390,10 @@ common::Status InferenceSession::TransformGraph(onnxruntime::Graph& graph, bool
   //   auto tensor_proto_to_add = utils::TensorToTensorProto(ort_value.Get<Tensor>(), tensor_proto.name(),
   //                                                        use_tensor_buffer_true);
   //   ORT_RETURN_IF_ERROR(graph.ReplaceInitializedTensor(tensor_proto_to_add, ort_value));
-  ORT_RETURN_IF_ERROR_SESSIONID_(graph.ConvertInitializersIntoOrtValues());
+  InlinedVector<std::filesystem::path> whitelisted_external_data_folders;
+  ORT_RETURN_IF_ERROR_SESSIONID_(utils::ParseWhiteListedPaths(session_options_.whitelisted_data_folders,
+                                                              whitelisted_external_data_folders));
+  ORT_RETURN_IF_ERROR_SESSIONID_(graph.ConvertInitializersIntoOrtValues(whitelisted_external_data_folders));
 
   auto apply_transformer_once = [](const GraphTransformer& transformer, const logging::Logger& logger,
                                    Graph& graph, bool* is_graph_modified = nullptr) -> onnxruntime::common::Status {

onnxruntime/core/session/provider_bridge_ort.cc

@@ -1295,11 +1295,6 @@ struct ProviderHostImpl : ProviderHost {
     return onnxruntime::utils::HasExternalDataInMemory(ten_proto);
   }
 
-  Status Utils__ValidateExternalDataPath(const std::filesystem::path& base_path,
-                                         const std::filesystem::path& location) override {
-    return onnxruntime::utils::ValidateExternalDataPath(base_path, location);
-  }
-
   // Model (wrapped)
   std::unique_ptr<Model> Model__construct(ONNX_NAMESPACE::ModelProto&& model_proto, const PathString& model_path,
                                           const IOnnxRuntimeOpSchemaRegistryList* local_registries,

onnxruntime/test/framework/tensorutils_test.cc

@@ -511,18 +511,22 @@ class PathValidationTest : public ::testing::Test {
     // Create a temporary directory for the tests.
     base_dir_ = std::filesystem::temp_directory_path() / "PathValidationTest";
     outside_dir_ = std::filesystem::temp_directory_path() / "outside";
+    whitelisted_dir_ = std::filesystem::temp_directory_path() / "whitelisted";
     std::filesystem::create_directories(base_dir_);
     std::filesystem::create_directories(outside_dir_);
+    std::filesystem::create_directories(whitelisted_dir_);
   }
 
   void TearDown() override {
     // Clean up the temporary directory.
     std::filesystem::remove_all(base_dir_);
     std::filesystem::remove_all(outside_dir_);
+    std::filesystem::remove_all(whitelisted_dir_);
   }
 
   std::filesystem::path base_dir_;
   std::filesystem::path outside_dir_;
+  std::filesystem::path whitelisted_dir_;
 };
 
 // Test cases for ValidateExternalDataPath.
@@ -586,6 +590,85 @@ TEST_F(PathValidationTest, ValidateExternalDataPathWithSymlinkOutside) {
   ASSERT_FALSE(utils::ValidateExternalDataPath(base_dir_, "outside_link.bin").IsOK());
 }
 
+TEST_F(PathValidationTest, ValidateExternalDataPathWithWhitelistedFolder) {
+  // Path is valid under the whitelisted folder directly.
+  std::vector<std::filesystem::path> whitelist = {outside_dir_};
+  ASSERT_STATUS_OK(utils::ValidateExternalDataPath(outside_dir_, "data.bin", whitelist));
+}
+
+TEST_F(PathValidationTest, ValidateExternalDataPathEscapesBaseButMatchesWhitelist) {
+  std::vector<std::filesystem::path> whitelist = {whitelisted_dir_};
+  // "data.bin" is valid under base_dir_, no need for whitelist
+  ASSERT_STATUS_OK(utils::ValidateExternalDataPath(base_dir_, "data.bin", whitelist));
+
+  // Create a subdirectory of whitelisted_dir_ and use that as the whitelisted folder.
+  auto whitelisted_sub = whitelisted_dir_ / "sub";
+  std::filesystem::create_directories(whitelisted_sub);
+  std::vector<std::filesystem::path> whitelist2 = {whitelisted_sub};
+
+  // "../data.bin" escapes base_dir_ and also escapes whitelisted_sub
+  ASSERT_FALSE(utils::ValidateExternalDataPath(base_dir_, "../data.bin").IsOK());
+  ASSERT_FALSE(utils::ValidateExternalDataPath(base_dir_, "../data.bin", whitelist2).IsOK());
+}
+
+TEST_F(PathValidationTest, ValidateExternalDataPathWhitelistSavesEscapingPath) {
+  // Location "../outside/data.bin" escapes base_dir_ but resolves under outside_dir_.
+  auto relative_to_outside = std::filesystem::path("..") / "outside" / "data.bin";
+  std::vector<std::filesystem::path> whitelist = {outside_dir_};
+
+  // Without whitelist, it should fail.
+  ASSERT_FALSE(utils::ValidateExternalDataPath(base_dir_, relative_to_outside).IsOK());
+
+  // With whitelist containing outside_dir_, it should succeed because
+  // outside_dir_ / "../outside/data.bin" resolves under outside_dir_.
+  ASSERT_STATUS_OK(utils::ValidateExternalDataPath(base_dir_, relative_to_outside, whitelist));
+}
+
+TEST_F(PathValidationTest, ValidateExternalDataPathWhitelistDoesNotMatchEither) {
+  // Location escapes both base_dir and all whitelisted folders.
+  auto unrelated_dir = std::filesystem::temp_directory_path() / "unrelated_PathValidationTest";
+  std::filesystem::create_directories(unrelated_dir);
+  auto cleanup = [&]() { std::filesystem::remove_all(unrelated_dir); };
+
+  std::vector<std::filesystem::path> whitelist = {whitelisted_dir_};
+  auto escaping_location = std::filesystem::path("..") / "unrelated_PathValidationTest" / "data.bin";
+  ASSERT_FALSE(utils::ValidateExternalDataPath(base_dir_, escaping_location, whitelist).IsOK());
+
+  cleanup();
+}
+
+TEST_F(PathValidationTest, ValidateExternalDataPathEmptyWhitelist) {
+  // Empty whitelist should behave the same as no whitelist.
+  std::vector<std::filesystem::path> empty_whitelist;
+  ASSERT_STATUS_OK(utils::ValidateExternalDataPath(base_dir_, "data.bin", empty_whitelist));
+  ASSERT_FALSE(utils::ValidateExternalDataPath(base_dir_, "../data.bin", empty_whitelist).IsOK());
+}
+
+TEST_F(PathValidationTest, ValidateExternalDataPathMultipleWhitelistedFolders) {
+  // First whitelisted folder doesn't match, second one does.
+  auto another_dir = std::filesystem::temp_directory_path() / "another_PathValidationTest";
+  std::filesystem::create_directories(another_dir);
+  auto cleanup = [&]() { std::filesystem::remove_all(another_dir); };
+
+  auto relative_to_outside = std::filesystem::path("..") / "outside" / "data.bin";
+  std::vector<std::filesystem::path> whitelist = {another_dir, outside_dir_};
+
+  // Escapes base_dir_ but outside_dir_ (second whitelist entry) should match.
+  ASSERT_STATUS_OK(utils::ValidateExternalDataPath(base_dir_, relative_to_outside, whitelist));
+
+  cleanup();
+}
+
+TEST_F(PathValidationTest, ValidateExternalDataPathAbsoluteLocationRejectsEvenWithWhitelist) {
+  // Absolute paths are always rejected, regardless of whitelist.
+  std::vector<std::filesystem::path> whitelist = {outside_dir_};
+#ifdef _WIN32
+  ASSERT_FALSE(utils::ValidateExternalDataPath(base_dir_, "C:\\data.bin", whitelist).IsOK());
+#else
+  ASSERT_FALSE(utils::ValidateExternalDataPath(base_dir_, "/data.bin", whitelist).IsOK());
+#endif
+}
+
 // Test fixture for ParseWhiteListedPaths tests.
 class ParseWhiteListedPathsTest : public ::testing::Test {
  protected:
@@ -773,5 +856,6 @@ TEST_F(ParseWhiteListedPathsTest, OutParamUnchangedOnError) {
   ASSERT_EQ(paths.size(), 1u);
   EXPECT_EQ(paths[0], sub_dir_a_);
 }
+
 }  // namespace test
 }  // namespace onnxruntime

onnxruntime/test/ir/graph_test.cc

@@ -2817,7 +2817,7 @@ TEST_F(GraphTest, ShapeInferenceAfterInitializerExternalization) {
       << "We no longer externalize data in the Graph constructor.";
 
   // Now externalize explicitly to trigger the bug scenario
-  ASSERT_STATUS_OK(graph.ConvertInitializersIntoOrtValues());
+  ASSERT_STATUS_OK(graph.ConvertInitializersIntoOrtValues({}));
   ASSERT_TRUE(graph.GetInitializedTensor("split_sizes", initializer_after));
   ASSERT_NE(initializer_after, nullptr);
   ASSERT_TRUE(utils::HasExternalDataInMemory(*initializer_after))