From 4c97097757919580a380b70785e6c161de6b03cc Mon Sep 17 00:00:00 2001
From: AAGaming <aagaming@riseup.net>
Date: Thu, 8 Aug 2024 14:18:19 -0400
Subject: [PATCH] fix(utils/react): fix potential race condition in
 injectFCTrampoline

---
 src/utils/react/fc.ts | 50 ++++++++++++++++++++++++++++++++++---------
 1 file changed, 40 insertions(+), 10 deletions(-)

diff --git a/src/utils/react/fc.ts b/src/utils/react/fc.ts
index f4add16..0acd9b5 100644
--- a/src/utils/react/fc.ts
+++ b/src/utils/react/fc.ts
@@ -52,6 +52,7 @@ export function injectFCTrampoline(component: FC, customHooks?: any): FCTrampoli
             // we have to redirect this to return an object with component's prototype as a constructor returning a value overrides its prototype
             window.SP_REACT.createElement = () => {
                 loggingEnabled && logger.debug("createElement hook called");
+                loggingEnabled && console.trace("createElement trace");
                 return Object.create(component.prototype);
             };
         }
@@ -66,12 +67,16 @@ export function injectFCTrampoline(component: FC, customHooks?: any): FCTrampoli
         }
     }
 
+    let renderHookStep = 0;
+
     // Accessed two times, once directly before class instantiation, and again in some extra logic we don't need to worry about that we hanlde below just in case.
     Object.defineProperty(component, "contextType", {
         configurable: true,
         get: function () {
-            loggingEnabled && logger.debug("get contexttype", this, stubsApplied);
-            applyStubsIfNeeded();
+            loggingEnabled && logger.debug("get contexttype", this, stubsApplied, renderHookStep);
+            loggingEnabled && console.trace("contextType trace");
+            if (renderHookStep == 0) renderHookStep = 1;
+            else if (renderHookStep == 3) renderHookStep = 4;
             return this._contextType;
         },
         set: function (value) {
@@ -79,16 +84,20 @@ export function injectFCTrampoline(component: FC, customHooks?: any): FCTrampoli
         }
     });
 
-    // Undoes the second contextType access we can't detect shortly before render before it's able to cause any damage
-    Object.defineProperty(component, "getDerivedStateFromProps", {
+    // Always accessed directly after contextType for the path we want to catch.
+    Object.defineProperty(component, "contextTypes", {
         configurable: true,
         get: function () {
-            loggingEnabled && logger.debug("get getDerivedStateFromProps", this, stubsApplied);
-            removeStubsIfNeeded();
-            return this._getDerivedStateFromProps;
+            loggingEnabled && logger.debug("get contexttypes", this, stubsApplied, renderHookStep);
+            loggingEnabled && console.trace("contextTypes trace");
+            if (renderHookStep == 1) {
+                renderHookStep = 2;
+                applyStubsIfNeeded();
+            };
+            return this._contextTypes;
         },
         set: function (value) {
-            this._getDerivedStateFromProps = value;
+            this._contextTypes = value;
         }
     });
 
@@ -99,11 +108,32 @@ export function injectFCTrampoline(component: FC, customHooks?: any): FCTrampoli
             return this._updater;
         },
         set: function (value) {
-            loggingEnabled && logger.debug("set updater", this, value, stubsApplied);
-            removeStubsIfNeeded();
+            loggingEnabled && logger.debug("set updater", this, value, stubsApplied, renderHookStep);
+            loggingEnabled && console.trace("updater trace");
+            if (renderHookStep == 2) {
+                renderHookStep = 0;
+                removeStubsIfNeeded();
+            }
             return this._updater = value;
         }
     });
 
+    // Prevents the second contextType+contextTypes access from leaving its hooks around
+    Object.defineProperty(component, "getDerivedStateFromProps", {
+        configurable: true,
+        get: function () {
+            loggingEnabled && logger.debug("get getDerivedStateFromProps", this, stubsApplied, renderHookStep);
+            loggingEnabled && console.trace("getDerivedStateFromProps trace");
+            if (renderHookStep == 2) {
+                renderHookStep = 0;
+                removeStubsIfNeeded();
+            }
+            return this._getDerivedStateFromProps;
+        },
+        set: function (value) {
+            this._getDerivedStateFromProps = value;
+        }
+    });
+
     return userComponent;
 }
\ No newline at end of file