[PATCH] D118749: [mte] work around lifetime issue with setjmp.

[Florian Mayer via Phabricator via llvm-commits]

fmayer created this revision.
Herald added a subscriber: hiraditya.
fmayer requested review of this revision.
Herald added a project: LLVM.
Herald added a subscriber: llvm-commits.

setjmp can return twice, but PostDominatorTree is unaware of this. as
such, it overestimates postdominance, leaving some cases where memory
does not get untagged on return. this causes false positives later in
the program execution.

this is a workaround for now, in the longer term PostDominatorTree
should be made aware of returns_twice, as this may cause problems
elsewhere.

See D118647 <https://reviews.llvm.org/D118647> for equivalent fix to HWASan.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D118749

Files:
  llvm/lib/Target/AArch64/AArch64StackTagging.cpp
  llvm/test/CodeGen/AArch64/stack-tagging-setjmp.ll


Index: llvm/test/CodeGen/AArch64/stack-tagging-setjmp.ll
===================================================================
--- /dev/null
+++ llvm/test/CodeGen/AArch64/stack-tagging-setjmp.ll
@@ -0,0 +1,44 @@
+; RUN: opt -S -aarch64-stack-tagging %s -o - | FileCheck %s
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-unknown-linux-android29"
+
+ at stackbuf = dso_local local_unnamed_addr global i8* null, align 8
+ at jbuf = dso_local global [32 x i64] zeroinitializer, align 8
+
+declare void @may_jump()
+
+define dso_local noundef i1 @_Z6targetv() sanitize_memtag {
+entry:
+  %buf = alloca [4096 x i8], align 1
+  %call = call i32 @setjmp(i64* noundef getelementptr inbounds ([32 x i64], [32 x i64]* @jbuf, i64 0, i64 0))
+  switch i32 %call, label %while.body [
+    i32 1, label %return
+    i32 2, label %sw.bb1
+  ]
+
+sw.bb1:                                           ; preds = %entry
+  br label %return
+
+while.body:                                       ; preds = %entry
+  %0 = getelementptr inbounds [4096 x i8], [4096 x i8]* %buf, i64 0, i64 0
+  call void @llvm.lifetime.start.p0i8(i64 4096, i8* nonnull %0) #10
+  store i8* %0, i8** @stackbuf, align 8
+  ; may_jump may call longjmp, going back to the switch (and then the return),
+  ; bypassing the lifetime.end. This is why we need to untag on the return,
+  ; rather than the lifetime.end.
+  call void @may_jump()
+  call void @llvm.lifetime.end.p0i8(i64 4096, i8* nonnull %0) #10
+  br label %return
+
+; CHECK-LABEL: return:
+; CHECK: call void @llvm.aarch64.settag
+return:                                           ; preds = %entry, %while.body, %sw.bb1
+  %retval.0 = phi i1 [ true, %while.body ], [ true, %sw.bb1 ], [ false, %entry ]
+  ret i1 %retval.0
+}
+
+declare i32 @setjmp(i64* noundef) returns_twice
+
+declare void @llvm.lifetime.start.p0i8(i64 immarg, i8* nocapture)
+declare void @llvm.lifetime.end.p0i8(i64 immarg, i8* nocapture)
+
Index: llvm/lib/Target/AArch64/AArch64StackTagging.cpp
===================================================================
--- llvm/lib/Target/AArch64/AArch64StackTagging.cpp
+++ llvm/lib/Target/AArch64/AArch64StackTagging.cpp
@@ -536,8 +536,14 @@
   SmallVector<Instruction *, 8> RetVec;
   SmallVector<Instruction *, 4> UnrecognizedLifetimes;
 
+  bool CallsReturnTwice = false;
   for (auto &BB : *F) {
     for (Instruction &I : BB) {
+      if (CallInst *CI = dyn_cast<CallInst>(&I)) {
+        if (CI->canReturnTwice()) {
+          CallsReturnTwice = true;
+        }
+      }
       if (auto *AI = dyn_cast<AllocaInst>(&I)) {
         Allocas[AI].AI = AI;
         Allocas[AI].OldAI = AI;
@@ -639,8 +645,12 @@
     Info.AI->replaceAllUsesWith(TagPCall);
     TagPCall->setOperand(0, Info.AI);
 
+    // Calls to functions that may return twice (e.g. setjmp) confuse the
+    // postdominator analysis, and will leave us to keep memory tagged after
+    // function return. Work around this by always untagging at every return
+    // statement if return_twice functions are called.
     if (UnrecognizedLifetimes.empty() && Info.LifetimeStart.size() == 1 &&
-        Info.LifetimeEnd.size() == 1) {
+        Info.LifetimeEnd.size() == 1 && !CallsReturnTwice) {
       IntrinsicInst *Start = Info.LifetimeStart[0];
       IntrinsicInst *End = Info.LifetimeEnd[0];
       uint64_t Size =


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D118749.405131.patch
Type: text/x-patch
Size: 3367 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20220202/06e09bae/attachment.bin>