Optimize binary value checking

Instead of copying all typed arrays that we serialize into a string,
we only do that if the length of the data is the length of any tainted
binary value first.

You're not really supposed to use this with arbitrary binary data but
mostly things like hashes or tokens that end up with the same length.

Author: sebmarkbage

packages/react-server/src/ReactFlightServer.js

@@ -206,6 +206,7 @@ export type Request = {
 const {
   TaintRegistryObjects,
   TaintRegistryValues,
+  TaintRegistryByteLengths,
   ReactCurrentDispatcher,
   ReactCurrentCache,
 } = ReactServerSharedInternals;
@@ -796,14 +797,15 @@ function serializeTypedArray(
   typedArray: $ArrayBufferView,
 ): string {
   if (enableTaint) {
-    // TODO: This is way too slow for the common case. We should first check if
-    // there even could be a match such as if there are any tainted arrays of
-    // equal size.
-    const tainted = TaintRegistryValues.get(
-      binaryToComparableString(typedArray),
-    );
-    if (tainted !== undefined) {
-      throwTaintViolation(tainted.message);
+    if (TaintRegistryByteLengths.has(typedArray.byteLength)) {
+      // If we have had any tainted values of this length, we check
+      // to see if these bytes matches any entries in the registry.
+      const tainted = TaintRegistryValues.get(
+        binaryToComparableString(typedArray),
+      );
+      if (tainted !== undefined) {
+        throwTaintViolation(tainted.message);
+      }
     }
   }
   request.pendingChunks += 2;

packages/react/src/ReactServerSharedInternals.js

@@ -7,7 +7,11 @@
 
 import ReactCurrentDispatcher from './ReactCurrentDispatcher';
 import ReactCurrentCache from './ReactCurrentCache';
-import {TaintRegistryObjects, TaintRegistryValues} from './ReactTaintRegistry';
+import {
+  TaintRegistryObjects,
+  TaintRegistryValues,
+  TaintRegistryByteLengths,
+} from './ReactTaintRegistry';
 
 import {enableTaint} from 'shared/ReactFeatureFlags';
 
@@ -19,6 +23,8 @@ const ReactServerSharedInternals = {
 if (enableTaint) {
   ReactServerSharedInternals.TaintRegistryObjects = TaintRegistryObjects;
   ReactServerSharedInternals.TaintRegistryValues = TaintRegistryValues;
+  ReactServerSharedInternals.TaintRegistryByteLengths =
+    TaintRegistryByteLengths;
 }
 
 export default ReactServerSharedInternals;

packages/react/src/ReactTaint.js

@@ -12,7 +12,8 @@ import {enableTaint, enableBinaryFlight} from 'shared/ReactFeatureFlags';
 import binaryToComparableString from 'shared/binaryToComparableString';
 
 import ReactServerSharedInternals from './ReactServerSharedInternals';
-const {TaintRegistryObjects, TaintRegistryValues} = ReactServerSharedInternals;
+const {TaintRegistryObjects, TaintRegistryValues, TaintRegistryByteLengths} =
+  ReactServerSharedInternals;
 
 interface Reference {}
 
@@ -74,6 +75,7 @@ export function taintValue(
     // hashing in the Map implementation. It doesn't really matter what form the string
     // take as long as it's the same when we look it up.
     // We're not too worried about collisions since this should be a high entropy value.
+    TaintRegistryByteLengths.add(value.byteLength);
     entryValue = binaryToComparableString(value);
   } else {
     const kind = value === null ? 'null' : typeof value;

packages/react/src/ReactTaintRegistry.js

@@ -16,3 +16,6 @@ type TaintEntry = {
 
 export const TaintRegistryObjects: WeakMap<Reference, string> = new WeakMap();
 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();