Fix throwing operator new in ExecuTorch FileDataLoader (#20400)

Summary:

`et_aligned_alloc` in file_data_loader.cpp used the throwing form of
`::operator new(size, alignment)`, which throws `std::bad_alloc` on
allocation failure. The ExecuTorch runtime is built exception-free and
uses a `Result<>`/`Error` convention: `FileDataLoader::load` already
guards with `if (aligned_buffer == nullptr) return Error::MemoryAllocationFailed;`,
but that null-check was dead code because the throwing `operator new`
never returns null. On memory-constrained iOS devices, loading a large
model segment (e.g. SceneX V7) could throw, unwind with no landing pad,
and abort the process (SIGABRT) -- see the linked task.

Switch to the nothrow form `::operator new(size, alignment, std::nothrow)`
so the existing null-check fires and `load` returns
`Error::MemoryAllocationFailed`, which propagates cleanly through
`Program::LoadSegment` -> `Method::init` -> `-[ExecuTorchModule loadMethod:error:]`
and surfaces as a catchable NSError instead of aborting.

Also add a unit test that forces the segment allocation to fail (by
replacing the global nothrow aligned operator new for the test binary)
and asserts `load` returns `Error::MemoryAllocationFailed`.

Reviewed By: Gasoonjia

Differential Revision: D109072812

Author: metascroy

extension/data_loader/file_data_loader.cpp

@@ -13,6 +13,7 @@
 #include <cstddef>
 #include <cstring>
 #include <limits>
+#include <new>
 
 #include <executorch/runtime/platform/compat_unistd.h>
 #include <fcntl.h>
@@ -44,7 +45,12 @@ namespace extension {
 
 namespace {
 inline void* et_aligned_alloc(size_t size, std::align_val_t alignment) {
-  return ::operator new(size, alignment);
+  // Use the nothrow form so allocation failure returns nullptr instead of
+  // throwing std::bad_alloc. ExecuTorch is built exception-free and callers
+  // (e.g. FileDataLoader::load) check for nullptr and return
+  // Error::MemoryAllocationFailed; a throw here would unwind with no landing
+  // pad and abort the process.
+  return ::operator new(size, alignment, std::nothrow);
 }
 
 inline void et_aligned_free(void* ptr, std::align_val_t alignment) {

extension/data_loader/test/file_data_loader_test.cpp

@@ -8,12 +8,15 @@
 
 #include <executorch/extension/data_loader/file_data_loader.h>
 
+#include <atomic>
 #include <cstring>
+#include <new>
 
 #include <gtest/gtest.h>
 
 #include <executorch/extension/testing_util/temp_file.h>
 #include <executorch/runtime/core/result.h>
+#include <executorch/runtime/platform/compiler.h>
 #include <executorch/runtime/platform/runtime.h>
 #include <executorch/test/utils/alignment.h>
 
@@ -25,6 +28,57 @@ using executorch::runtime::Error;
 using executorch::runtime::FreeableBuffer;
 using executorch::runtime::Result;
 
+namespace {
+// When set, the replacement nothrow aligned operator new below returns nullptr,
+// simulating an allocation failure without needing a real OOM.
+std::atomic<bool> g_fail_aligned_nothrow_alloc{false};
+
+// RAII guard to ensure flag is reset even if test asserts early.
+struct FailAllocGuard {
+  FailAllocGuard() {
+    g_fail_aligned_nothrow_alloc.store(true, std::memory_order_relaxed);
+  }
+  ~FailAllocGuard() {
+    g_fail_aligned_nothrow_alloc.store(false, std::memory_order_relaxed);
+  }
+};
+} // namespace
+
+// Detect ASAN to avoid multiple definition link error and to skip test when
+// ASAN runtime provides its own strong operator new.
+#if defined(__SANITIZE_ADDRESS__) || (defined(__has_feature) && __has_feature(address_sanitizer))
+#define ET_TEST_ASAN_ENABLED 1
+#else
+#define ET_TEST_ASAN_ENABLED 0
+#endif
+
+#if !ET_TEST_ASAN_ENABLED
+// Replaces the global nothrow aligned allocation function for this test binary
+// so FileDataLoader's segment allocation can be made to fail on demand. When
+// the toggle is off it forwards to the real aligned allocator. We call the
+// throwing aligned new inside a try/catch and convert exceptions to nullptr
+// to emulate nothrow semantics without recursing into this same nothrow
+// overload (calling ::operator new(size, alignment, std::nothrow) here would
+// infinite-loop). Memory allocated here is released through the default
+// operator delete, which is not replaced.
+// Marked weak to avoid conflict with ASAN runtime which provides its own strong
+// definition. Under ASAN the test is skipped.
+ET_WEAK
+void* operator new(
+    std::size_t size,
+    std::align_val_t alignment,
+    const std::nothrow_t& /* tag */) noexcept {
+  if (g_fail_aligned_nothrow_alloc.load(std::memory_order_relaxed)) {
+    return nullptr;
+  }
+  try {
+    return ::operator new(size, alignment);
+  } catch (...) {
+    return nullptr;
+  }
+}
+#endif // !ET_TEST_ASAN_ENABLED
+
 class FileDataLoaderTest : public ::testing::TestWithParam<size_t> {
  protected:
   void SetUp() override {
@@ -147,6 +201,46 @@ TEST_P(FileDataLoaderTest, OutOfBoundsLoadFails) {
   }
 }
 
+#if !ET_TEST_ASAN_ENABLED
+TEST_P(FileDataLoaderTest, AllocationFailureDuringLoadReturnsError) {
+  // Create a temp file; contents don't matter.
+  uint8_t data[256] = {};
+  TempFile tf(data, sizeof(data));
+
+  Result<FileDataLoader> fdl =
+      FileDataLoader::from(tf.path().c_str(), alignment());
+  ASSERT_EQ(fdl.error(), Error::Ok);
+
+  // Force the segment allocation inside load() to fail. The loader must surface
+  // Error::MemoryAllocationFailed rather than letting std::bad_alloc escape,
+  // which would abort the process in the exception-free runtime.
+  FailAllocGuard fail_guard;
+  Result<FreeableBuffer> fb = fdl->load(
+      /*offset=*/0,
+      /*size=*/sizeof(data),
+      DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program));
+
+  EXPECT_EQ(fb.error(), Error::MemoryAllocationFailed);
+}
+#endif // ET_TEST_ASAN_ENABLED
+
+#if !ET_TEST_ASAN_ENABLED
+TEST_P(FileDataLoaderTest, AllocationFailureDuringFromReturnsError) {
+  // Create a temp file; contents don't matter.
+  uint8_t data[256] = {};
+  TempFile tf(data, sizeof(data));
+
+  // Force the filename allocation inside from() to fail. FileDataLoader::from
+  // copies the filename using et_aligned_alloc and must return
+  // Error::MemoryAllocationFailed on nullptr rather than throwing.
+  FailAllocGuard fail_guard;
+  Result<FileDataLoader> fdl =
+      FileDataLoader::from(tf.path().c_str(), alignment());
+
+  EXPECT_EQ(fdl.error(), Error::MemoryAllocationFailed);
+}
+#endif // ET_TEST_ASAN_ENABLED
+
 TEST_P(FileDataLoaderTest, FromMissingFileFails) {
   // Wrapping a file that doesn't exist should fail.
   Result<FileDataLoader> fdl = FileDataLoader::from(