[Flight]: Client-side registerServerReference must not break .bind() (#32565)

Author: unstubbable

packages/react-client/src/ReactFlightReplyClient.js

@@ -63,10 +63,14 @@ export type EncodeFormActionCallback = <A>(
 
 export type ServerReferenceId = any;
 
-const knownServerReferences: WeakMap<
-  Function,
-  {id: ServerReferenceId, bound: null | Thenable<Array<any>>},
-> = new WeakMap();
+type ServerReferenceClosure = {
+  id: ServerReferenceId,
+  originalBind: Function,
+  bound: null | Thenable<Array<any>>,
+};
+
+const knownServerReferences: WeakMap<Function, ServerReferenceClosure> =
+  new WeakMap();
 
 // Serializable values
 export type ReactServerValue =
@@ -760,16 +764,17 @@ export function processReply(
     }
 
     if (typeof value === 'function') {
-      const metaData = knownServerReferences.get(value);
-      if (metaData !== undefined) {
-        const metaDataJSON = JSON.stringify(metaData, resolveToJSON);
+      const referenceClosure = knownServerReferences.get(value);
+      if (referenceClosure !== undefined) {
+        const {id, bound} = referenceClosure;
+        const referenceClosureJSON = JSON.stringify({id, bound}, resolveToJSON);
         if (formData === null) {
           // Upgrade to use FormData to allow us to stream this value.
           formData = new FormData();
         }
         // The reference to this function came from the same client so we can pass it back.
         const refId = nextPartId++;
-        formData.set(formFieldPrefix + refId, metaDataJSON);
+        formData.set(formFieldPrefix + refId, referenceClosureJSON);
         return serializeServerReferenceID(refId);
       }
       if (temporaryReferences !== undefined && key.indexOf(':') === -1) {
@@ -864,7 +869,7 @@ export function processReply(
 }
 
 const boundCache: WeakMap<
-  {id: ServerReferenceId, bound: null | Thenable<Array<any>>},
+  ServerReferenceClosure,
   Thenable<FormData>,
 > = new WeakMap();
 
@@ -905,21 +910,22 @@ function defaultEncodeFormAction(
   this: any => Promise<any>,
   identifierPrefix: string,
 ): ReactCustomFormAction {
-  const reference = knownServerReferences.get(this);
-  if (!reference) {
+  const referenceClosure = knownServerReferences.get(this);
+  if (!referenceClosure) {
     throw new Error(
       'Tried to encode a Server Action from a different instance than the encoder is from. ' +
         'This is a bug in React.',
     );
   }
   let data: null | FormData = null;
   let name;
-  const boundPromise = reference.bound;
+  const boundPromise = referenceClosure.bound;
   if (boundPromise !== null) {
-    let thenable = boundCache.get(reference);
+    let thenable = boundCache.get(referenceClosure);
     if (!thenable) {
-      thenable = encodeFormData(reference);
-      boundCache.set(reference, thenable);
+      const {id, bound} = referenceClosure;
+      thenable = encodeFormData({id, bound});
+      boundCache.set(referenceClosure, thenable);
     }
     if (thenable.status === 'rejected') {
       throw thenable.reason;
@@ -941,7 +947,7 @@ function defaultEncodeFormAction(
     name = '$ACTION_REF_' + identifierPrefix;
   } else {
     // This is the simple case so we can just encode the ID.
-    name = '$ACTION_ID_' + reference.id;
+    name = '$ACTION_ID_' + referenceClosure.id;
   }
   return {
     name: name,
@@ -952,42 +958,42 @@ function defaultEncodeFormAction(
 }
 
 function customEncodeFormAction(
-  proxy: any => Promise<any>,
+  reference: any => Promise<any>,
   identifierPrefix: string,
   encodeFormAction: EncodeFormActionCallback,
 ): ReactCustomFormAction {
-  const reference = knownServerReferences.get(proxy);
-  if (!reference) {
+  const referenceClosure = knownServerReferences.get(reference);
+  if (!referenceClosure) {
     throw new Error(
       'Tried to encode a Server Action from a different instance than the encoder is from. ' +
         'This is a bug in React.',
     );
   }
-  let boundPromise: Promise<Array<any>> = (reference.bound: any);
+  let boundPromise: Promise<Array<any>> = (referenceClosure.bound: any);
   if (boundPromise === null) {
     boundPromise = Promise.resolve([]);
   }
-  return encodeFormAction(reference.id, boundPromise);
+  return encodeFormAction(referenceClosure.id, boundPromise);
 }
 
 function isSignatureEqual(
   this: any => Promise<any>,
   referenceId: ServerReferenceId,
   numberOfBoundArgs: number,
 ): boolean {
-  const reference = knownServerReferences.get(this);
-  if (!reference) {
+  const referenceClosure = knownServerReferences.get(this);
+  if (!referenceClosure) {
     throw new Error(
       'Tried to encode a Server Action from a different instance than the encoder is from. ' +
         'This is a bug in React.',
     );
   }
-  if (reference.id !== referenceId) {
+  if (referenceClosure.id !== referenceId) {
     // These are different functions.
     return false;
   }
   // Now check if the number of bound arguments is the same.
-  const boundPromise = reference.bound;
+  const boundPromise = referenceClosure.bound;
   if (boundPromise === null) {
     // No bound arguments.
     return numberOfBoundArgs === 0;
@@ -1131,10 +1137,20 @@ export function registerBoundServerReference<T: Function>(
   bound: null | Thenable<Array<any>>,
   encodeFormAction: void | EncodeFormActionCallback,
 ): void {
+  if (knownServerReferences.has(reference)) {
+    return;
+  }
+
+  knownServerReferences.set(reference, {
+    id,
+    originalBind: reference.bind,
+    bound,
+  });
+
   // Expose encoder for use by SSR, as well as a special bind that can be used to
   // keep server capabilities.
   if (usedWithSSR) {
-    // Only expose this in builds that would actually use it. Not needed on the client.
+    // Only expose this in builds that would actually use it. Not needed in the browser.
     const $$FORM_ACTION =
       encodeFormAction === undefined
         ? defaultEncodeFormAction
@@ -1154,7 +1170,6 @@ export function registerBoundServerReference<T: Function>(
       bind: {value: bind},
     });
   }
-  knownServerReferences.set(reference, {id, bound});
 }
 
 export function registerServerReference<T: Function>(
@@ -1171,43 +1186,54 @@ const FunctionBind = Function.prototype.bind;
 // $FlowFixMe[method-unbinding]
 const ArraySlice = Array.prototype.slice;
 function bind(this: Function): Function {
-  // $FlowFixMe[unsupported-syntax]
-  // $FlowFixMe[prop-missing]
-  const newFn = FunctionBind.apply(this, arguments);
-  const reference = knownServerReferences.get(this);
-  if (reference) {
-    if (__DEV__) {
-      const thisBind = arguments[0];
-      if (thisBind != null) {
-        // This doesn't warn in browser environments since it's not instrumented outside
-        // usedWithSSR. This makes this an SSR only warning which we don't generally do.
-        // TODO: Consider a DEV only instrumentation in the browser.
-        console.error(
-          'Cannot bind "this" of a Server Action. Pass null or undefined as the first argument to .bind().',
-        );
-      }
-    }
-    const args = ArraySlice.call(arguments, 1);
-    let boundPromise = null;
-    if (reference.bound !== null) {
-      boundPromise = Promise.resolve((reference.bound: any)).then(boundArgs =>
-        boundArgs.concat(args),
+  const referenceClosure = knownServerReferences.get(this);
+
+  if (!referenceClosure) {
+    // $FlowFixMe[prop-missing]
+    return FunctionBind.apply(this, arguments);
+  }
+
+  const newFn = referenceClosure.originalBind.apply(this, arguments);
+
+  if (__DEV__) {
+    const thisBind = arguments[0];
+    if (thisBind != null) {
+      // This doesn't warn in browser environments since it's not instrumented outside
+      // usedWithSSR. This makes this an SSR only warning which we don't generally do.
+      // TODO: Consider a DEV only instrumentation in the browser.
+      console.error(
+        'Cannot bind "this" of a Server Action. Pass null or undefined as the first argument to .bind().',
       );
-    } else {
-      boundPromise = Promise.resolve(args);
-    }
-    // Expose encoder for use by SSR, as well as a special bind that can be used to
-    // keep server capabilities.
-    if (usedWithSSR) {
-      // Only expose this in builds that would actually use it. Not needed on the client.
-      Object.defineProperties((newFn: any), {
-        $$FORM_ACTION: {value: this.$$FORM_ACTION},
-        $$IS_SIGNATURE_EQUAL: {value: isSignatureEqual},
-        bind: {value: bind},
-      });
     }
-    knownServerReferences.set(newFn, {id: reference.id, bound: boundPromise});
   }
+
+  const args = ArraySlice.call(arguments, 1);
+  let boundPromise = null;
+  if (referenceClosure.bound !== null) {
+    boundPromise = Promise.resolve((referenceClosure.bound: any)).then(
+      boundArgs => boundArgs.concat(args),
+    );
+  } else {
+    boundPromise = Promise.resolve(args);
+  }
+
+  knownServerReferences.set(newFn, {
+    id: referenceClosure.id,
+    originalBind: newFn.bind,
+    bound: boundPromise,
+  });
+
+  // Expose encoder for use by SSR, as well as a special bind that can be used to
+  // keep server capabilities.
+  if (usedWithSSR) {
+    // Only expose this in builds that would actually use it. Not needed on the client.
+    Object.defineProperties((newFn: any), {
+      $$FORM_ACTION: {value: this.$$FORM_ACTION},
+      $$IS_SIGNATURE_EQUAL: {value: isSignatureEqual},
+      bind: {value: bind},
+    });
+  }
+
   return newFn;
 }
 

packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMEdge-test.js

@@ -309,6 +309,13 @@ describe('ReactFlightDOMEdge', () => {
       greet,
     });
 
+    // Registering the server reference also with the client must not break
+    // subsequent `.bind` calls.
+    ReactServerDOMClient.registerServerReference(
+      ServerModule.greet,
+      ServerModule.greet.$$id,
+    );
+
     const stream = await serverAct(() =>
       ReactServerDOMServer.renderToReadableStream(
         {