Fix dfsan handling of musttail calls.

[browneee]

Summary

Without this change, a callsite like:

[[clang::musttail]] return func_call(x);

will cause an error like:

fatal error: error in backend: failed to perform tail call elimination
on a call site marked musttail

due to DFSan inserting instrumentation between the musttail call and
the return.

[stephan.yichao.zhao]

llvm/test/Instrumentation/DataFlowSanitizer/musttailcall.ll

5

@gbalats: FYI for removing the non-fast-8-test from the new test.

9–13

Please remove the pattern check because this test is not related.

[browneee]

Fix review comments.

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

const DataLayout &getDataLayout() const {
     return DFSF.F->getParent()->getDataLayout();
   }
 
   // Combines shadow values and origins for all of I's operands.
   void visitInstOperands(Instruction &I);
 
   void visitUnaryOperator(UnaryOperator &UO);
   void visitBinaryOperator(BinaryOperator &BO);
+  void visitBitCastInst(BitCastInst &BCI);
   void visitCastInst(CastInst &CI);
   void visitCmpInst(CmpInst &CI);
   void visitGetElementPtrInst(GetElementPtrInst &GEPI);
   void visitLoadInst(LoadInst &LI);
   void visitStoreInst(StoreInst &SI);
   void visitAtomicRMWInst(AtomicRMWInst &I);
   void visitAtomicCmpXchgInst(AtomicCmpXchgInst &I);
   void visitReturnInst(ReturnInst &RI);

 void DFSanVisitor::visitUnaryOperator(UnaryOperator &UO) {
   visitInstOperands(UO);
 }
 
 void DFSanVisitor::visitBinaryOperator(BinaryOperator &BO) {
   visitInstOperands(BO);
 }
 
+void DFSanVisitor::visitBitCastInst(BitCastInst &BCI) {
+  if (DFSF.DFS.getInstrumentedABI() == DataFlowSanitizer::IA_TLS) {
+    // Special case: if this is the bitcast (there is exactly 1 allowed) between
+    // a musttail call and a ret, don't instrument. New instructions are not
+    // allowed after a musttail call.
+    if (auto *CI = dyn_cast<CallInst>(BCI.getOperand(0)))
+      if (CI->isMustTailCall())
+        return;
+  }
+  // TODO: handle musttail call returns for IA_Args.
+  visitInstOperands(BCI);
+}
+
 void DFSanVisitor::visitCastInst(CastInst &CI) { visitInstOperands(CI); }
 
 void DFSanVisitor::visitCmpInst(CmpInst &CI) {
   visitInstOperands(CI);
   if (ClEventCallbacks) {
     IRBuilder<> IRB(&CI);
     Value *CombinedShadow = DFSF.getShadow(&CI);
     IRB.CreateCall(DFSF.DFS.DFSanCmpCallbackFn, CombinedShadow);
void DFSanVisitor::visitMemTransferInst(MemTransferInst &I) {
   }
   if (ClEventCallbacks) {
     IRB.CreateCall(DFSF.DFS.DFSanMemTransferCallbackFn,
                    {RawDestShadow,
                     IRB.CreateZExtOrTrunc(I.getLength(), DFSF.DFS.IntptrTy)});
   }
 }
 
+static bool isAMustTailRetVal(Value *RetVal) {
+  // Tail call may have a bitcast between return.
+  if (auto *I = dyn_cast<BitCastInst>(RetVal)) {
+    RetVal = I->getOperand(0);
+  }
+  if (auto *I = dyn_cast<CallInst>(RetVal)) {
+    return I->isMustTailCall();
+  }
+  return false;
+}
+
 void DFSanVisitor::visitReturnInst(ReturnInst &RI) {
   if (!DFSF.IsNativeABI && RI.getReturnValue()) {
     switch (DFSF.IA) {
     case DataFlowSanitizer::IA_TLS: {
+      // Don't emit the instrumentation for musttail call returns.
+      if (isAMustTailRetVal(RI.getReturnValue()))
+        return;
+
       Value *S = DFSF.getShadow(RI.getReturnValue());
       IRBuilder<> IRB(&RI);
       Type *RT = DFSF.F->getFunctionType()->getReturnType();
       unsigned Size =
           getDataLayout().getTypeAllocSize(DFSF.DFS.getShadowTy(RT));
       if (Size <= RetvalTLSSize) {
         // If the size overflows, stores nothing. At callsite, oversized return
         // shadows are set to zero.
         IRB.CreateAlignedStore(S, DFSF.getRetvalTLS(RT, IRB),
                                ShadowTLSAlignment);
       }
       if (DFSF.DFS.shouldTrackOrigins()) {
         Value *O = DFSF.getOrigin(RI.getReturnValue());
         IRB.CreateStore(O, DFSF.getRetvalOriginTLS());
       }
       break;
     }
     case DataFlowSanitizer::IA_Args: {
+      // TODO: handle musttail call returns for IA_Args.
+
       IRBuilder<> IRB(&RI);
       Type *RT = DFSF.F->getFunctionType()->getReturnType();
       Value *InsVal =
           IRB.CreateInsertValue(UndefValue::get(RT), RI.getReturnValue(), 0);
       Value *InsShadow =
           IRB.CreateInsertValue(InsVal, DFSF.getShadow(RI.getReturnValue()), 1);
       RI.setOperand(0, InsShadow);
       break;
if (InvokeInst *II = dyn_cast<InvokeInst>(&CB)) {
         Next = &NewBB->front();
       }
     } else {
       assert(CB.getIterator() != CB.getParent()->end());
       Next = CB.getNextNode();
     }
 
     if (DFSF.DFS.getInstrumentedABI() == DataFlowSanitizer::IA_TLS) {
+      // Don't emit the epilogue for musttail call returns.
+      if (isa<CallInst>(CB) && cast<CallInst>(CB).isMustTailCall())
+        return;
+
       // Loads the return value shadow.
       IRBuilder<> NextIRB(Next);
       const DataLayout &DL = getDataLayout();
       unsigned Size = DL.getTypeAllocSize(DFSF.DFS.getShadowTy(&CB));
       if (Size > RetvalTLSSize) {
         // Set overflowed return shadow to be zero.
         DFSF.setShadow(&CB, DFSF.DFS.getZeroShadow(&CB));
       } else {
if (DFSF.DFS.getInstrumentedABI() == DataFlowSanitizer::IA_TLS) {
         DFSF.setOrigin(&CB, LI);
       }
     }
   }
 
   // Do all instrumentation for IA_Args down here to defer tampering with the
   // CFG in a way that SplitEdge may be able to detect.
   if (DFSF.DFS.getInstrumentedABI() == DataFlowSanitizer::IA_Args) {
+    // TODO: handle musttail call returns for IA_Args.
+
     FunctionType *NewFT = DFSF.DFS.getArgsFunctionType(FT);
     Value *Func =
         IRB.CreateBitCast(CB.getCalledOperand(), PointerType::getUnqual(NewFT));
 
     const unsigned NumParams = FT->getNumParams();
 
     // Copy original arguments.
     auto *ArgIt = CB.arg_begin(), *ArgEnd = CB.arg_end();

llvm/test/Instrumentation/DataFlowSanitizer/musttailcall.ll

+; RUN: opt < %s -dfsan -S | FileCheck %s
+; RUN: opt < %s -dfsan -dfsan-fast-16-labels -S | FileCheck %s
+; RUN: opt < %s -dfsan -dfsan-fast-8-labels -S | FileCheck %s
+; RUN: opt < %s -dfsan -dfsan-fast-16-labels -dfsan-track-origins=1 -S | FileCheck %s --check-prefixes=CHECK,CHECK_ORIGIN
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+declare i32 @f(i32)
+
+; CHECK-LABEL: @"dfs$inner_callee"
+define i32 @inner_callee(i32) {
+  %r = call i32 @f(i32 %0)
+
+  ; COMM: Store here will be loaded in @outer_caller
+  ; CHECK: store{{.*}}__dfsan_retval_tls
+  ; CHECK_ORIGIN-NEXT: store{{.*}}__dfsan_retval_origin_tls
+  ; CHECK-NEXT: ret i32
+  ret i32 %r
+}
+
+; CHECK-LABEL: @"dfs$musttail_call"
+define i32 @musttail_call(i32) {
+  ; CHECK: store{{.*}}__dfsan_arg_tls
+  ; CHECK-NEXT: musttail call i32 @"dfs$inner_callee"
+  %r = musttail call i32 @inner_callee(i32 %0)
+
+  ; For "musttail" calls we can not insert any shadow manipulating code between
+  ; call and the return instruction. And we don't need to, because everything is
+  ; taken care of in the callee.
+  ; This is similar to the function above, but the last load and store of
+  ; __dfsan_retval_tls can be elided because we know about the musttail.
+
+  ; CHECK-NEXT: ret i32
+  ret i32 %r
+}
+
+; CHECK-LABEL: @"dfs$outer_caller"
+define i32 @outer_caller() {
+  ; CHECK: call{{.*}}@"dfs$musttail_call"
+  ; CHECK-NEXT: load{{.*}}__dfsan_retval_tls
+  ; CHECK_ORIGIN-NEXT: load{{.*}}__dfsan_retval_origin_tls
+  %r = call i32 @musttail_call(i32 0)
+
+  ; CHECK-NEXT: store{{.*}}__dfsan_retval_tls
+  ; CHECK_ORIGIN-NEXT: store{{.*}}__dfsan_retval_origin_tls
+  ; CHECK-NEXT: ret i32
+  ret i32 %r
+}
+
+declare i32* @mismatching_callee(i32)
+
+; CHECK-LABEL: define i8* @"dfs$mismatching_musttail_call"
+define i8* @mismatching_musttail_call(i32) {
+  %r = musttail call i32* @mismatching_callee(i32 %0)
+  ; CHECK: musttail call i32* @"dfs$mismatching_callee"
+  ; COMM: No instrumentation between call and ret.
+  ; CHECK-NEXT: bitcast i32* {{.*}} to i8*
+  %c = bitcast i32* %r to i8*
+  ; CHECK-NEXT: ret i8*
+  ret i8* %c
+}