Eliminate promise-punning UB from getErrorHandle

Summary:
As discussed in the preceding D82907354, it is quite hard to eliminate promise-punning UB from `TaskWrapper::get_return_object`, so we just leave it be with the safeguard of `is_promise_type_punning_safe`.

However, there's no inherent reason for `getErrorHandle` to interact with `coroutine_handle<Promise>`. It actually only either needs the `Promise&` (which may be stored as a member of a wrapper promise without issue), or the type-erased `coroutine_handle<>`. Furthermore, the latter is only needed on the "no-op" path, i.e. the per-coro implementation does not need to see a handle at all.

This diff renames `getErrorHandle` implementations to `getErrorHandleImpl` for clarity, and switches them all to the new, UB-free design.

Reviewed By: ispeters

Differential Revision: D82926069

fbshipit-source-id: 3230102ffa12ae466ca822ce6a058092c720ac5d

Author: snarkmaster

third-party/folly/src/folly/coro/AsyncGenerator.h

@@ -565,9 +565,8 @@ class AsyncGeneratorPromise final
   class YieldAwaiter {
    public:
     bool await_ready() noexcept { return false; }
-    coroutine_handle<> await_suspend(
-        coroutine_handle<AsyncGeneratorPromise> h) noexcept {
-      AsyncGeneratorPromise& promise = h.promise();
+    coroutine_handle<> await_suspend_promise(
+        AsyncGeneratorPromise& promise) noexcept {
       // Pop AsyncStackFrame first as clearContext() clears the frame state.
       folly::popAsyncStackFrameCallee(promise.getAsyncFrame());
       promise.clearContext();
@@ -578,6 +577,10 @@ class AsyncGeneratorPromise final
       }
       return promise.continuation_.getHandle();
     }
+    coroutine_handle<> await_suspend(
+        coroutine_handle<AsyncGeneratorPromise> h) noexcept {
+      return await_suspend_promise(h.promise());
+    }
     void await_resume() noexcept {}
   };
 
@@ -815,18 +818,17 @@ class AsyncGeneratorPromise final
 
   folly::AsyncStackFrame& getAsyncFrame() noexcept { return asyncFrame_; }
 
-  static ExtendedCoroutineHandle::ErrorHandle getErrorHandle(
+  static std::optional<ExtendedCoroutineHandle::ErrorHandle> getErrorHandleImpl(
       AsyncGeneratorPromise& me, exception_wrapper& ex) {
     if (me.bypassExceptionThrowing_ == BypassExceptionThrowing::ACTIVE) {
       auto yieldAwaiter = me.yield_value(co_error(std::move(ex)));
       DCHECK(!yieldAwaiter.await_ready());
-      return {
-          yieldAwaiter.await_suspend(
-              coroutine_handle<AsyncGeneratorPromise>::from_promise(me)),
+      return ExtendedCoroutineHandle::ErrorHandle{
+          yieldAwaiter.await_suspend_promise(me),
           // yieldAwaiter.await_suspend pops a frame
           me.getAsyncFrame().getParentFrame()};
     }
-    return {coroutine_handle<AsyncGeneratorPromise>::from_promise(me), nullptr};
+    return std::nullopt;
   }
 
  private:

third-party/folly/src/folly/coro/Coroutine.h

@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include <optional>
 #include <type_traits>
 
 #if __has_include(<variant>)
@@ -298,9 +299,53 @@ inline bool detect_promise_return_object_eager_conversion() {
 template <typename>
 class ExtendedCoroutinePromiseCrtp;
 
+namespace detail {
+template <typename, typename, typename>
+class TaskPromiseWrapperBase;
+}
+
 // Extended version of coroutine_handle<void>
 // Assumes (and enforces) assumption that coroutine_handle is a pointer
 class ExtendedCoroutineHandle {
+ protected:
+  template <typename>
+  friend class ExtendedCoroutinePromiseCrtp;
+  template <typename, typename, typename>
+  friend class detail::TaskPromiseWrapperBase;
+  // This passkey aims to stop end users from calling `getPromiseBase`, which
+  // is an unsafe implementation detail, and to prevent overload ambiguity.
+  //
+  // It also doubles as the sigil for `use_extended_handle_concept`, another
+  // private detail.
+  class PrivateTag {
+   private:
+    friend ExtendedCoroutineHandle;
+    PrivateTag() = default;
+  };
+
+ private:
+  // SFINAE detection for the `use_extended_handle_concept` member type alias
+  // that classes implementing `getErrorHandle` must expose.  We don't want to
+  // use any kind of common base on `TaskWrapperPromise`, be it non-empty
+  // `PromiseBase`, or a dedicated empty tag, since either one would break
+  // empty-base optimization.
+
+  template <typename T>
+  using use_extended_handle_of_ = typename T::use_extended_handle_concept;
+
+  template <typename T, typename Void = void>
+  struct use_extended_handle {
+    static_assert(
+        require_sizeof<T>, "`use_extended_handle` on incomplete type");
+    static constexpr bool value = false;
+  };
+
+  template <typename T>
+  struct use_extended_handle<T, void_t<use_extended_handle_of_<T>>> {
+    static constexpr bool value =
+        std::is_same_v<use_extended_handle_of_<T>, PrivateTag>;
+  };
+
  public:
   using ErrorHandle = std::pair<ExtendedCoroutineHandle, AsyncStackFrame*>;
 
@@ -310,15 +355,17 @@ class ExtendedCoroutineHandle {
     template <typename>
     friend class ExtendedCoroutinePromiseCrtp;
 
-    using Fn = ErrorHandle(PromiseBase*, exception_wrapper& ex);
+    using Fn = std::optional<ErrorHandle>(PromiseBase*, exception_wrapper& ex);
 
     explicit PromiseBase(Fn* fn) : getErrorHandlePtr_(fn) {}
     ~PromiseBase() = default;
 
-    // A manual vtable with 1 function.  The benefit over virtual inheritance
-    // is that derived classes like `TaskPromise` don't have to be `final` in
-    // order to for the compiler to treat them as non-polymorphic.
-    // Specifically, this enables `SafeTask`, a more type-safe `Task`.
+    // A manual vtable with 1 function. Benefits over virtual inheritance:
+    //   - `TaskWrapperPromise` can implement `getErrorHandle` without bloating
+    //     itself with with a vtable it does not need.
+    //   - A tiny binary size win.
+    //   - Derived classes like `TaskPromise` don't have to be `final` in order
+    //     for the compiler to treat them as non-polymorphic.
     Fn* getErrorHandlePtr_;
   };
 
@@ -332,10 +379,10 @@ class ExtendedCoroutineHandle {
 
   template <
       typename Promise,
-      std::enable_if_t<std::is_base_of_v<PromiseBase, Promise>, int> = 0>
+      std::enable_if_t<use_extended_handle<Promise>::value, int> = 0>
   /*implicit*/ ExtendedCoroutineHandle(Promise* promise) noexcept
       : basic_(coroutine_handle<Promise>::from_promise(*promise)),
-        extended_(static_cast<PromiseBase*>(promise)) {}
+        extended_(Promise::getPromiseBase(PrivateTag{}, promise)) {}
 
   ExtendedCoroutineHandle() noexcept = default;
 
@@ -347,7 +394,9 @@ class ExtendedCoroutineHandle {
 
   ErrorHandle getErrorHandle(exception_wrapper& ex) {
     if (extended_) {
-      return extended_->getErrorHandlePtr_(extended_, ex);
+      if (auto res = extended_->getErrorHandlePtr_(extended_, ex)) {
+        return *res;
+      }
     }
     return {basic_, nullptr};
   }
@@ -357,8 +406,8 @@ class ExtendedCoroutineHandle {
  private:
   template <typename Promise>
   static auto fromBasic(coroutine_handle<Promise> handle) noexcept {
-    if constexpr (std::is_convertible_v<Promise*, PromiseBase*>) {
-      return static_cast<PromiseBase*>(&handle.promise());
+    if constexpr (use_extended_handle<Promise>::value) {
+      return Promise::getPromiseBase(PrivateTag{}, &handle.promise());
     } else {
       return nullptr;
     }
@@ -368,24 +417,42 @@ class ExtendedCoroutineHandle {
   PromiseBase* extended_{nullptr};
 };
 
-// folly::coro types are expected to implement this extended promise interface:
-//   (1) Publicly inherit from `ExtendedCoroutinePromiseCrtp<YourPromise>`,
-//   (2) Implement this static method on `YourPromise`:
+// folly::coro types are expected to implement this extended promise interface.
+//
+// It allows types to provide a more efficient resumption path when they know
+// they will be receiving an error result from the awaitee.
+//
+// First, publicly inherit from `ExtendedCoroutinePromiseCrtp<YourPromise>`,
+// Second, implement this static method on `YourPromise`:
 //
-//  static ExtendedCoroutineHandle::ErrorHandle getErrorHandle(
-//      YourPromise&, exception_wrapper&)
+//   static std::optional<ExtendedCoroutineHandle::ErrorHandle>
+//   getErrorHandleImpl(YourPromise&, exception_wrapper&);
 //
-// Rationale: Types may provide a more efficient resumption path when they
-// know they will be receiving an error result from the awaitee.  If they
-// do, they might also update the active stack frame.
+// Return `std::nullopt` to avoid changing the resumption path.  Otherwise,
+// return the `ExtendedCoroutineHandle` to resume & the active stack frame.
+//
+// DANGER: `YourPromise& promise` is a promise instance, but it might NOT
+// directly correspond to a coro frame.  For example, if your coro is wrapped,
+// that promise is a **member** inside a larger wrapper promise for the coro.
+// Therefore, you must NOT call `coroutine_handle<...>::from_promise(promise)`.
+// In the future, the true handle could be supplied, but none of the current
+// coros required it.
 template <typename Promise>
 class ExtendedCoroutinePromiseCrtp
     : public ExtendedCoroutineHandle::PromiseBase {
+ public:
+  using use_extended_handle_concept = ExtendedCoroutineHandle::PrivateTag;
+
+  static ExtendedCoroutineHandle::PromiseBase* getPromiseBase(
+      ExtendedCoroutineHandle::PrivateTag, ExtendedCoroutinePromiseCrtp* me) {
+    return me;
+  }
+
  protected:
   using PromiseBase = typename ExtendedCoroutineHandle::PromiseBase;
   ExtendedCoroutinePromiseCrtp()
-      : PromiseBase(+[](PromiseBase* promise, exception_wrapper& ex) {
-          return Promise::getErrorHandle(*static_cast<Promise*>(promise), ex);
+      : PromiseBase(+[](PromiseBase* p, exception_wrapper& ex) {
+          return Promise::getErrorHandleImpl(*static_cast<Promise*>(p), ex);
         }) {}
   ~ExtendedCoroutinePromiseCrtp() = default;
 };

third-party/folly/src/folly/coro/Task.h

@@ -82,9 +82,7 @@ class TaskPromiseBase {
     bool await_ready() noexcept { return false; }
 
     template <typename Promise>
-    FOLLY_CORO_AWAIT_SUSPEND_NONTRIVIAL_ATTRIBUTES coroutine_handle<>
-    await_suspend(coroutine_handle<Promise> coro) noexcept {
-      auto& promise = coro.promise();
+    coroutine_handle<> await_suspend_promise(Promise& promise) noexcept {
       // If ScopeExitTask has been attached, then we expect that the
       // ScopeExitTask will handle the lifetime of the async stack. See
       // ScopeExitTaskPromise's FinalAwaiter for more details.
@@ -114,6 +112,12 @@ class TaskPromiseBase {
       return promise.continuationRef(privateTag()).getHandle();
     }
 
+    template <typename Promise>
+    FOLLY_CORO_AWAIT_SUSPEND_NONTRIVIAL_ATTRIBUTES coroutine_handle<>
+    await_suspend(coroutine_handle<Promise> coro) noexcept {
+      return await_suspend_promise(coro.promise());
+    }
+
     [[noreturn]] void await_resume() noexcept { folly::assume_unreachable(); }
   };
 
@@ -276,18 +280,17 @@ class TaskPromiseCrtpBase
     return do_safe_point(*this);
   }
 
-  static ExtendedCoroutineHandle::ErrorHandle getErrorHandle(
+  static std::optional<ExtendedCoroutineHandle::ErrorHandle> getErrorHandleImpl(
       Promise& me, exception_wrapper& ex) {
     if (me.bypassExceptionThrowing_ == BypassExceptionThrowing::ACTIVE) {
       auto finalAwaiter = me.yield_value(co_error(std::move(ex)));
       DCHECK(!finalAwaiter.await_ready());
-      return {
-          finalAwaiter.await_suspend(
-              coroutine_handle<Promise>::from_promise(me)),
+      return ExtendedCoroutineHandle::ErrorHandle{
+          finalAwaiter.await_suspend_promise(me),
           // finalAwaiter.await_suspend pops a frame
           me.getAsyncFrame().getParentFrame()};
     }
-    return {coroutine_handle<Promise>::from_promise(me), nullptr};
+    return std::nullopt;
   }
 
  protected:
@@ -1003,6 +1006,10 @@ Task<drop_unit_t<T>> makeResultTask(Try<T> t) {
 template <typename Promise, typename T>
 inline Task<T>
 detail::TaskPromiseCrtpBase<Promise, T>::get_return_object() noexcept {
+  // Watch out: When used with `TaskWrapper`, this relies on "practically safe"
+  // UB wherein this handle is only valid because `TaskPromise` and the true
+  // "wrapper promise" of the wrapper coro coincide in layout exactly.
+  // Documented in `TaskPromiseWrapperBase::is_promise_type_punning_safe`.
   return Task<T>{
       coroutine_handle<Promise>::from_promise(*static_cast<Promise*>(this))};
 }

third-party/folly/src/folly/coro/TaskWrapper.h

@@ -169,8 +169,10 @@ class TaskPromiseWrapperBase {
   using TaskWrapperInnerPromise = Promise;
 
   WrapperTask get_return_object() noexcept {
-    // NB: See the function doc.  It'd be nice to have the `static_assert` at
-    // class scope, but the type is still incomplete at that point.
+    // CRITICAL: This assert justifies why it is practically safe to rely on
+    // the `from_promise` UB in `TaskPromiseCrtpBase::get_return_object`.
+    //
+    // PS The assert isn't at class scope, since the type would be incomplete.
     static_assert(is_promise_type_punning_safe());
     return WrapperTask{promise_.get_return_object()};
   }
@@ -231,6 +233,9 @@ class TaskPromiseWrapper
  public:
   template <typename U = T> // see "`co_return` with implicit ctor" test
   auto return_value(U&& value) {
+    static_assert( // See `is_promise_type_punning_safe` for rationale
+        require_sizeof<TaskPromiseWrapper> ==
+        require_sizeof<TaskPromiseWrapperBase<T, WrapperTask, Promise>>);
     return this->promise_.return_value(std::forward<U>(value));
   }
 };
@@ -243,7 +248,12 @@ class TaskPromiseWrapper<void, WrapperTask, Promise>
   ~TaskPromiseWrapper() = default;
 
  public:
-  void return_void() noexcept { this->promise_.return_void(); }
+  void return_void() noexcept {
+    static_assert( // See `is_promise_type_punning_safe` for rationale
+        require_sizeof<TaskPromiseWrapper> ==
+        require_sizeof<TaskPromiseWrapperBase<void, WrapperTask, Promise>>);
+    this->promise_.return_void();
+  }
 };
 
 // Mixin for TaskWrapper.h configs for `Task` & `TaskWithExecutor` types

third-party/folly/src/folly/coro/ViaIfAsync.h

@@ -154,14 +154,14 @@ class ViaCoroutine {
 
     folly::AsyncStackFrame& getLeafFrame() noexcept { return leafFrame_; }
 
-    static ExtendedCoroutineHandle::ErrorHandle getErrorHandle(
-        promise_type& me, exception_wrapper& ex) {
+    static std::optional<ExtendedCoroutineHandle::ErrorHandle>
+    getErrorHandleImpl(promise_type& me, exception_wrapper& ex) {
       auto [handle, frame] = me.continuation_.getErrorHandle(ex);
       me.setContinuation(handle);
       if (frame && IsStackAware) {
         me.leafFrame_.setParentFrame(*frame);
       }
-      return {coroutine_handle<promise_type>::from_promise(me), nullptr};
+      return std::nullopt;
     }
   };