[hwasan] work around lifetime issue with setjmp.

[fmayer]

Summary

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

this is a crude workaround to unblock use-after-scope for now, in the
longer term PostDominatorTree should bemade aware of returns_twice
function, as this may cause problems elsewhere.

[fmayer]

missing newline at eof

[fmayer]

comment

[fmayer]

even more comments

[fmayer]

cosmetic changes

[fmayer]

more cosmetic changes

[eugenis]

This is a great find. @pcc
It seems like it would affect AArch64StackTagging, and some of the sanitizers as well (ASan in particular).
Also I don't see anything in CodeGen/StackColoring.cpp to avoid this case - is it possible for it to cause erroneous stack slot reuse?

llvm/test/Instrumentation/HWAddressSanitizer/use-after-scope-setjmp.ll

27

so the callee of this may longjmp, bypassing the lifetime.end. Effectively, this adds a DT edge from any call site in this function to immediately after the setjmp call.

Please add some comments here explaining this.

[eugenis]

llvm/test/Instrumentation/HWAddressSanitizer/use-after-scope-setjmp.ll

27

I meant to say "call graph edge".

[fmayer]

cosmetic changes

[fmayer]

In D118647#3285567, @eugenis wrote:

This is a great find. @pcc
It seems like it would affect AArch64StackTagging, and some of the sanitizers as well (ASan in particular).
Also I don't see anything in CodeGen/StackColoring.cpp to avoid this case - is it possible for it to cause erroneous stack slot reuse?

yep, as explained in the commit message i sent this CL to unblock the use-after-scope rollout while we figure out how to deal with the potentially bigger problem.

[fmayer]

stylistic change

[eugenis]

Did you forget to include a file?

[eugenis]

Oh I get it you are shadowing a class member. Let's not do that, it's confusing.

[fmayer]

naming

[fmayer]

In D118647#3288135, @eugenis wrote:

Oh I get it you are shadowing a class member. Let's not do that, it's confusing.

Fixed.

[eugenis]

LGTM

compiler-rt/test/hwasan/TestCases/use-after-scope-setjmp.cpp

5

I don't see any reason to limit the test to aarch64.

[fmayer]

compiler-rt/test/hwasan/TestCases/use-after-scope-setjmp.cpp

5

I'll leave this for now for consistency with the other hwasan compiler-rt tests. We can bulk remove them after.

compiler-rt/test/hwasan/TestCases/use-after-scope-setjmp.cpp

+// RUN: %clangxx_hwasan  -mllvm -hwasan-use-stack-safety=0 -mllvm -hwasan-use-after-scope -O2 %s -o %t && \
+// RUN:     %run %t 2>&1
+
+// REQUIRES: aarch64-target-arch
+// REQUIRES: stable-runtime
+
+#include <sanitizer/hwasan_interface.h>
+#include <setjmp.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include <sys/types.h>
+#include <unistd.h>
+
+volatile const char *stackbuf = nullptr;
+jmp_buf jbuf;
+
+__attribute__((noinline)) bool jump() {
+  // Fool the compiler so it cannot deduce noreturn.
+  if (getpid() != 0) {
+    longjmp(jbuf, 1);
+    return true;
+  }
+  return false;
+}
+
+bool target() {
+  switch (setjmp(jbuf)) {
+  case 1:
+    return false;
+  default:
+    break;
+  }
+
+  while (true) {
+    char buf[4096];
+    stackbuf = buf;
+    if (!jump()) {
+      break;
+    };
+  }
+  return true;
+}
+
+int main() {
+  target();
+
+  void *untagged = __hwasan_tag_pointer(stackbuf, 0);
+  if (stackbuf == untagged) {
+    // The buffer wasn't tagged in the first place, so the test will not work
+    // as expected.
+    return 2;
+  }
+  if (__hwasan_test_shadow(untagged, 4096) != -1) {
+    return 1;
+  }
+
+  return 0;
+}

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

public:
 
   bool isInterestingAlloca(const AllocaInst &AI);
   void tagAlloca(IRBuilder<> &IRB, AllocaInst *AI, Value *Tag, size_t Size);
   Value *tagPointer(IRBuilder<> &IRB, Type *Ty, Value *PtrLong, Value *Tag);
   Value *untagPointer(IRBuilder<> &IRB, Value *PtrLong);
   static bool isStandardLifetime(const AllocaInfo &AllocaInfo,
                                  const DominatorTree &DT);
   bool instrumentStack(
+      bool ShouldDetectUseAfterScope,
       MapVector<AllocaInst *, AllocaInfo> &AllocasToInstrument,
       SmallVector<Instruction *, 4> &UnrecognizedLifetimes,
       DenseMap<AllocaInst *, std::vector<DbgVariableIntrinsic *>> &AllocaDbgMap,
       SmallVectorImpl<Instruction *> &RetVec, Value *StackTag,
       llvm::function_ref<const DominatorTree &()> GetDT,
       llvm::function_ref<const PostDominatorTree &()> GetPDT);
   Value *readRegister(IRBuilder<> &IRB, StringRef Name);
   bool instrumentLandingPads(SmallVectorImpl<Instruction *> &RetVec);
bool HWAddressSanitizer::isStandardLifetime(const AllocaInfo &AllocaInfo,
   // at most one of them is actually used for each execution of the function.
   return AllocaInfo.LifetimeStart.size() == 1 &&
          (AllocaInfo.LifetimeEnd.size() == 1 ||
           (AllocaInfo.LifetimeEnd.size() > 0 &&
            !maybeReachableFromEachOther(AllocaInfo.LifetimeEnd, DT)));
 }
 
 bool HWAddressSanitizer::instrumentStack(
+    bool ShouldDetectUseAfterScope,
     MapVector<AllocaInst *, AllocaInfo> &AllocasToInstrument,
     SmallVector<Instruction *, 4> &UnrecognizedLifetimes,
     DenseMap<AllocaInst *, std::vector<DbgVariableIntrinsic *>> &AllocaDbgMap,
     SmallVectorImpl<Instruction *> &RetVec, Value *StackTag,
     llvm::function_ref<const DominatorTree &()> GetDT,
     llvm::function_ref<const PostDominatorTree &()> GetPDT) {
   // Ideally, we want to calculate tagged stack base pointer, and rewrite all
   // alloca addresses using that. Unfortunately, offsets are not known yet
for (auto &KV : AllocasToInstrument) {
     size_t AlignedSize = alignTo(Size, Mapping.getObjectAlignment());
     auto TagEnd = [&](Instruction *Node) {
       IRB.SetInsertPoint(Node);
       Value *UARTag = getUARTag(IRB, StackTag);
       tagAlloca(IRB, AI, UARTag, AlignedSize);
     };
     bool StandardLifetime =
         UnrecognizedLifetimes.empty() && isStandardLifetime(Info, GetDT());
-    if (DetectUseAfterScope && StandardLifetime) {
+    if (ShouldDetectUseAfterScope && StandardLifetime) {
       IntrinsicInst *Start = Info.LifetimeStart[0];
       IRB.SetInsertPoint(Start->getNextNode());
       tagAlloca(IRB, AI, Tag, Size);
       if (!forAllReachableExits(GetDT(), GetPDT(), Start, Info.LifetimeEnd,
                                 RetVec, TagEnd)) {
         for (auto *End : Info.LifetimeEnd)
           End->eraseFromParent();
       }
bool HWAddressSanitizer::sanitizeFunction(
 
   SmallVector<InterestingMemoryOperand, 16> OperandsToInstrument;
   SmallVector<MemIntrinsic *, 16> IntrinToInstrument;
   MapVector<AllocaInst *, AllocaInfo> AllocasToInstrument;
   SmallVector<Instruction *, 8> RetVec;
   SmallVector<Instruction *, 8> LandingPadVec;
   SmallVector<Instruction *, 4> UnrecognizedLifetimes;
   DenseMap<AllocaInst *, std::vector<DbgVariableIntrinsic *>> AllocaDbgMap;
+  bool CallsReturnTwice = false;
   for (auto &BB : F) {
     for (auto &Inst : BB) {
+      if (CallInst *CI = dyn_cast<CallInst>(&Inst)) {
+        if (CI->canReturnTwice()) {
+          CallsReturnTwice = true;
+        }
+      }
       if (InstrumentStack) {
         if (AllocaInst *AI = dyn_cast<AllocaInst>(&Inst)) {
           if (isInterestingAlloca(*AI))
             AllocasToInstrument.insert({AI, {}});
           continue;
         }
         auto *II = dyn_cast<IntrinsicInst>(&Inst);
         if (II && (II->getIntrinsicID() == Intrinsic::lifetime_start ||
bool HWAddressSanitizer::sanitizeFunction(
   IRBuilder<> EntryIRB(InsertPt);
   emitPrologue(EntryIRB,
                /*WithFrameRecord*/ ClRecordStackHistory &&
                    Mapping.WithFrameRecord && !AllocasToInstrument.empty());
 
   if (!AllocasToInstrument.empty()) {
     Value *StackTag =
         ClGenerateTagsWithCalls ? nullptr : getStackBaseTag(EntryIRB);
-    instrumentStack(AllocasToInstrument, UnrecognizedLifetimes, AllocaDbgMap,
+    // 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.
+    instrumentStack(DetectUseAfterScope && !CallsReturnTwice,
+                    AllocasToInstrument, UnrecognizedLifetimes, AllocaDbgMap,
                     RetVec, StackTag, GetDT, GetPDT);
   }
   // Pad and align each of the allocas that we instrumented to stop small
   // uninteresting allocas from hiding in instrumented alloca's padding and so
   // that we have enough space to store real tags for short granules.
   DenseMap<AllocaInst *, AllocaInst *> AllocaToPaddedAllocaMap =
       padInterestingAllocas(AllocasToInstrument);
 

llvm/test/Instrumentation/HWAddressSanitizer/use-after-scope-setjmp.ll

+; RUN: opt -passes=hwasan -hwasan-use-stack-safety=0 -hwasan-use-after-scope -S < %s | 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_hwaddress {
+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: void @llvm.memset.p0i8.i64({{.*}}, i8 0, i64 256, i1 false)
+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)