[ASAN] Properly deal with musttail calls in ASAN

[lxfind]

Summary

When address sanitizing a function, stack unpinsoning code is inserted before each ret instruction. However if the ret instruciton is preceded by a musttail call, such transformation broke the musttail call contract and generates invalid IR.
This patch fixes the issue by moving the insertion point prior to the musttail call if there is one.

[lxfind]

comments

[lxfind]

add test

[lxfind]

fix test

[wenlei]

LGTM except a minor comment.

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

562

Is there a better place for this helper and avoid duplication? We have a few copies of it now, TSAN, ASAN and EntryExitInstrumenter...

[lxfind]

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

562

Yeah I have been wondering about this too. There are more duplicates than TSAN and ASAN.
Do you have any suggestions? I couldn't find a good place for this.

[wenlei]

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

562

Was thinking about something like building this into SetInsertPoint of IRBuilder, perhaps with an extra boolean argument with default value. But doesn't look like good layering there - might be too specific for IRBuilder. I'm fine with the way it is now.

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

 static uint64_t GetCtorAndDtorPriority(Triple &TargetTriple) {
   if (TargetTriple.isOSEmscripten()) {
     return kAsanEmscriptenCtorAndDtorPriority;
   } else {
     return kAsanCtorAndDtorPriority;
   }
 }
 
+// For a ret instruction followed by a musttail call, we cannot insert anything
+// in between. Instead we use the musttail call instruction as the insertion
+// point.
+static Instruction *adjustForMusttailCall(Instruction *I) {
+  ReturnInst *RI = dyn_cast<ReturnInst>(I);
+  if (!RI)
+    return I;
+  Instruction *Prev = RI->getPrevNode();
+  if (BitCastInst *BCI = dyn_cast_or_null<BitCastInst>(Prev))
+    Prev = BCI->getPrevNode();
+  if (CallInst *CI = dyn_cast_or_null<CallInst>(Prev))
+    if (CI->isMustTailCall())
+      return CI;
+  return RI;
+}
+
 namespace {
 
 /// Module analysis for getting various metadata about the module.
 class ASanGlobalsMetadataWrapperPass : public ModulePass {
 public:
   static char ID;
 
   ASanGlobalsMetadataWrapperPass() : ModulePass(ID) {
void unpoisonDynamicAllocasBeforeInst(Instruction *InstBefore,
 
     IRB.CreateCall(
         AsanAllocasUnpoisonFunc,
         {IRB.CreateLoad(IntptrTy, DynamicAllocaLayout), DynamicAreaPtr});
   }
 
   // Unpoison dynamic allocas redzones.
   void unpoisonDynamicAllocas() {
-    for (auto &Ret : RetVec)
-      unpoisonDynamicAllocasBeforeInst(Ret, DynamicAllocaLayout);
+    for (Instruction *Ret : RetVec)
+      unpoisonDynamicAllocasBeforeInst(adjustForMusttailCall(Ret),
+                                       DynamicAllocaLayout);
 
-    for (auto &StackRestoreInst : StackRestoreVec)
+    for (Instruction *StackRestoreInst : StackRestoreVec)
       unpoisonDynamicAllocasBeforeInst(StackRestoreInst,
                                        StackRestoreInst->getOperand(0));
   }
 
   // Deploy and poison redzones around dynamic alloca call. To do this, we
   // should replace this call with another one with changed parameters and
   // replace all its uses with new address, so
   //   addr = alloca type, old_size, align
for (const auto &APC : StaticAllocaPoisonCallVec) {
                    IRB, ShadowBase);
     }
   }
 
   SmallVector<uint8_t, 64> ShadowClean(ShadowAfterScope.size(), 0);
   SmallVector<uint8_t, 64> ShadowAfterReturn;
 
   // (Un)poison the stack before all ret instructions.
-  for (auto Ret : RetVec) {
-    IRBuilder<> IRBRet(Ret);
+  for (Instruction *Ret : RetVec) {
+    Instruction *Adjusted = adjustForMusttailCall(Ret);
+    IRBuilder<> IRBRet(Adjusted);
     // Mark the current frame as retired.
     IRBRet.CreateStore(ConstantInt::get(IntptrTy, kRetiredStackFrameMagic),
                        BasePlus0);
     if (DoStackMalloc) {
       assert(StackMallocIdx >= 0);
       // if FakeStack != 0  // LocalStackBase == FakeStack
       //     // In use-after-return mode, poison the whole stack frame.
       //     if StackMallocIdx <= 4
       //         // For small sizes inline the whole thing:
       //         memset(ShadowBase, kAsanStackAfterReturnMagic, ShadowSize);
       //         **SavedFlagPtr(FakeStack) = 0
       //     else
       //         __asan_stack_free_N(FakeStack, LocalStackSize)
       // else
       //     <This is not a fake stack; unpoison the redzones>
       Value *Cmp =
           IRBRet.CreateICmpNE(FakeStack, Constant::getNullValue(IntptrTy));
       Instruction *ThenTerm, *ElseTerm;
-      SplitBlockAndInsertIfThenElse(Cmp, Ret, &ThenTerm, &ElseTerm);
+      SplitBlockAndInsertIfThenElse(Cmp, Adjusted, &ThenTerm, &ElseTerm);
 
       IRBuilder<> IRBPoison(ThenTerm);
       if (StackMallocIdx <= 4) {
         int ClassSize = kMinStackMallocSize << StackMallocIdx;
         ShadowAfterReturn.resize(ClassSize / L.Granularity,
                                  kAsanStackUseAfterReturnMagic);
         copyToShadow(ShadowAfterReturn, ShadowAfterReturn, IRBPoison,
                      ShadowBase);

llvm/test/Instrumentation/AddressSanitizer/musttail.ll

+; To test that asan does not break the musttail call contract.
+;
+; RUN: opt < %s -asan -asan-module -S -enable-new-pm=0 | FileCheck %s
+; RUN: opt < %s -passes='asan-pipeline' -S | FileCheck %s
+
+define internal i32 @foo(i32* %p) sanitize_address {
+  %rv = load i32, i32* %p
+  ret i32 %rv
+}
+
+declare void @alloca_test_use([10 x i8]*)
+define i32 @call_foo(i32* %a) sanitize_address {
+  %x = alloca [10 x i8], align 1
+  call void @alloca_test_use([10 x i8]* %x)
+  %r = musttail call i32 @foo(i32* %a)
+  ret i32 %r
+}
+
+; CHECK-LABEL:  define i32 @call_foo(i32* %a) 
+; CHECK:          %r = musttail call i32 @foo(i32* %a)
+; CHECK-NEXT:     ret i32 %r
+
+
+define i32 @call_foo_cast(i32* %a) sanitize_address {
+  %x = alloca [10 x i8], align 1
+  call void @alloca_test_use([10 x i8]* %x)
+  %r = musttail call i32 @foo(i32* %a)
+  %t = bitcast i32 %r to i32
+  ret i32 %t
+}
+
+; CHECK-LABEL:  define i32 @call_foo_cast(i32* %a)
+; CHECK:          %r = musttail call i32 @foo(i32* %a)
+; CHECK-NEXT:     %t = bitcast i32 %r to i32
+; CHECK-NEXT:     ret i32 %t