[llvm] 1145374 - [ASAN] Properly deal with musttail calls in ASAN

[Xun Li via llvm-commits]

Author: Xun Li
Date: 2020-09-18T23:10:34-07:00
New Revision: 11453740bc6fb7c9a5fba3e0761060fc2182dc14

URL: https://github.com/llvm/llvm-project/commit/11453740bc6fb7c9a5fba3e0761060fc2182dc14
DIFF: https://github.com/llvm/llvm-project/commit/11453740bc6fb7c9a5fba3e0761060fc2182dc14.diff

LOG: [ASAN] Properly deal with musttail calls in ASAN

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.

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

Added: 
    llvm/test/Instrumentation/AddressSanitizer/musttail.ll

Modified: 
    llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
index 5d52a4b0fc18..802f6759265c 100644
--- a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
@@ -556,6 +556,22 @@ static uint64_t GetCtorAndDtorPriority(Triple &TargetTriple) {
   }
 }
 
+// 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.
@@ -999,10 +1015,11 @@ struct FunctionStackPoisoner : public InstVisitor<FunctionStackPoisoner> {
 
   // 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));
   }
@@ -3303,8 +3320,9 @@ void FunctionStackPoisoner::processStaticAllocas() {
   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);
@@ -3323,7 +3341,7 @@ void FunctionStackPoisoner::processStaticAllocas() {
       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) {

diff  --git a/llvm/test/Instrumentation/AddressSanitizer/musttail.ll b/llvm/test/Instrumentation/AddressSanitizer/musttail.ll
new file mode 100644
index 000000000000..83294d3e4502
--- /dev/null
+++ b/llvm/test/Instrumentation/AddressSanitizer/musttail.ll
@@ -0,0 +1,35 @@
+; 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