From a4329f4bf8065183c813bff9a1ca46f5b596441b Mon Sep 17 00:00:00 2001
From: Josh Story <story@hey.com>
Date: Thu, 10 Mar 2022 10:55:41 -0800
Subject: [PATCH] add lazy initialization to resolveModelToJson

adding lazying initialization toResolveModelToJson means we use attemptResolveElement's full logic on whatever the resolved type ends up being. This better aligns handling of misued Lazy types like a lazy element being used as a Component or a lazy Component being used as an element.
---
 .../src/__tests__/ReactFlight-test.js         | 74 +++++++++++++++++++
 .../react-server/src/ReactFlightServer.js     |  9 +--
 2 files changed, 77 insertions(+), 6 deletions(-)

diff --git a/packages/react-client/src/__tests__/ReactFlight-test.js b/packages/react-client/src/__tests__/ReactFlight-test.js
index df8476130e32a..6213f0b72e08e 100644
--- a/packages/react-client/src/__tests__/ReactFlight-test.js
+++ b/packages/react-client/src/__tests__/ReactFlight-test.js
@@ -184,6 +184,43 @@ describe('ReactFlight', () => {
     );
   });
 
+  it('errors on a Lazy element being used in Component position', async () => {
+    function SharedComponent({text}) {
+      return (
+        <div>
+          shared<span>{text}</span>
+        </div>
+      );
+    }
+
+    let load = null;
+
+    const LazyElementDisguisedAsComponent = React.lazy(() => {
+      return new Promise(res => {
+        load = () => res({default: <SharedComponent text={'a'} />});
+      });
+    });
+
+    function ServerComponent() {
+      return (
+        <React.Suspense fallback={'Loading...'}>
+          <LazyElementDisguisedAsComponent text={'b'} />
+        </React.Suspense>
+      );
+    }
+
+    const transport = ReactNoopFlightServer.render(<ServerComponent />);
+
+    act(() => {
+      const rootModel = ReactNoopFlightClient.read(transport);
+      ReactNoop.render(rootModel);
+    });
+    expect(ReactNoop).toMatchRenderedOutput('Loading...');
+    spyOnDevAndProd(console, 'error');
+    await load();
+    expect(console.error).toHaveBeenCalledTimes(1);
+  });
+
   it('can render a lazy element', async () => {
     function SharedComponent({text}) {
       return (
@@ -229,6 +266,43 @@ describe('ReactFlight', () => {
     );
   });
 
+  it('errors with lazy value in element position that resolves to Component', async () => {
+    function SharedComponent({text}) {
+      return (
+        <div>
+          shared<span>{text}</span>
+        </div>
+      );
+    }
+
+    let load = null;
+
+    const componentDisguisedAsElement = React.lazy(() => {
+      return new Promise(res => {
+        load = () => res({default: SharedComponent});
+      });
+    });
+
+    function ServerComponent() {
+      return (
+        <React.Suspense fallback={'Loading...'}>
+          {componentDisguisedAsElement}
+        </React.Suspense>
+      );
+    }
+
+    const transport = ReactNoopFlightServer.render(<ServerComponent />);
+
+    act(() => {
+      const rootModel = ReactNoopFlightClient.read(transport);
+      ReactNoop.render(rootModel);
+    });
+    expect(ReactNoop).toMatchRenderedOutput('Loading...');
+    spyOnDevAndProd(console, 'error');
+    await load();
+    expect(console.error).toHaveBeenCalledTimes(1);
+  });
+
   it('can render a lazy module reference', async () => {
     function ClientComponent() {
       return <div>I am client</div>;
diff --git a/packages/react-server/src/ReactFlightServer.js b/packages/react-server/src/ReactFlightServer.js
index 7ea3e1ef83894..e08f308a32fd5 100644
--- a/packages/react-server/src/ReactFlightServer.js
+++ b/packages/react-server/src/ReactFlightServer.js
@@ -210,11 +210,6 @@ function attemptResolveElement(
         const render = type.render;
         return render(props, undefined);
       }
-      case REACT_ELEMENT_TYPE: {
-        // this can happen when a lazy component resolves to an element instead of
-        // a Component.
-        return attemptResolveElement(type.type, type.key, type.ref, type.props);
-      }
       case REACT_MEMO_TYPE: {
         return attemptResolveElement(type.type, key, ref, props);
       }
@@ -508,7 +503,9 @@ export function resolveModelToJSON(
           break;
         }
         case REACT_LAZY_TYPE: {
-          value = attemptResolveElement(value, null, null, {});
+          const payload = (value: any)._payload;
+          const init = (value: any)._init;
+          value = init(payload);
           break;
         }
       }