Copy ArrayBuffer arguments when an ObjC TurboModule method takes a block (#57983)

packages/react-native/React/Base/RCTArrayBuffer.h

@@ -17,7 +17,11 @@ NS_ASSUME_NONNULL_BEGIN
  *  - `YES` — safe to retain and use from any thread (synchronize if aliasing JS
  *     memory).
  *  - `NO` — valid only during the synchronous call on the calling thread; copy with
- *    `arrayBufferWithCopiedBytes:length:` to keep the bytes.
+ *    `arrayBufferWithCopiedBytes:length:` to keep the bytes. Such a buffer must not be captured in
+ *    a block or handed to a callback or promise resolve block, which deliver after the call
+ *    returns. The TurboModule framework copies rather than aliases whenever the method signature
+ *    exposes a block parameter, and buffers nested inside `NSArray` or `NSDictionary` arguments
+ *    always own their bytes.
  */
 @interface RCTArrayBuffer : NSObject
 

packages/react-native/ReactCommon/react/nativemodule/core/platform/ios/ReactCommon/RCTTurboModule.mm

@@ -210,11 +210,29 @@ size_t size() const override
   };
 }
 
+// A block parameter means the module can hand an argument to a callback it invokes after the call
+// returns, so nothing passed to such a method may alias the JS heap.
+static BOOL methodSignatureTakesBlock(NSMethodSignature *methodSignature)
+{
+  const char *BLOCK_TYPE = @encode(
+      __typeof__(^{
+      }));
+
+  for (NSUInteger i = 2; i < methodSignature.numberOfArguments; i++) {
+    const std::string objCArgType = [methodSignature getArgumentTypeAtIndex:i];
+    if (objCArgType == BLOCK_TYPE) {
+      return YES;
+    }
+  }
+  return NO;
+}
+
 // Native-backed buffers are aliased and keep their backing store alive, so they stay valid for
 // as long as the module holds them. JS-heap buffers have nothing to retain — their bytes are
 // freed when the ArrayBuffer is collected or detached — so `mustCopyBytes` is set whenever the
-// invocation may outlive the JS call, and they are aliased only when it cannot. Engines that
-// don't hand out the backing MutableBuffer fall back to those JS-heap rules for every buffer.
+// invocation may outlive the JS call (an async return kind, or a block parameter the module can
+// call later), and they are aliased only when it cannot. Engines that don't hand out the backing
+// MutableBuffer fall back to those JS-heap rules for every buffer.
 static RCTArrayBuffer *
 convertJSIArrayBufferToRCTArrayBuffer(jsi::Runtime &rt, const jsi::ArrayBuffer &arrayBuffer, BOOL mustCopyBytes)
 {
@@ -813,10 +831,15 @@ TraceSection s(
   NSInvocation *inv = [NSInvocation invocationWithMethodSignature:methodSignature];
   [inv setSelector:selector];
 
+  // A block argument can outlive the call, and JS drains it asynchronously, so an ArrayBuffer the
+  // module hands to it must own its bytes even when the method itself returns synchronously.
+  BOOL mustCopyArrayBufferBytes = mustCopyBytes || methodSignatureTakesBlock(methodSignature);
+
   for (size_t i = 0; i < count; i++) {
     const jsi::Value &arg = args[i];
     const std::string objCArgType = [methodSignature getArgumentTypeAtIndex:i + 2];
-    setInvocationArg(runtime, methodName, objCArgType, arg, i, inv, retainedObjectsForInvocation, mustCopyBytes);
+    setInvocationArg(
+        runtime, methodName, objCArgType, arg, i, inv, retainedObjectsForInvocation, mustCopyArrayBufferBytes);
   }
 
   if (isSync) {

packages/rn-tester/RNTesterUnitTests/RCTTurboModuleArrayBufferTests.mm

@@ -14,6 +14,7 @@
 #import <hermes/hermes.h>
 #import <react/bridging/ArrayBuffer.h>
 
+#import <array>
 #import <list>
 #import <vector>
 
@@ -142,6 +143,15 @@ - (void)testMethodWhichStoresArrayBuffer:(RCTArrayBuffer *)payload
   self.lastReceivedPayload = [NSData dataWithBytes:payload.mutableBytes length:payload.length];
 }
 
+- (NSNumber *)testMethodWhichPassesArrayBufferToCallback:(RCTArrayBuffer *)buffer
+                                                callback:(RCTResponseSenderBlock)callback
+{
+  // The buffer escapes into a callback JS drains after this method returns, so the framework must
+  // have handed us a copy rather than an alias of the JS heap.
+  callback(@[ buffer ]);
+  return @(buffer.isOwningBytes);
+}
+
 - (void)testMethodWhichCallsBackWithArrayBuffer:(double)size callback:(RCTResponseSenderBlock)callback
 {
   callback(@[ createIntegerSequenceBuffer(static_cast<NSUInteger>(size)) ]);
@@ -444,6 +454,69 @@ - (void)testCallbackDeliversArrayBuffer
   XCTAssertEqual(callbackBytes[2], 2);
 }
 
+// A sync method that hands its argument to a callback lets the buffer escape the call, so the
+// argument must be copied even though the return kind would otherwise allow aliasing.
+- (void)testArrayBufferIsCopiedWhenAMethodTakesACallback
+{
+  auto hermesRuntime = createHermesRuntime();
+  facebook::jsi::Runtime *rt = hermesRuntime.get();
+  auto jsInvoker = std::make_shared<TestCallInvoker>(*rt);
+  auto *instance = [RCTTestArrayBufferTurboModule new];
+
+  ObjCTurboModule::InitParams params = {
+      .moduleName = "TestModule",
+      .instance = instance,
+      .jsInvoker = jsInvoker,
+      .nativeMethodCallInvoker = std::make_shared<ImmediateNativeMethodCallInvoker>(),
+      .isSyncModule = false,
+  };
+  ObjCTurboModule module(params);
+
+  auto sourceBuffer = rt->global()
+                          .getPropertyAsFunction(*rt, "eval")
+                          .call(*rt, "new Uint8Array([1, 2, 3]).buffer")
+                          .asObject(*rt)
+                          .getArrayBuffer(*rt);
+
+  std::vector<uint8_t> callbackBytes;
+  auto onCallback = facebook::jsi::Function::createFromHostFunction(
+      *rt,
+      facebook::jsi::PropNameID::forAscii(*rt, "onCallback"),
+      1,
+      [&callbackBytes](
+          facebook::jsi::Runtime &runtime,
+          const facebook::jsi::Value &,
+          const facebook::jsi::Value *callbackArgs,
+          size_t count) -> facebook::jsi::Value {
+        if (count == 1) {
+          const auto &callbackArg = *callbackArgs;
+          if (callbackArg.isObject() && callbackArg.asObject(runtime).isArrayBuffer(runtime)) {
+            callbackBytes = bytesFromArrayBuffer(runtime, callbackArg.asObject(runtime).getArrayBuffer(runtime));
+          }
+        }
+        return facebook::jsi::Value::undefined();
+      });
+  std::array<facebook::jsi::Value, 2> args = {
+      facebook::jsi::Value(*rt, sourceBuffer), facebook::jsi::Value(*rt, onCallback)};
+
+  auto result = module.invokeObjCMethod(
+      *rt,
+      BooleanKind,
+      "testMethodWhichPassesArrayBufferToCallback",
+      @selector(testMethodWhichPassesArrayBufferToCallback:callback:),
+      args.data(),
+      2);
+
+  XCTAssertTrue(result.getBool(), @"A buffer that can escape into a callback must own its bytes");
+
+  jsInvoker->flushQueue();
+
+  XCTAssertEqual(callbackBytes.size(), 3u);
+  XCTAssertEqual(callbackBytes[0], 1);
+  XCTAssertEqual(callbackBytes[1], 2);
+  XCTAssertEqual(callbackBytes[2], 3);
+}
+
 - (void)testPromiseResolvesArrayBuffer
 {
   auto hermesRuntime = createHermesRuntime(true);