From c0b288ca3fedfd58f332876d0677eca5feedab73 Mon Sep 17 00:00:00 2001
From: Nick Lefever <lefever@meta.com>
Date: Tue, 16 Jul 2024 09:23:42 -0700
Subject: [PATCH] Use weak ptr for runtime shadow node reference (#45463)

Summary:
Pull Request resolved: https://github.com/facebook/react-native/pull/45463

Changelog: [Internal]

State updates will clone shadow nodes for an shadow tree revision that is outdated.

This can lead to accessing deallocated shadow node references because the JS renderer committed a newer revision and deallocated the one used by the pending state update.

By using a weak pointer to hold a reference to the runtime shadow node reference, we can only update references for wrappers that are still valid.

Reviewed By: javache

Differential Revision: D59804999

fbshipit-source-id: 89c9967d139d3cac7d7252994beae419bc591e79
---
 .../ReactCommon/react/renderer/core/ShadowNode.cpp     |  7 ++++---
 .../ReactCommon/react/renderer/core/ShadowNode.h       | 10 +++++-----
 .../react/renderer/core/tests/ShadowNodeTest.cpp       |  2 +-
 .../ReactCommon/react/renderer/uimanager/primitives.h  |  2 +-
 4 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.cpp b/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.cpp
index 438d5fbb9a6..44cca7436b9 100644
--- a/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.cpp
+++ b/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.cpp
@@ -310,7 +310,8 @@ bool ShadowNode::progressStateIfNecessary() {
 }
 
 void ShadowNode::setRuntimeShadowNodeReference(
-    ShadowNodeWrapper* runtimeShadowNodeReference) const {
+    const std::shared_ptr<ShadowNodeWrapper>& runtimeShadowNodeReference)
+    const {
   runtimeShadowNodeReference_ = runtimeShadowNodeReference;
 }
 
@@ -319,8 +320,8 @@ void ShadowNode::transferRuntimeShadowNodeReference(
   destinationShadowNode->runtimeShadowNodeReference_ =
       runtimeShadowNodeReference_;
 
-  if (runtimeShadowNodeReference_ != nullptr) {
-    runtimeShadowNodeReference_->shadowNode = destinationShadowNode;
+  if (auto reference = runtimeShadowNodeReference_.lock()) {
+    reference->shadowNode = destinationShadowNode;
   }
 }
 
diff --git a/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.h b/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.h
index ea5dae23eeb..48196b4475b 100644
--- a/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.h
+++ b/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.h
@@ -190,11 +190,11 @@ class ShadowNode : public Sealable,
   bool progressStateIfNecessary();
 
   /*
-   * Bind the runtime reference to this `ShadowNode` with a raw pointer,
+   * Bind the runtime reference to this `ShadowNode` with a weak pointer,
    * allowing to update the reference to this `ShadowNode` when cloned.
    */
-  void setRuntimeShadowNodeReference(
-      ShadowNodeWrapper* runtimeShadowNodeReference) const;
+  void setRuntimeShadowNodeReference(const std::shared_ptr<ShadowNodeWrapper>&
+                                         runtimeShadowNodeReference) const;
 
   /*
    * Transfer the runtime reference to this `ShadowNode` to a new instance,
@@ -269,9 +269,9 @@ class ShadowNode : public Sealable,
   ShadowNodeTraits traits_;
 
   /*
-   * Pointer to the runtime reference to this `ShadowNode`.
+   * Weak pointer to the runtime reference to this `ShadowNode`.
    */
-  mutable ShadowNodeWrapper* runtimeShadowNodeReference_{};
+  mutable std::weak_ptr<ShadowNodeWrapper> runtimeShadowNodeReference_{};
 };
 
 static_assert(
diff --git a/packages/react-native/ReactCommon/react/renderer/core/tests/ShadowNodeTest.cpp b/packages/react-native/ReactCommon/react/renderer/core/tests/ShadowNodeTest.cpp
index 4a33373969d..1430eb027df 100644
--- a/packages/react-native/ReactCommon/react/renderer/core/tests/ShadowNodeTest.cpp
+++ b/packages/react-native/ReactCommon/react/renderer/core/tests/ShadowNodeTest.cpp
@@ -353,7 +353,7 @@ TEST_P(ShadowNodeTest, testCloneTree) {
 TEST_P(ShadowNodeTest, handleRuntimeReferenceTransferOnClone) {
   auto nodeABRev1 = nodeAB_->clone({});
   auto wrappedShadowNode = std::make_shared<ShadowNodeWrapper>(nodeABRev1);
-  nodeABRev1->setRuntimeShadowNodeReference(&*wrappedShadowNode);
+  nodeABRev1->setRuntimeShadowNodeReference(wrappedShadowNode);
 
   auto nodeABRev2 = nodeABRev1->clone({});
 
diff --git a/packages/react-native/ReactCommon/react/renderer/uimanager/primitives.h b/packages/react-native/ReactCommon/react/renderer/uimanager/primitives.h
index 0a6b29dab58..b88b288c99b 100644
--- a/packages/react-native/ReactCommon/react/renderer/uimanager/primitives.h
+++ b/packages/react-native/ReactCommon/react/renderer/uimanager/primitives.h
@@ -52,7 +52,7 @@ inline static jsi::Value valueFromShadowNode(
 
   if (assignRuntimeShadowNodeReference) {
     wrappedShadowNode->shadowNode->setRuntimeShadowNodeReference(
-        &*wrappedShadowNode);
+        wrappedShadowNode);
   }
 
   jsi::Object obj(runtime);