From 4af40315c864c4467aa1a2fb61b9c05b7d180e3e 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 a02d2b0c0fd81..7e7091016ed4e 100644
--- a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReplyEdge-test.js
+++ b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReplyEdge-test.js
@@ -250,4 +250,31 @@ describe('ReactFlightDOMReplyEdge', () => {
       ),
     );
   });
+
+  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 95746e6e16464..f6c817eafc4f7 100644
--- a/packages/react-server/src/ReactFlightReplyServer.js
+++ b/packages/react-server/src/ReactFlightReplyServer.js
@@ -173,6 +173,8 @@ export type Response = {
   _prefix: string,
   _formData: FormData,
   _chunks: Map<number, SomeChunk<any>>,
+  _closed: boolean,
+  _closedReason: mixed,
   _temporaryReferences: void | TemporaryReferenceSet,
 };
 
@@ -261,6 +263,14 @@ function createResolvedModelChunk<T>(
   return new Chunk(RESOLVED_MODEL, value, id, 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,
@@ -501,6 +511,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
@@ -522,6 +534,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), id);
+    } 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);
@@ -1102,6 +1118,8 @@ export function createResponse(
     _prefix: formFieldPrefix,
     _formData: backingFormData,
     _chunks: chunks,
+    _closed: false,
+    _closedReason: null,
     _temporaryReferences: temporaryReferences,
   };
   return response;