From 8680d6db1e45193bd3595f974fa2ef43bb85ca1e Mon Sep 17 00:00:00 2001
From: Florian Mayer <fmayer@google.com>
Date: Tue, 1 Feb 2022 16:44:23 -0800
Subject: [PATCH] [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   | 12 ++++++-
 llvm/test/CodeGen/AArch64/stack-tagging-setjmp.ll | 44 +++++++++++++++++++++++
 2 files changed, 55 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/AArch64/stack-tagging-setjmp.ll

diff --git a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
index 13ae502..31a00a4 100644
--- a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
+++ b/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 =
diff --git a/llvm/test/CodeGen/AArch64/stack-tagging-setjmp.ll b/llvm/test/CodeGen/AArch64/stack-tagging-setjmp.ll
new file mode 100644
index 0000000..e5a6e8f
--- /dev/null
+++ b/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)
+
-- 
2.7.4