From 3a713c5d343a7a6886ed70a953df6310bf4a4055 Mon Sep 17 00:00:00 2001
From: Sebastian Markbage <sebastian@calyptus.eu>
Date: Wed, 18 Dec 2024 15:36:14 -0500
Subject: [PATCH] Reject any new Chunks not yet discovered at the time of
 reportGlobalError

---
 .../__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 0c6ec9599e1e2..f6157dff171d7 100644
--- a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReplyEdge-test.js
+++ b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReplyEdge-test.js
@@ -245,4 +245,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 930998ea3d889..4db7571bb66a6 100644
--- a/packages/react-server/src/ReactFlightReplyServer.js
+++ b/packages/react-server/src/ReactFlightReplyServer.js
@@ -169,6 +169,8 @@ export type Response = {
   _prefix: string,
   _formData: FormData,
   _chunks: Map<number, SomeChunk<any>>,
+  _closed: boolean,
+  _closedReason: mixed,
   _temporaryReferences: void | TemporaryReferenceSet,
 };
 
@@ -255,6 +257,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,
@@ -493,6 +503,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
@@ -514,6 +526,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);
@@ -1082,6 +1098,8 @@ export function createResponse(
     _prefix: formFieldPrefix,
     _formData: backingFormData,
     _chunks: chunks,
+    _closed: false,
+    _closedReason: null,
     _temporaryReferences: temporaryReferences,
   };
   return response;