fix: resolve memory management issues in JNI bindings, JSI lambdas, and delegate wrappers

CAMOBAP · CAMOBAP · commit f4a22914e9d3 · 2026-03-15T22:41:56.000+01:00
diff --git a/packages/react-native-sandbox/android/src/main/java/io/callstack/rnsandbox/SandboxReactNativeDelegate.kt b/packages/react-native-sandbox/android/src/main/java/io/callstack/rnsandbox/SandboxReactNativeDelegate.kt
@@ -59,6 +59,7 @@ class SandboxReactNativeDelegate(
     private var jsiStateHandle: Long = 0
     private var sandboxReactContext: ReactContext? = null
     private var ownsReactHost = false
+    private var instanceEventListener: ReactInstanceEventListener? = null
 
     @OptIn(UnstableReactNativeAPI::class)
     fun loadReactNativeView(
@@ -86,8 +87,7 @@ class SandboxReactNativeDelegate(
                 ownsReactHost = false
                 Log.d(TAG, "Reusing shared ReactHost for origin '$origin' (refCount=${shared.refCount})")
             } else {
-                val sandboxId = System.identityHashCode(this).toString(16)
-                sandboxContext = SandboxContextWrapper(context, sandboxId)
+                sandboxContext = SandboxContextWrapper(context, origin)
 
                 val packages: List<ReactPackage> =
                     listOf(
@@ -132,7 +132,7 @@ class SandboxReactNativeDelegate(
 
             reactHost = host
 
-            host.addReactInstanceEventListener(
+            val listener =
                 object : ReactInstanceEventListener {
                     override fun onReactContextInitialized(reactContext: ReactContext) {
                         sandboxReactContext = reactContext
@@ -142,8 +142,9 @@ class SandboxReactNativeDelegate(
                             }
                         }
                     }
-                },
-            )
+                }
+            instanceEventListener = listener
+            host.addReactInstanceEventListener(listener)
 
             val surface = host.createSurface(sandboxContext, componentName, initialProperties)
             reactSurface = surface
@@ -304,6 +305,10 @@ class SandboxReactNativeDelegate(
         reactSurface = null
 
         val host = reactHost
+        instanceEventListener?.let { listener ->
+            host?.removeReactInstanceEventListener(listener)
+        }
+        instanceEventListener = null
         if (host != null) {
             if (origin.isNotEmpty()) {
                 val shared = sharedHosts[origin]
diff --git a/packages/react-native-sandbox/android/src/main/jni/SandboxJSIInstaller.cpp b/packages/react-native-sandbox/android/src/main/jni/SandboxJSIInstaller.cpp
@@ -33,17 +33,44 @@ struct SandboxJSIState {
 static std::mutex gRegistryMutex;
 static std::unordered_map<jlong, std::shared_ptr<SandboxJSIState>> gStates;
 
+static JNIEnv* getJNIEnv() {
+  JNIEnv* env = nullptr;
+  if (gJavaVM) {
+    gJavaVM->GetEnv(reinterpret_cast<void**>(&env), JNI_VERSION_1_6);
+    if (!env) {
+      gJavaVM->AttachCurrentThread(&env, nullptr);
+    }
+  }
+  return env;
+}
+
 /**
  * ISandboxDelegate that dispatches postMessage through the Kotlin delegate
  * via JNI. The Kotlin postMessage uses runOnJSQueueThread to safely access
  * the JSI runtime on the correct thread.
+ *
+ * Holds its own JNI global reference which must be released via invalidate().
  */
 class JNISandboxDelegate : public rnsandbox::ISandboxDelegate {
  public:
-  JNISandboxDelegate(jobject globalDelegateRef)
-      : globalDelegateRef_(globalDelegateRef) {}
+  explicit JNISandboxDelegate(JNIEnv* env, jobject delegateRef)
+      : globalDelegateRef_(env->NewGlobalRef(delegateRef)) {}
+
+  ~JNISandboxDelegate() override { invalidate(); }
+
+  void invalidate() {
+    std::lock_guard<std::mutex> lock(mutex_);
+    if (globalDelegateRef_) {
+      JNIEnv* env = getJNIEnv();
+      if (env) {
+        env->DeleteGlobalRef(globalDelegateRef_);
+      }
+      globalDelegateRef_ = nullptr;
+    }
+  }
 
   void postMessage(const std::string& message) override {
+    std::lock_guard<std::mutex> lock(mutex_);
     JNIEnv* env = getJNIEnv();
     if (!env || !globalDelegateRef_)
       return;
@@ -74,30 +101,9 @@ class JNISandboxDelegate : public rnsandbox::ISandboxDelegate {
 
  private:
   jobject globalDelegateRef_;
-
-  static JNIEnv* getJNIEnv() {
-    JNIEnv* env = nullptr;
-    if (gJavaVM) {
-      gJavaVM->GetEnv(reinterpret_cast<void**>(&env), JNI_VERSION_1_6);
-      if (!env) {
-        gJavaVM->AttachCurrentThread(&env, nullptr);
-      }
-    }
-    return env;
-  }
+  std::mutex mutex_;
 };
 
-static JNIEnv* getJNIEnv() {
-  JNIEnv* env = nullptr;
-  if (gJavaVM) {
-    gJavaVM->GetEnv(reinterpret_cast<void**>(&env), JNI_VERSION_1_6);
-    if (!env) {
-      gJavaVM->AttachCurrentThread(&env, nullptr);
-    }
-  }
-  return env;
-}
-
 static std::string safeGetStringProperty(
     jsi::Runtime& rt,
     const jsi::Object& obj,
@@ -124,7 +130,7 @@ stubJsiFunction(jsi::Runtime& runtime, jsi::Object& object, const char* name) {
 
 static void setupErrorHandler(
     jsi::Runtime& runtime,
-    jobject globalDelegateRef) {
+    std::weak_ptr<SandboxJSIState> stateWeak) {
   jsi::Object global = runtime.global();
   jsi::Value errorUtilsVal = global.getProperty(runtime, "ErrorUtils");
   if (!errorUtilsVal.isObject())
@@ -142,23 +148,27 @@ static void setupErrorHandler(
       runtime,
       jsi::PropNameID::forAscii(runtime, "sandboxGlobalErrorHandler"),
       2,
-      [globalDelegateRef, originalHandler = std::move(originalHandler)](
+      [stateWeak, originalHandler = std::move(originalHandler)](
           jsi::Runtime& rt,
           const jsi::Value&,
           const jsi::Value* args,
           size_t count) -> jsi::Value {
         if (count < 2)
           return jsi::Value::undefined();
 
+        auto state = stateWeak.lock();
+        if (!state || !state->delegateRef)
+          return jsi::Value::undefined();
+
         JNIEnv* jniEnv = getJNIEnv();
         if (!jniEnv)
           return jsi::Value::undefined();
 
-        jclass cls = jniEnv->GetObjectClass(globalDelegateRef);
+        jclass cls = jniEnv->GetObjectClass(state->delegateRef);
         jfieldID hasHandlerField =
             jniEnv->GetFieldID(cls, "hasOnErrorHandler", "Z");
         jboolean hasHandler =
-            jniEnv->GetBooleanField(globalDelegateRef, hasHandlerField);
+            jniEnv->GetBooleanField(state->delegateRef, hasHandlerField);
 
         if (hasHandler) {
           const jsi::Object& error = args[0].asObject(rt);
@@ -175,7 +185,7 @@ static void setupErrorHandler(
           jstring jMsg = jniEnv->NewStringUTF(message.c_str());
           jstring jStack = jniEnv->NewStringUTF(stack.c_str());
           jniEnv->CallVoidMethod(
-              globalDelegateRef,
+              state->delegateRef,
               emitMethod,
               jName,
               jMsg,
@@ -223,7 +233,7 @@ jlong installSandboxJSIBindings(
       runtime,
       jsi::PropNameID::forAscii(runtime, "postMessage"),
       2,
-      [globalDelegateRef](
+      [stateWeak = std::weak_ptr<SandboxJSIState>(state)](
           jsi::Runtime& rt,
           const jsi::Value&,
           const jsi::Value* args,
@@ -238,6 +248,10 @@ jlong installSandboxJSIBindings(
               rt, "postMessage: first argument must be an object");
         }
 
+        auto statePtr = stateWeak.lock();
+        if (!statePtr || !statePtr->delegateRef)
+          return jsi::Value::undefined();
+
         jsi::Object jsonObj = rt.global().getPropertyAsObject(rt, "JSON");
         jsi::Function stringify =
             jsonObj.getPropertyAsFunction(rt, "stringify");
@@ -255,7 +269,6 @@ jlong installSandboxJSIBindings(
           }
           std::string targetOrigin = args[1].getString(rt).utf8(rt);
 
-          // Fan out to all delegates registered for the target origin
           auto& registry = rnsandbox::SandboxRegistry::getInstance();
           auto targets = registry.findAll(targetOrigin);
           if (!targets.empty()) {
@@ -266,11 +279,11 @@ jlong installSandboxJSIBindings(
             LOGW("postMessage: target '%s' not found", targetOrigin.c_str());
           }
         } else {
-          jclass cls = jniEnv->GetObjectClass(globalDelegateRef);
+          jclass cls = jniEnv->GetObjectClass(statePtr->delegateRef);
           jmethodID mid = jniEnv->GetMethodID(
               cls, "emitOnMessageFromJS", "(Ljava/lang/String;)V");
           jstring jMsg = jniEnv->NewStringUTF(messageJson.c_str());
-          jniEnv->CallVoidMethod(globalDelegateRef, mid, jMsg);
+          jniEnv->CallVoidMethod(statePtr->delegateRef, mid, jMsg);
           jniEnv->DeleteLocalRef(jMsg);
           jniEnv->DeleteLocalRef(cls);
         }
@@ -356,7 +369,7 @@ jlong installSandboxJSIBindings(
       makePropDesc(std::move(setOnMessageFn)));
 
   try {
-    setupErrorHandler(runtime, globalDelegateRef);
+    setupErrorHandler(runtime, std::weak_ptr<SandboxJSIState>(state));
   } catch (const std::exception& e) {
     LOGW("Failed to setup error handler: %s", e.what());
   }
@@ -379,7 +392,7 @@ jlong installSandboxJSIBindings(
         if (!origin.empty()) {
           state->origin = origin;
           auto delegate =
-              std::make_shared<JNISandboxDelegate>(globalDelegateRef);
+              std::make_shared<JNISandboxDelegate>(jniEnv, globalDelegateRef);
           state->registryDelegate = delegate;
           rnsandbox::SandboxRegistry::getInstance().registerSandbox(
               origin, delegate, std::set<std::string>());
@@ -481,7 +494,7 @@ Java_io_callstack_rnsandbox_SandboxJSIInstaller_nativeInstallErrorHandler(
     return;
 
   try {
-    setupErrorHandler(*state->runtime, state->delegateRef);
+    setupErrorHandler(*state->runtime, std::weak_ptr<SandboxJSIState>(state));
   } catch (const std::exception& e) {
     LOGW("Failed to setup error handler post-bundle: %s", e.what());
   }
@@ -497,19 +510,28 @@ Java_io_callstack_rnsandbox_SandboxJSIInstaller_nativeDestroy(
   if (it != gStates.end()) {
     std::string origin;
     std::shared_ptr<rnsandbox::ISandboxDelegate> delegate;
+    jobject delegateRef = nullptr;
     {
       std::lock_guard<std::mutex> stateLock(it->second->mutex);
       origin = it->second->origin;
       delegate = it->second->registryDelegate;
+      delegateRef = it->second->delegateRef;
       it->second->onMessageCallback.reset();
       it->second->pendingMessages.clear();
       it->second->runtime = nullptr;
       it->second->registryDelegate.reset();
+      it->second->delegateRef = nullptr;
     }
     if (!origin.empty() && delegate) {
       rnsandbox::SandboxRegistry::getInstance().unregisterDelegate(
           origin, delegate);
     }
+    if (delegateRef) {
+      JNIEnv* env = getJNIEnv();
+      if (env) {
+        env->DeleteGlobalRef(delegateRef);
+      }
+    }
     gStates.erase(it);
   }
 }
diff --git a/packages/react-native-sandbox/cxx/SandboxDelegateWrapper.h b/packages/react-native-sandbox/cxx/SandboxDelegateWrapper.h
@@ -1,11 +1,8 @@
 #pragma once
 
-#include <memory>
 #include <string>
-#include <vector>
 #include "ISandboxDelegate.h"
 
-// Forward declaration to avoid including Objective-C headers in C++ interface
 #ifdef __OBJC__
 @class SandboxReactNativeDelegate;
 #else
@@ -18,21 +15,18 @@ namespace rnsandbox {
  * C++ wrapper for SandboxReactNativeDelegate that implements ISandboxDelegate.
  * This allows the C++ registry to work with Objective-C++ objects through
  * a proper C++ interface.
+ *
+ * The wrapped delegate pointer is non-owning. Call invalidate() before the
+ * delegate is deallocated to prevent dangling pointer access.
  */
 class SandboxDelegateWrapper : public ISandboxDelegate {
  public:
-  /**
-   * Creates a wrapper around the given Objective-C++ delegate.
-   * @param delegate The Objective-C++ delegate to wrap (must not be null)
-   */
   explicit SandboxDelegateWrapper(SandboxReactNativeDelegate* delegate);
 
-  /**
-   * Destructor - does not delete the wrapped delegate
-   */
   ~SandboxDelegateWrapper() override = default;
 
-  // ISandboxDelegate implementation
+  void invalidate() { delegate_ = nullptr; }
+
   void postMessage(const std::string& message) override;
   bool routeMessage(const std::string& message, const std::string& targetId)
       override;
@@ -41,7 +35,6 @@ class SandboxDelegateWrapper : public ISandboxDelegate {
   void setAllowedTurboModules(const std::set<std::string>& modules) override;
 
  private:
-  // Non-owning pointer to the Objective-C++ delegate
   SandboxReactNativeDelegate* delegate_;
 };
 
diff --git a/packages/react-native-sandbox/cxx/StubTurboModuleCxx.cpp b/packages/react-native-sandbox/cxx/StubTurboModuleCxx.cpp
@@ -34,15 +34,15 @@ facebook::jsi::Function StubTurboModuleCxx::createStubFunction(
       runtime,
       facebook::jsi::PropNameID::forAscii(runtime, methodName.c_str()),
       0,
-      [this, methodName](
-          facebook::jsi::Runtime& rt,
-          const facebook::jsi::Value& thisVal,
-          const facebook::jsi::Value* args,
-          size_t count) -> facebook::jsi::Value {
+      [moduleName = moduleName_, methodName](
+          facebook::jsi::Runtime&,
+          const facebook::jsi::Value&,
+          const facebook::jsi::Value*,
+          size_t) -> facebook::jsi::Value {
         SANDBOX_LOG_WARN(
             "[StubTurboModuleCxx] Method call '%s' blocked on module '%s'.",
             methodName.c_str(),
-            this->moduleName_.c_str());
+            moduleName.c_str());
         return facebook::jsi::Value::undefined();
       });
 }
diff --git a/packages/react-native-sandbox/ios/SandboxReactNativeDelegate.mm b/packages/react-native-sandbox/ios/SandboxReactNativeDelegate.mm
