From 6189e5d82f3a90b62661fe082c5c9d8bd81ea85f Mon Sep 17 00:00:00 2001
From: Boris Zbarsky <bzbarsky@apple.com>
Date: Wed, 26 Oct 2022 15:26:58 -0400
Subject: [PATCH] Add a way to avoid heap pools crashing on exit if items are
 still live. (#23359)

Currently a pool's destructor crashes if there are still live objects in the
pool.  For heap pools, add a way to not crash if the pool's destructor is
running under exit(), for consumers that don't care about exit-time-only leaks.
---
 .../Framework/CHIP/MTRControllerFactory.mm    |  6 +++++
 src/lib/support/Pool.cpp                      | 27 +++++++++++++++++++
 src/lib/support/Pool.h                        | 26 +++++++++++++++---
 3 files changed, 56 insertions(+), 3 deletions(-)

diff --git a/src/darwin/Framework/CHIP/MTRControllerFactory.mm b/src/darwin/Framework/CHIP/MTRControllerFactory.mm
index 3d77b893229d7b..88f0cc70401856 100644
--- a/src/darwin/Framework/CHIP/MTRControllerFactory.mm
+++ b/src/darwin/Framework/CHIP/MTRControllerFactory.mm
@@ -39,6 +39,7 @@
 #include <credentials/attestation_verifier/DefaultDeviceAttestationVerifier.h>
 #include <credentials/attestation_verifier/DeviceAttestationVerifier.h>
 #include <crypto/PersistentStorageOperationalKeystore.h>
+#include <lib/support/Pool.h>
 #include <lib/support/TestPersistentStorageDelegate.h>
 #include <platform/PlatformManager.h>
 
@@ -309,6 +310,11 @@ - (BOOL)startup:(MTRControllerFactoryParams *)startupParams
             return;
         }
 
+        // This needs to happen after DeviceControllerFactory::Init,
+        // because that creates (lazily, by calling functions with
+        // static variables in them) some static-lifetime objects.
+        chip::HeapObjectPoolExitHandling::IgnoreLeaksOnExit();
+
         // Make sure we don't leave a system state running while we have no
         // controllers started.  This is working around the fact that a system
         // state is brought up live on factory init, and not when it comes time
diff --git a/src/lib/support/Pool.cpp b/src/lib/support/Pool.cpp
index e7f5f576a59b59..c32b46be45e819 100644
--- a/src/lib/support/Pool.cpp
+++ b/src/lib/support/Pool.cpp
@@ -22,6 +22,33 @@
 
 namespace chip {
 
+#if CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+
+bool HeapObjectPoolExitHandling::sIgnoringLeaksOnExit   = false;
+bool HeapObjectPoolExitHandling::sExitHandlerRegistered = false;
+
+void HeapObjectPoolExitHandling::IgnoreLeaksOnExit()
+{
+    if (sExitHandlerRegistered)
+    {
+        return;
+    }
+
+    int ret = atexit(ExitHandler);
+    if (ret != 0)
+    {
+        ChipLogError(Controller, "IgnoreLeaksOnExit: atexit failed: %d\n", ret);
+    }
+    sExitHandlerRegistered = true;
+}
+
+void HeapObjectPoolExitHandling::ExitHandler()
+{
+    sIgnoringLeaksOnExit = true;
+}
+
+#endif // CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+
 namespace internal {
 
 StaticAllocatorBitmap::StaticAllocatorBitmap(void * storage, std::atomic<tBitChunkType> * usage, size_t capacity,
diff --git a/src/lib/support/Pool.h b/src/lib/support/Pool.h
index 54914c7d91d141..72d1a1842b0e36 100644
--- a/src/lib/support/Pool.h
+++ b/src/lib/support/Pool.h
@@ -291,13 +291,30 @@ class BitMapObjectPool : public internal::StaticAllocatorBitmap, public internal
 
 #if CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
 
+class HeapObjectPoolExitHandling
+{
+public:
+    // If IgnoreLeaksOnExit is called, some time after all static initializers have
+    // run, HeapObjectPool will not assert that everything in it has been
+    // released if its destructor runs under exit() (i.e. when the application
+    // is quitting anyway).
+    static void IgnoreLeaksOnExit();
+
+protected:
+    static bool sIgnoringLeaksOnExit;
+
+private:
+    static void ExitHandler();
+    static bool sExitHandlerRegistered;
+};
+
 /**
  * A class template used for allocating objects from the heap.
  *
  *  @tparam     T   type to be allocated.
  */
 template <class T>
-class HeapObjectPool : public internal::Statistics, public internal::PoolCommon<T>
+class HeapObjectPool : public internal::Statistics, public internal::PoolCommon<T>, public HeapObjectPoolExitHandling
 {
 public:
     HeapObjectPool() {}
@@ -314,8 +331,11 @@ class HeapObjectPool : public internal::Statistics, public internal::PoolCommon<
         // Free all remaining objects so that ASAN can catch specific use-after-free cases.
         ReleaseAll();
 #else  // __SANITIZE_ADDRESS__
-       // Verify that no live objects remain, to prevent potential use-after-free.
-        VerifyOrDie(Allocated() == 0);
+        if (!sIgnoringLeaksOnExit)
+        {
+            // Verify that no live objects remain, to prevent potential use-after-free.
+            VerifyOrDie(Allocated() == 0);
+        }
 #endif // __SANITIZE_ADDRESS__
     }