[clang] [Coroutines] Mark parameter allocas with coro.outside.frame metadata (PR #127653)

[Hans Wennborg via cfe-commits]

https://github.com/zmodem created https://github.com/llvm/llvm-project/pull/127653

Parameters to a coroutine get copied (moved) to coroutine-local instances which code inside the coroutine then uses.

The original parameters should not be part of the frame. Normally CoroSplit figures that out by itself, but for [[clang::trivial_abi]] parameters which, get destructed at the end of the ramp function, it does not (see bug), causing use-after-free's if the frame is destroyed before the end of the ramp (as happens if it doesn't suspend).

Since Clang knows these should never be part of the frame, use metadata to make it so.

Fixes #127499

>From cde82a27139c39406a9afb5b471fa527e52e5bca Mon Sep 17 00:00:00 2001
From: Hans Wennborg <hans at chromium.org>
Date: Tue, 18 Feb 2025 15:27:37 +0100
Subject: [PATCH] [Coroutines] Mark parameter allocas with coro.outside.frame
 metadata

Parameters to a coroutine get copied (moved) to coroutine-local
instances which code inside the coroutine then uses.

The original parameters should not be part of the frame. Normally
CoroSplit figures that out by itself, but for [[clang::trivial_abi]]
parameters which, get destructed at the end of the ramp function,
it does not (see bug), causing use-after-free's if the frame is
destroyed before the end of the ramp (as happens if it doesn't suspend).

Since Clang knows these should never be part of the frame, use metadata
to make it so.

Fixes #127499
---
 clang/lib/CodeGen/CGCoroutine.cpp            | 10 +++++
 clang/test/CodeGenCoroutines/coro-params.cpp | 39 +++++++++++++++-----
 2 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/clang/lib/CodeGen/CGCoroutine.cpp b/clang/lib/CodeGen/CGCoroutine.cpp
index 9abf2e8c9190d..cdc61676524b4 100644
--- a/clang/lib/CodeGen/CGCoroutine.cpp
+++ b/clang/lib/CodeGen/CGCoroutine.cpp
@@ -855,6 +855,16 @@ void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
     // Create parameter copies. We do it before creating a promise, since an
     // evolution of coroutine TS may allow promise constructor to observe
     // parameter copies.
+    for (const ParmVarDecl *Parm : FnArgs) {
+      // If the original param is in an alloca, exclude it from the coroutine
+      // frame. The parameter copy will be part of the frame.
+      Address ParmAddr = GetAddrOfLocalVar(Parm);
+      if (auto *ParmAlloca =
+              dyn_cast<llvm::AllocaInst>(ParmAddr.getBasePointer())) {
+        ParmAlloca->setMetadata(llvm::LLVMContext::MD_coro_outside_frame,
+                                llvm::MDNode::get(CGM.getLLVMContext(), {}));
+      }
+    }
     for (auto *PM : S.getParamMoves()) {
       EmitStmt(PM);
       ParamReplacer.addCopy(cast<DeclStmt>(PM));
diff --git a/clang/test/CodeGenCoroutines/coro-params.cpp b/clang/test/CodeGenCoroutines/coro-params.cpp
index b318f2f52ac09..9e640fb2c5c62 100644
--- a/clang/test/CodeGenCoroutines/coro-params.cpp
+++ b/clang/test/CodeGenCoroutines/coro-params.cpp
@@ -59,13 +59,22 @@ struct MoveAndCopy {
   ~MoveAndCopy();
 };
 
-void consume(int,int,int) noexcept;
+struct [[clang::trivial_abi]] TrivialABI {
+  int val;
+  TrivialABI(MoveAndCopy&&) noexcept;
+  ~TrivialABI();
+};
+
+void consume(int,int,int,int) noexcept;
 
 // TODO: Add support for CopyOnly params
-// CHECK: define{{.*}} void @_Z1fi8MoveOnly11MoveAndCopy(i32 noundef %val, ptr noundef %[[MoParam:.+]], ptr noundef %[[McParam:.+]]) #0 personality ptr @__gxx_personality_v0
-void f(int val, MoveOnly moParam, MoveAndCopy mcParam) {
+// CHECK: define{{.*}} void @_Z1fi8MoveOnly11MoveAndCopy10TrivialABI(i32 noundef %val, ptr noundef %[[MoParam:.+]], ptr noundef %[[McParam:.+]], i32 %[[TrivialParam:.+]]) #0 personality ptr @__gxx_personality_v0
+void f(int val, MoveOnly moParam, MoveAndCopy mcParam, TrivialABI trivialParam) {
+  // CHECK: %[[TrivialAlloca:.+]] = alloca %struct.TrivialABI,
+  // CHECK-SAME: !coro.outside.frame
   // CHECK: %[[MoCopy:.+]] = alloca %struct.MoveOnly,
   // CHECK: %[[McCopy:.+]] = alloca %struct.MoveAndCopy,
+  // CHECK: %[[TrivialCopy:.+]] = alloca %struct.TrivialABI,
   // CHECK: store i32 %val, ptr %[[ValAddr:.+]]
 
   // CHECK: call ptr @llvm.coro.begin(
@@ -73,25 +82,33 @@ void f(int val, MoveOnly moParam, MoveAndCopy mcParam) {
   // CHECK-NEXT: call void @llvm.lifetime.start.p0(
   // CHECK-NEXT: call void @_ZN11MoveAndCopyC1EOS_(ptr {{[^,]*}} %[[McCopy]], ptr noundef nonnull align 4 dereferenceable(4) %[[McParam]]) #
   // CHECK-NEXT: call void @llvm.lifetime.start.p0(
-  // CHECK-NEXT: invoke void @_ZNSt16coroutine_traitsIJvi8MoveOnly11MoveAndCopyEE12promise_typeC1Ev(
+  // CHECK-NEXT: call void @llvm.memcpy
+  // CHECK-SAME: %[[TrivialCopy]]
+  // CHECK-SAME: %[[TrivialAlloca]]
+  // CHECK-NEXT: call void @llvm.lifetime.start.p0(
+  // CHECK-NEXT: invoke void @_ZNSt16coroutine_traitsIJvi8MoveOnly11MoveAndCopy10TrivialABIEE12promise_typeC1Ev(
 
   // CHECK: call void @_ZN14suspend_always12await_resumeEv(
   // CHECK: %[[IntParam:.+]] = load i32, ptr %{{.*}}
   // CHECK: %[[MoGep:.+]] = getelementptr inbounds nuw %struct.MoveOnly, ptr %[[MoCopy]], i32 0, i32 0
   // CHECK: %[[MoVal:.+]] = load i32, ptr %[[MoGep]]
-  // CHECK: %[[McGep:.+]] =  getelementptr inbounds nuw %struct.MoveAndCopy, ptr %[[McCopy]], i32 0, i32 0
+  // CHECK: %[[McGep:.+]] = getelementptr inbounds nuw %struct.MoveAndCopy, ptr %[[McCopy]], i32 0, i32 0
   // CHECK: %[[McVal:.+]] = load i32, ptr %[[McGep]]
-  // CHECK: call void @_Z7consumeiii(i32 noundef %[[IntParam]], i32 noundef %[[MoVal]], i32 noundef %[[McVal]])
+  // CHECK: %[[TrivialGep:.+]] = getelementptr inbounds nuw %struct.TrivialABI, ptr %[[TrivialCopy]], i32 0, i32 0
+  // CHECK: %[[TrivialVal:.+]] = load i32, ptr %[[TrivialGep]]
+  // CHECK: call void @_Z7consumeiiii(i32 noundef %[[IntParam]], i32 noundef %[[MoVal]], i32 noundef %[[McVal]], i32 noundef %[[TrivialVal]])
 
-  consume(val, moParam.val, mcParam.val);
+  consume(val, moParam.val, mcParam.val, trivialParam.val);
   co_return;
 
   // Skip to final suspend:
-  // CHECK: call void @_ZNSt16coroutine_traitsIJvi8MoveOnly11MoveAndCopyEE12promise_type13final_suspendEv(
+  // CHECK: call void @_ZNSt16coroutine_traitsIJvi8MoveOnly11MoveAndCopy10TrivialABIEE12promise_type13final_suspendEv(
   // CHECK: call void @_ZN14suspend_always12await_resumeEv(
 
   // Destroy promise, then parameter copies:
-  // CHECK: call void @_ZNSt16coroutine_traitsIJvi8MoveOnly11MoveAndCopyEE12promise_typeD1Ev(ptr {{[^,]*}} %__promise)
+  // CHECK: call void @_ZNSt16coroutine_traitsIJvi8MoveOnly11MoveAndCopy10TrivialABIEE12promise_typeD1Ev(ptr {{[^,]*}} %__promise)
+  // CHECK-NEXT: call void @llvm.lifetime.end.p0(
+  // CHECK-NEXT: call void @_ZN10TrivialABID1Ev(ptr {{[^,]*}} %[[TrivialCopy]])
   // CHECK-NEXT: call void @llvm.lifetime.end.p0(
   // CHECK-NEXT: call void @_ZN11MoveAndCopyD1Ev(ptr {{[^,]*}} %[[McCopy]])
   // CHECK-NEXT: call void @llvm.lifetime.end.p0(
@@ -99,6 +116,10 @@ void f(int val, MoveOnly moParam, MoveAndCopy mcParam) {
   // CHECK-NEXT: call void @llvm.lifetime.end.p0(
   // CHECK-NEXT: call void @llvm.lifetime.end.p0(
   // CHECK-NEXT: call ptr @llvm.coro.free(
+
+  // The original trivial_abi parameter is destroyed when returning from the ramp.
+  // CHECK: call i1 @llvm.coro.end
+  // CHECK: call void @_ZN10TrivialABID1Ev(ptr {{[^,]*}} %[[TrivialAlloca]])
 }
 
 // CHECK-LABEL: void @_Z16dependent_paramsI1A1BEvT_T0_S3_(ptr noundef %x, ptr noundef %0, ptr noundef %y)