[clang] [coroutines] Detect lifetime issues with coroutine lambda captures (PR #77066)

Fri Jan 5 02:23:34 PST 2024

llvmbot wrote:




@llvm/pr-subscribers-clang

Author: Utkarsh Saxena (usx95)

<details>
<summary>Changes</summary>

### Problem

```cpp
co_task<int> coro() {
    int a = 1;
    auto lamb = [a]() -> co_task<int> {
        co_return a; // 'a' in the lambda object dies after the iniital_suspend in the lambda coroutine.
    }();
    co_return co_await lamb;
}
```
[use-after-free](https://godbolt.org/z/GWPEovWWc)

Lambda captures (even by value) are prone to use-after-free once the lambda object dies. In the above example, the lambda object appears only as a temporary in the call expression. It dies after the first suspension (`initial_suspend`) in the lambda.
On resumption in `co_await lamb`, the lambda accesses `a` which is part of the already-dead lambda object.

---

### Solution

This problem can be formulated by saying that the `this` parameter of the lambda call operator is a lifetimebound parameter. The lambda object argument should therefore live atleast as long as the return object.
That said, this requirement does not hold if the lambda does not have a capture list. In principle, the coroutine frame still has a reference to a dead lambda object, but it is easy to see that the object would not be used in the lambda-coroutine body due to no capture list.

Using this pattern inside a`co_await` expression due to the lifetime extension of temporaries. Example:

```cpp
co_task<int> coro() {
    int a = 1;
    int res = co_await [a]() -> co_task<int> { co_return a; }();
    co_return res;
}
```
---
### Background

This came up in the discussion with seastar folks on [RFC](https://discourse.llvm.org/t/rfc-lifetime-bound-check-for-parameters-of-coroutines/74253/19?u=usx95). This is a fairly common pattern in continuation-style-passing (CSP) async programming involving futures and continuations. Document ["Lambda coroutine fiasco"](https://github.com/scylladb/seastar/blob/master/doc/lambda-coroutine-fiasco.md) by Seastar captures the problem.
This pattern makes the migration from CSP-style async programming to coroutines very bugprone.


Fixes https://github.com/llvm/llvm-project/issues/76995 

---
Full diff: https://github.com/llvm/llvm-project/pull/77066.diff


2 Files Affected:

- (modified) clang/lib/Sema/SemaInit.cpp (+17-3) 
- (modified) clang/test/SemaCXX/coro-lifetimebound.cpp (+59-5) 


``````````diff

diff --git a/clang/lib/Sema/SemaInit.cpp b/clang/lib/Sema/SemaInit.cpp
index 60c0e3e74204ec..c100bf11454786 100644
--- a/clang/lib/Sema/SemaInit.cpp
+++ b/clang/lib/Sema/SemaInit.cpp
@@ -12,6 +12,7 @@
 
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/DeclObjC.h"
+#include "clang/AST/Expr.h"
 #include "clang/AST/ExprCXX.h"
 #include "clang/AST/ExprObjC.h"
 #include "clang/AST/ExprOpenMP.h"
@@ -33,6 +34,7 @@
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/Support/Casting.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/raw_ostream.h"
 
@@ -7575,15 +7577,27 @@ static void visitLifetimeBoundArguments(IndirectLocalPath &Path, Expr *Call,
     Path.pop_back();
   };
 
-  if (ObjectArg && implicitObjectParamIsLifetimeBound(Callee))
-    VisitLifetimeBoundArg(Callee, ObjectArg);
-
   bool CheckCoroCall = false;
   if (const auto *RD = Callee->getReturnType()->getAsRecordDecl()) {
     CheckCoroCall = RD->hasAttr<CoroLifetimeBoundAttr>() &&
                     RD->hasAttr<CoroReturnTypeAttr>() &&
                     !Callee->hasAttr<CoroDisableLifetimeBoundAttr>();
   }
+
+  if (ObjectArg) {
+    bool CheckCoroObjArg = CheckCoroCall;
+    // Ignore `__promise.get_return_object()` as it not lifetimebound.
+    if (Callee->getDeclName().isIdentifier() &&
+        Callee->getName() == "get_return_object")
+      CheckCoroObjArg = false;
+    // Coroutine lambda objects with empty capture list are not lifetimebound.
+    if (auto *LE = dyn_cast<LambdaExpr>(ObjectArg->IgnoreImplicit());
+        LE && LE->captures().empty())
+      CheckCoroObjArg = false;
+    if (implicitObjectParamIsLifetimeBound(Callee) || CheckCoroObjArg)
+      VisitLifetimeBoundArg(Callee, ObjectArg);
+  }
+
   for (unsigned I = 0,
                 N = std::min<unsigned>(Callee->getNumParams(), Args.size());
        I != N; ++I) {
diff --git a/clang/test/SemaCXX/coro-lifetimebound.cpp b/clang/test/SemaCXX/coro-lifetimebound.cpp
index 3fc7ca70a14a12..319134450e4b6f 100644
--- a/clang/test/SemaCXX/coro-lifetimebound.cpp
+++ b/clang/test/SemaCXX/coro-lifetimebound.cpp
@@ -64,6 +64,10 @@ Co<int> bar_coro(const int &b, int c) {
       : bar_coro(0, 1); // expected-warning {{returning address of local temporary object}}
 }
 
+// =============================================================================
+// Lambdas
+// =============================================================================
+namespace lambdas {
 void lambdas() {
   auto unsafe_lambda = [] [[clang::coro_wrapper]] (int b) {
     return foo_coro(b); // expected-warning {{address of stack memory associated with parameter}}
@@ -84,15 +88,47 @@ void lambdas() {
     co_return x + co_await foo_coro(b);
   };
 }
+
+Co<int> lambda_captures() {
+  int a = 1;
+  // Temporary lambda object dies.
+  auto lamb = [a](int x, const int& y) -> Co<int> { // expected-warning {{temporary whose address is used as value of local variable 'lamb'}}
+    co_return x + y + a;
+  }(1, a);
+  // Object dies but it has no capture.
+  auto no_capture = []() -> Co<int> { co_return 1; }();
+  auto bad_no_capture = [](const int& a) -> Co<int> { co_return a; }(1); // expected-warning {{temporary}}
+  // Temporary lambda object with lifetime extension under co_await.
+  int res = co_await [a](int x, const int& y) -> Co<int> {
+    co_return x + y + a;
+  }(1, a);
+  co_return 1;
+}
+} // namespace lambdas
+
 // =============================================================================
-// Safe usage when parameters are value
+// Member coroutines
 // =============================================================================
-namespace by_value {
-Co<int> value_coro(int b) { co_return co_await foo_coro(b); }
-[[clang::coro_wrapper]] Co<int> wrapper1(int b) { return value_coro(b); }
-[[clang::coro_wrapper]] Co<int> wrapper2(const int& b) { return value_coro(b); }
+namespace member_coroutines{
+struct S {
+  Co<int> member(const int& a) { co_return a; }  
+};
+
+Co<int> use() {
+  S s;
+  int a = 1;
+  auto test1 = s.member(1);  // expected-warning {{temporary whose address is used as value of local variable}}
+  auto test2 = s.member(a);
+  auto test3 = S{}.member(a);  // expected-warning {{temporary whose address is used as value of local variable}}
+  co_return 1;
 }
 
+[[clang::coro_wrapper]] Co<int> wrapper(const int& a) {
+  S s;
+  return s.member(a); // expected-warning {{address of stack memory}}
+}
+} // member_coroutines
+
 // =============================================================================
 // Lifetime bound but not a Coroutine Return Type: No analysis.
 // =============================================================================
@@ -129,4 +165,22 @@ Co<int> foo_wrapper(const int& x) { return foo(x); }
   // The call to foo_wrapper is wrapper is safe.
   return foo_wrapper(1);
 }
+
+struct S{
+[[clang::coro_wrapper, clang::coro_disable_lifetimebound]] 
+Co<int> member(const int& x) { return foo(x); }
+};
+
+Co<int> use() {
+  S s;
+  int a = 1;
+  auto test1 = s.member(1); // param is not flagged.
+  auto test2 = S{}.member(a); // 'this' is not flagged.
+  co_return 1;
+}
+
+[[clang::coro_wrapper]] Co<int> return_stack_addr(const int& a) {
+  S s;
+  return s.member(a); // return of stack addr is not flagged.
+}
 } // namespace disable_lifetimebound

``````````

</details>


https://github.com/llvm/llvm-project/pull/77066
