Keep values alive until the end of any pending requests

A request is the only thing that is expected to do any work.
The principle is that you can derive values from out of a tainted
entry during a request. Including stashing it in a per request cache.

What you can't do is store a derived value in a global module level
cache. At least not without also tainting the object.

Author: sebmarkbage

packages/react-client/src/__tests__/ReactFlight-test.js

@@ -10,6 +10,23 @@
 
 'use strict';
 
+const heldValues = [];
+let finalizationCallback;
+function FinalizationRegistryMock(callback) {
+  finalizationCallback = callback;
+}
+FinalizationRegistryMock.prototype.register = function (target, heldValue) {
+  heldValues.push(heldValue);
+};
+global.FinalizationRegistry = FinalizationRegistryMock;
+
+function gc() {
+  for (let i = 0; i < heldValues.length; i++) {
+    finalizationCallback(heldValues[i]);
+  }
+  heldValues.length = 0;
+}
+
 let act;
 let use;
 let startTransition;
@@ -1589,4 +1606,58 @@ describe('ReactFlight', () => {
 
     expect(errors).toEqual(['Cannot pass a secret token to the client']);
   });
+
+  // @gate enableTaint
+  it('keep a tainted value tainted until the end of any pending requests', async () => {
+    function UserClient({user}) {
+      return <span>{user.name}</span>;
+    }
+    const User = clientReference(UserClient);
+
+    function getUser() {
+      const user = {
+        name: 'Seb',
+        token: '3e971ecc1485fe78625598bf9b6f85db',
+      };
+      ReactServer.unstable_taintValue(
+        'Cannot pass a secret token to the client',
+        user,
+        user.token,
+      );
+      return user;
+    }
+
+    function App() {
+      const user = getUser();
+      const derivedValue = {...user};
+      // A garbage collection can happen at any time. Even before the end of
+      // this request. This would clean up the user object.
+      gc();
+      // We should still block the tainted value.
+      return <User user={derivedValue} />;
+    }
+
+    let errors = [];
+    ReactNoopFlightServer.render(<App />, {
+      onError(x) {
+        errors.push(x.message);
+      },
+    });
+
+    expect(errors).toEqual(['Cannot pass a secret token to the client']);
+
+    // After the previous requests finishes, the token can be rendered again.
+
+    errors = [];
+    ReactNoopFlightServer.render(
+      <User user={{token: '3e971ecc1485fe78625598bf9b6f85db'}} />,
+      {
+        onError(x) {
+          errors.push(x.message);
+        },
+      },
+    );
+
+    expect(errors).toEqual([]);
+  });
 });

packages/react-server/src/ReactFlightServer.js

@@ -198,6 +198,7 @@ export type Request = {
   writtenProviders: Map<string, number>,
   identifierPrefix: string,
   identifierCount: number,
+  taintCleanupQueue: Array<string | bigint>,
   onError: (error: mixed) => ?string,
   onPostpone: (reason: string) => void,
   toJSON: (key: string, value: ReactClientValue) => ReactJSONValue,
@@ -207,6 +208,7 @@ const {
   TaintRegistryObjects,
   TaintRegistryValues,
   TaintRegistryByteLengths,
+  TaintRegistryPendingRequests,
   ReactCurrentDispatcher,
   ReactCurrentCache,
 } = ReactServerSharedInternals;
@@ -216,6 +218,23 @@ function throwTaintViolation(message: string) {
   throw new Error(message);
 }
 
+function cleanupTaintQueue(request: Request): void {
+  const cleanupQueue = request.taintCleanupQueue;
+  TaintRegistryPendingRequests.delete(cleanupQueue);
+  for (let i = 0; i < cleanupQueue.length; i++) {
+    const entryValue = cleanupQueue[i];
+    const entry = TaintRegistryValues.get(entryValue);
+    if (entry !== undefined) {
+      if (entry.count === 1) {
+        TaintRegistryValues.delete(entryValue);
+      } else {
+        entry.count--;
+      }
+    }
+  }
+  cleanupQueue.length = 0;
+}
+
 function defaultErrorHandler(error: mixed) {
   console['error'](error);
   // Don't transform to our wrapper
@@ -250,6 +269,10 @@ export function createRequest(
 
   const abortSet: Set<Task> = new Set();
   const pingedTasks: Array<Task> = [];
+  const cleanupQueue: Array<string | bigint> = [];
+  if (enableTaint) {
+    TaintRegistryPendingRequests.add(cleanupQueue);
+  }
   const hints = createHints();
   const request: Request = {
     status: OPEN,
@@ -273,6 +296,7 @@ export function createRequest(
     writtenProviders: new Map(),
     identifierPrefix: identifierPrefix || '',
     identifierCount: 1,
+    taintCleanupQueue: cleanupQueue,
     onError: onError === undefined ? defaultErrorHandler : onError,
     onPostpone: onPostpone === undefined ? defaultPostponeHandler : onPostpone,
     // $FlowFixMe[missing-this-annot]
@@ -1249,6 +1273,9 @@ function logRecoverableError(request: Request, error: mixed): string {
 }
 
 function fatalError(request: Request, error: mixed): void {
+  if (enableTaint) {
+    cleanupTaintQueue(request);
+  }
   // This is called outside error handling code such as if an error happens in React internals.
   if (request.destination !== null) {
     request.status = CLOSED;
@@ -1573,6 +1600,9 @@ function flushCompletedChunks(
   flushBuffered(destination);
   if (request.pendingChunks === 0) {
     // We're done.
+    if (enableTaint) {
+      cleanupTaintQueue(request);
+    }
     close(destination);
   }
 }

packages/react/src/ReactServerSharedInternals.js

@@ -11,6 +11,7 @@ import {
   TaintRegistryObjects,
   TaintRegistryValues,
   TaintRegistryByteLengths,
+  TaintRegistryPendingRequests,
 } from './ReactTaintRegistry';
 
 import {enableTaint} from 'shared/ReactFeatureFlags';
@@ -25,6 +26,8 @@ if (enableTaint) {
   ReactServerSharedInternals.TaintRegistryValues = TaintRegistryValues;
   ReactServerSharedInternals.TaintRegistryByteLengths =
     TaintRegistryByteLengths;
+  ReactServerSharedInternals.TaintRegistryPendingRequests =
+    TaintRegistryPendingRequests;
 }
 
 export default ReactServerSharedInternals;

packages/react/src/ReactTaint.js

@@ -12,8 +12,12 @@ import {enableTaint, enableBinaryFlight} from 'shared/ReactFeatureFlags';
 import binaryToComparableString from 'shared/binaryToComparableString';
 
 import ReactServerSharedInternals from './ReactServerSharedInternals';
-const {TaintRegistryObjects, TaintRegistryValues, TaintRegistryByteLengths} =
-  ReactServerSharedInternals;
+const {
+  TaintRegistryObjects,
+  TaintRegistryValues,
+  TaintRegistryByteLengths,
+  TaintRegistryPendingRequests,
+} = ReactServerSharedInternals;
 
 interface Reference {}
 
@@ -29,6 +33,10 @@ const defaultMessage =
 function cleanup(entryValue: string | bigint): void {
   const entry = TaintRegistryValues.get(entryValue);
   if (entry !== undefined) {
+    TaintRegistryPendingRequests.forEach(function (requestQueue) {
+      requestQueue.push(entryValue);
+      entry.count++;
+    });
     if (entry.count === 1) {
       TaintRegistryValues.delete(entryValue);
     } else {

packages/react/src/ReactTaintRegistry.js

@@ -19,3 +19,9 @@ export const TaintRegistryValues: Map<string | bigint, TaintEntry> = new Map();
 // Byte lengths of all binary values we've ever seen. We don't both refcounting this.
 // We expect to see only a few lengths here such as the length of token.
 export const TaintRegistryByteLengths: Set<number> = new Set();
+
+// When a value is finalized, it means that it has been removed from any global caches.
+// No future requests can get a handle on it but any ongoing requests can still have
+// a handle on it. It's still tainted until that happens.
+type RequestCleanupQueue = Array<string | bigint>;
+export const TaintRegistryPendingRequests: Set<RequestCleanupQueue> = new Set();