[mte] work around lifetime issue with setjmp.

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 for equivalent fix to HWASan.

Reviewed By: eugenis

Differential Revision: https://reviews.llvm.org/D118749

llvm/lib/Target/AArch64/AArch64StackTagging.cpp

@@ -536,8 +536,14 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
   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 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
     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 =

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"
+
+@stackbuf = dso_local local_unnamed_addr global i8* null, align 8
+@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)
+