Allow to pass ObjC blocks into folly::Function

Summary:
In iOS blocks are initially allocated on the stack and only lazily copied to the heap (e.g when they are assigned to a variable). This means, if you pass a block as an rvalue to a C++ method that keeps moving it around instead of copy assigning it at some point, the block remains on the stack and will get freed once the original method is done (leading to use after free if the block is executed later).

This was mitigated by deleting the conversion from ObjC functions to folly functions. Given that all we need is to make sure that the block is allocated on the heap (that is it is an instance of NSMallocBlock rather than NSStackBlock), it seems drastic to ban the conversion. ObjC developers tend to be more familiar with ObjC blocks and will find it convenient to use this conversion.

This diff insteads implements the constructor and assignment operator by wrapping the ObjC block in a c++ lambda and capturing it by copy. ARC keeps track of the reference count and automatically releases the block when the lambda is deallocated. Moreover, copy only increase the retain count (instead of doing an actual copy) if the block was already stored on the heap (https://www.cocoawithlove.com/2009/10/how-blocks-are-implemented-and.html section NSMallocBlock never actually copies).

Reviewed By: ericniebler

Differential Revision: D6109932

fbshipit-source-id: 48bb446d3a66f46affba774cfe1cfb8a60c661de

diff --git a/folly/Function.h b/folly/Function.h
index 4bb7e6bea4e1fa90ac08dcd80928ae375b4b3fad..0ad0158292e0a26dad68db9d887350fda0a49cce 100644 (file)
--- a/folly/Function.h
+++ b/folly/Function.h
@@ -480,9 +480,11 @@ class Function final : private detail::function::FunctionTraits<FunctionType> {
   Function(const Function&) = delete;
 
 #if __OBJC__
-  // Delete conversion from Objective-C blocks
+  // Make sure Objective C blocks are copied
   template <class ReturnType, class... Args>
-  Function(ReturnType (^)(Args...)) = delete;
+  /*implicit*/ Function(ReturnType (^objCBlock)(Args... args))
+      : Function([blockCopy = (ReturnType (^)(Args...))[objCBlock copy]](
+                     Args... args) { return blockCopy(args...); }){};
 #endif
 
   /**
@@ -565,9 +567,13 @@ class Function final : private detail::function::FunctionTraits<FunctionType> {
   Function& operator=(const Function&) = delete;
 
 #if __OBJC__
-  // Delete conversion from Objective-C blocks
+  // Make sure Objective C blocks are copied
   template <class ReturnType, class... Args>
-  Function& operator=(ReturnType (^)(Args...)) = delete;
+  /* implicit */ Function &operator=(ReturnType (^objCBlock)(Args... args)) {
+    (*this) = [blockCopy = (ReturnType (^)(Args...))[objCBlock copy]](
+                  Args... args) { return blockCopy(args...); };
+    return *this;
+  }
 #endif
 
   /**