From 5ad02213a7678e3c13a67250648e60e2476e5e7c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Markb=C3=A5ge?= <sebastian@calyptus.eu>
Date: Wed, 18 Dec 2024 15:56:57 -0500
Subject: [PATCH] [Flight Reply] Reject any new Chunks not yet discovered at
 the time of reportGlobalError (#31840)

We might have already resolved models that are not pending and so are
not rejected by aborting the stream. When those later get parsed they
might discover new chunks which end up as pending. These should be
errored since they will never be able to resolve later.

This avoids infinitely hanging the stream.

This same fix needs to be ported to ReactFlightClient that has the same
issue.
---
 .../__tests__/ReactFlightDOMReplyEdge-test.js | 27 +++++++++++++++++++
 .../src/ReactFlightReplyServer.js             | 18 +++++++++++++
 2 files changed, 45 insertions(+)

diff --git a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReplyEdge-test.js b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReplyEdge-test.js
index 8e45472956294..83b127752de7a 100644
--- a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReplyEdge-test.js
+++ b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReplyEdge-test.js
@@ -45,4 +45,31 @@ describe('ReactFlightDOMReplyEdge', () => {
 
     expect(decoded).toEqual({some: 'object'});
   });
+
+  it('should abort when parsing an incomplete payload', async () => {
+    const infinitePromise = new Promise(() => {});
+    const controller = new AbortController();
+    const promiseForResult = ReactServerDOMClient.encodeReply(
+      {promise: infinitePromise},
+      {
+        signal: controller.signal,
+      },
+    );
+    controller.abort();
+    const body = await promiseForResult;
+
+    const decoded = await ReactServerDOMServer.decodeReply(
+      body,
+      webpackServerMap,
+    );
+
+    let error = null;
+    try {
+      await decoded.promise;
+    } catch (x) {
+      error = x;
+    }
+    expect(error).not.toBe(null);
+    expect(error.message).toBe('Connection closed.');
+  });
 });
diff --git a/packages/react-server/src/ReactFlightReplyServer.js b/packages/react-server/src/ReactFlightReplyServer.js
index 291da7870760e..e2056d7df2650 100644
--- a/packages/react-server/src/ReactFlightReplyServer.js
+++ b/packages/react-server/src/ReactFlightReplyServer.js
@@ -135,6 +135,8 @@ export type Response = {
   _formData: FormData,
   _chunks: Map<number, SomeChunk<any>>,
   _fromJSON: (key: string, value: JSONValue) => any,
+  _closed: boolean,
+  _closedReason: mixed,
 };
 
 export function getRoot<T>(response: Response): Thenable<T> {
@@ -198,6 +200,14 @@ function createResolvedModelChunk<T>(
   return new Chunk(RESOLVED_MODEL, value, null, response);
 }
 
+function createErroredChunk<T>(
+  response: Response,
+  reason: mixed,
+): ErroredChunk<T> {
+  // $FlowFixMe[invalid-constructor] Flow doesn't support functions as constructors
+  return new Chunk(ERRORED, null, reason, response);
+}
+
 function resolveModelChunk<T>(chunk: SomeChunk<T>, value: string): void {
   if (chunk.status !== PENDING) {
     // We already resolved. We didn't expect to see this.
@@ -297,6 +307,8 @@ function initializeModelChunk<T>(chunk: ResolvedModelChunk<T>): void {
 // Report that any missing chunks in the model is now going to throw this
 // error upon read. Also notify any pending promises.
 export function reportGlobalError(response: Response, error: Error): void {
+  response._closed = true;
+  response._closedReason = error;
   response._chunks.forEach(chunk => {
     // If this chunk was already resolved or errored, it won't
     // trigger an error but if it wasn't then we need to
@@ -318,6 +330,10 @@ function getChunk(response: Response, id: number): SomeChunk<any> {
     if (backingEntry != null) {
       // We assume that this is a string entry for now.
       chunk = createResolvedModelChunk(response, (backingEntry: any));
+    } else if (response._closed) {
+      // We have already errored the response and we're not going to get
+      // anything more streaming in so this will immediately error.
+      chunk = createErroredChunk(response, response._closedReason);
     } else {
       // We're still waiting on this entry to stream in.
       chunk = createPendingChunk(response);
@@ -519,6 +535,8 @@ export function createResponse(
       }
       return value;
     },
+    _closed: false,
+    _closedReason: null,
   };
   return response;
 }