[llvm] 64b448b - [dfsan] Refactor visitCallBase

Jianzhou Zhao via llvm-commits llvm-commits at lists.llvm.org
Mon Feb 8 11:56:27 PST 2021


Author: Jianzhou Zhao
Date: 2021-02-08T19:55:18Z
New Revision: 64b448b983b130bb59b1328473da0e9289d2e39d

URL: https://github.com/llvm/llvm-project/commit/64b448b983b130bb59b1328473da0e9289d2e39d
DIFF: https://github.com/llvm/llvm-project/commit/64b448b983b130bb59b1328473da0e9289d2e39d.diff

LOG: [dfsan] Refactor visitCallBase

To simplify the review of https://reviews.llvm.org/D95835.

Reviewed-by: morehouse

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

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp
index fe16e5b9555f..c1a9a32ba147 100644
--- a/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp
@@ -574,6 +574,10 @@ class DFSanVisitor : public InstVisitor<DFSanVisitor> {
   void visitSelectInst(SelectInst &I);
   void visitMemSetInst(MemSetInst &I);
   void visitMemTransferInst(MemTransferInst &I);
+
+private:
+  // Returns false when this is an invoke of a custom function.
+  bool visitWrappedCallBase(Function &F, CallBase &CB);
 };
 
 } // end anonymous namespace
@@ -1942,6 +1946,136 @@ void DFSanVisitor::visitReturnInst(ReturnInst &RI) {
   }
 }
 
+bool DFSanVisitor::visitWrappedCallBase(Function &F, CallBase &CB) {
+  IRBuilder<> IRB(&CB);
+  switch (DFSF.DFS.getWrapperKind(&F)) {
+  case DataFlowSanitizer::WK_Warning:
+    CB.setCalledFunction(&F);
+    IRB.CreateCall(DFSF.DFS.DFSanUnimplementedFn,
+                   IRB.CreateGlobalStringPtr(F.getName()));
+    DFSF.setShadow(&CB, DFSF.DFS.getZeroShadow(&CB));
+    return true;
+  case DataFlowSanitizer::WK_Discard:
+    CB.setCalledFunction(&F);
+    DFSF.setShadow(&CB, DFSF.DFS.getZeroShadow(&CB));
+    return true;
+  case DataFlowSanitizer::WK_Functional:
+    CB.setCalledFunction(&F);
+    visitOperandShadowInst(CB);
+    return true;
+  case DataFlowSanitizer::WK_Custom:
+    // Don't try to handle invokes of custom functions, it's too complicated.
+    // Instead, invoke the dfsw$ wrapper, which will in turn call the __dfsw_
+    // wrapper.
+    CallInst *CI = dyn_cast<CallInst>(&CB);
+    if (!CI)
+      return false;
+
+    FunctionType *FT = F.getFunctionType();
+    TransformedFunction CustomFn = DFSF.DFS.getCustomFunctionType(FT);
+    std::string CustomFName = "__dfsw_";
+    CustomFName += F.getName();
+    FunctionCallee CustomF = DFSF.DFS.Mod->getOrInsertFunction(
+        CustomFName, CustomFn.TransformedType);
+    if (Function *CustomFn = dyn_cast<Function>(CustomF.getCallee())) {
+      CustomFn->copyAttributesFrom(&F);
+
+      // Custom functions returning non-void will write to the return label.
+      if (!FT->getReturnType()->isVoidTy()) {
+        CustomFn->removeAttributes(AttributeList::FunctionIndex,
+                                   DFSF.DFS.ReadOnlyNoneAttrs);
+      }
+    }
+
+    std::vector<Value *> Args;
+
+    // Adds non-variable arguments.
+    auto *I = CB.arg_begin();
+    for (unsigned n = FT->getNumParams(); n != 0; ++I, --n) {
+      Type *T = (*I)->getType();
+      FunctionType *ParamFT;
+      if (isa<PointerType>(T) &&
+          (ParamFT = dyn_cast<FunctionType>(
+               cast<PointerType>(T)->getElementType()))) {
+        std::string TName = "dfst";
+        TName += utostr(FT->getNumParams() - n);
+        TName += "$";
+        TName += F.getName();
+        Constant *T = DFSF.DFS.getOrBuildTrampolineFunction(ParamFT, TName);
+        Args.push_back(T);
+        Args.push_back(
+            IRB.CreateBitCast(*I, Type::getInt8PtrTy(*DFSF.DFS.Ctx)));
+      } else {
+        Args.push_back(*I);
+      }
+    }
+
+    // Adds non-variable argument shadows.
+    I = CB.arg_begin();
+    const unsigned ShadowArgStart = Args.size();
+    for (unsigned N = FT->getNumParams(); N != 0; ++I, --N)
+      Args.push_back(DFSF.collapseToPrimitiveShadow(DFSF.getShadow(*I), &CB));
+
+    // Adds variable argument shadows.
+    if (FT->isVarArg()) {
+      auto *LabelVATy = ArrayType::get(DFSF.DFS.PrimitiveShadowTy,
+                                       CB.arg_size() - FT->getNumParams());
+      auto *LabelVAAlloca =
+          new AllocaInst(LabelVATy, getDataLayout().getAllocaAddrSpace(),
+                         "labelva", &DFSF.F->getEntryBlock().front());
+
+      for (unsigned N = 0; I != CB.arg_end(); ++I, ++N) {
+        auto *LabelVAPtr = IRB.CreateStructGEP(LabelVATy, LabelVAAlloca, N);
+        IRB.CreateStore(DFSF.collapseToPrimitiveShadow(DFSF.getShadow(*I), &CB),
+                        LabelVAPtr);
+      }
+
+      Args.push_back(IRB.CreateStructGEP(LabelVATy, LabelVAAlloca, 0));
+    }
+
+    // Adds the return value shadow.
+    if (!FT->getReturnType()->isVoidTy()) {
+      if (!DFSF.LabelReturnAlloca) {
+        DFSF.LabelReturnAlloca = new AllocaInst(
+            DFSF.DFS.PrimitiveShadowTy, getDataLayout().getAllocaAddrSpace(),
+            "labelreturn", &DFSF.F->getEntryBlock().front());
+      }
+      Args.push_back(DFSF.LabelReturnAlloca);
+    }
+
+    // Adds variable arguments.
+    append_range(Args, drop_begin(CB.args(), FT->getNumParams()));
+
+    CallInst *CustomCI = IRB.CreateCall(CustomF, Args);
+    CustomCI->setCallingConv(CI->getCallingConv());
+    CustomCI->setAttributes(TransformFunctionAttributes(
+        CustomFn, CI->getContext(), CI->getAttributes()));
+
+    // Update the parameter attributes of the custom call instruction to
+    // zero extend the shadow parameters. This is required for targets
+    // which consider PrimitiveShadowTy an illegal type.
+    for (unsigned N = 0; N < FT->getNumParams(); N++) {
+      const unsigned ArgNo = ShadowArgStart + N;
+      if (CustomCI->getArgOperand(ArgNo)->getType() ==
+          DFSF.DFS.PrimitiveShadowTy)
+        CustomCI->addParamAttr(ArgNo, Attribute::ZExt);
+    }
+
+    // Loads the return value shadow.
+    if (!FT->getReturnType()->isVoidTy()) {
+      LoadInst *LabelLoad =
+          IRB.CreateLoad(DFSF.DFS.PrimitiveShadowTy, DFSF.LabelReturnAlloca);
+      DFSF.setShadow(CustomCI, DFSF.expandFromPrimitiveShadow(
+                                   FT->getReturnType(), LabelLoad, &CB));
+    }
+
+    CI->replaceAllUsesWith(CustomCI);
+    CI->eraseFromParent();
+    return true;
+  }
+  return false;
+}
+
 void DFSanVisitor::visitCallBase(CallBase &CB) {
   Function *F = CB.getCalledFunction();
   if ((F && F->isIntrinsic()) || CB.isInlineAsm()) {
@@ -1954,137 +2088,17 @@ void DFSanVisitor::visitCallBase(CallBase &CB) {
   if (F == DFSF.DFS.DFSanVarargWrapperFn.getCallee()->stripPointerCasts())
     return;
 
-  IRBuilder<> IRB(&CB);
-
   DenseMap<Value *, Function *>::iterator i =
       DFSF.DFS.UnwrappedFnMap.find(CB.getCalledOperand());
-  if (i != DFSF.DFS.UnwrappedFnMap.end()) {
-    Function *F = i->second;
-    switch (DFSF.DFS.getWrapperKind(F)) {
-    case DataFlowSanitizer::WK_Warning:
-      CB.setCalledFunction(F);
-      IRB.CreateCall(DFSF.DFS.DFSanUnimplementedFn,
-                     IRB.CreateGlobalStringPtr(F->getName()));
-      DFSF.setShadow(&CB, DFSF.DFS.getZeroShadow(&CB));
-      return;
-    case DataFlowSanitizer::WK_Discard:
-      CB.setCalledFunction(F);
-      DFSF.setShadow(&CB, DFSF.DFS.getZeroShadow(&CB));
-      return;
-    case DataFlowSanitizer::WK_Functional:
-      CB.setCalledFunction(F);
-      visitOperandShadowInst(CB);
+  if (i != DFSF.DFS.UnwrappedFnMap.end())
+    if (visitWrappedCallBase(*i->second, CB))
       return;
-    case DataFlowSanitizer::WK_Custom:
-      // Don't try to handle invokes of custom functions, it's too complicated.
-      // Instead, invoke the dfsw$ wrapper, which will in turn call the __dfsw_
-      // wrapper.
-      if (CallInst *CI = dyn_cast<CallInst>(&CB)) {
-        FunctionType *FT = F->getFunctionType();
-        TransformedFunction CustomFn = DFSF.DFS.getCustomFunctionType(FT);
-        std::string CustomFName = "__dfsw_";
-        CustomFName += F->getName();
-        FunctionCallee CustomF = DFSF.DFS.Mod->getOrInsertFunction(
-            CustomFName, CustomFn.TransformedType);
-        if (Function *CustomFn = dyn_cast<Function>(CustomF.getCallee())) {
-          CustomFn->copyAttributesFrom(F);
-
-          // Custom functions returning non-void will write to the return label.
-          if (!FT->getReturnType()->isVoidTy()) {
-            CustomFn->removeAttributes(AttributeList::FunctionIndex,
-                                       DFSF.DFS.ReadOnlyNoneAttrs);
-          }
-        }
-
-        std::vector<Value *> Args;
-
-        auto i = CB.arg_begin();
-        for (unsigned n = FT->getNumParams(); n != 0; ++i, --n) {
-          Type *T = (*i)->getType();
-          FunctionType *ParamFT;
-          if (isa<PointerType>(T) &&
-              (ParamFT = dyn_cast<FunctionType>(
-                   cast<PointerType>(T)->getElementType()))) {
-            std::string TName = "dfst";
-            TName += utostr(FT->getNumParams() - n);
-            TName += "$";
-            TName += F->getName();
-            Constant *T = DFSF.DFS.getOrBuildTrampolineFunction(ParamFT, TName);
-            Args.push_back(T);
-            Args.push_back(
-                IRB.CreateBitCast(*i, Type::getInt8PtrTy(*DFSF.DFS.Ctx)));
-          } else {
-            Args.push_back(*i);
-          }
-        }
-
-        i = CB.arg_begin();
-        const unsigned ShadowArgStart = Args.size();
-        for (unsigned n = FT->getNumParams(); n != 0; ++i, --n)
-          Args.push_back(
-              DFSF.collapseToPrimitiveShadow(DFSF.getShadow(*i), &CB));
-
-        if (FT->isVarArg()) {
-          auto *LabelVATy = ArrayType::get(DFSF.DFS.PrimitiveShadowTy,
-                                           CB.arg_size() - FT->getNumParams());
-          auto *LabelVAAlloca = new AllocaInst(
-              LabelVATy, getDataLayout().getAllocaAddrSpace(),
-              "labelva", &DFSF.F->getEntryBlock().front());
-
-          for (unsigned n = 0; i != CB.arg_end(); ++i, ++n) {
-            auto LabelVAPtr = IRB.CreateStructGEP(LabelVATy, LabelVAAlloca, n);
-            IRB.CreateStore(
-                DFSF.collapseToPrimitiveShadow(DFSF.getShadow(*i), &CB),
-                LabelVAPtr);
-          }
 
-          Args.push_back(IRB.CreateStructGEP(LabelVATy, LabelVAAlloca, 0));
-        }
-
-        if (!FT->getReturnType()->isVoidTy()) {
-          if (!DFSF.LabelReturnAlloca) {
-            DFSF.LabelReturnAlloca =
-                new AllocaInst(DFSF.DFS.PrimitiveShadowTy,
-                               getDataLayout().getAllocaAddrSpace(),
-                               "labelreturn", &DFSF.F->getEntryBlock().front());
-          }
-          Args.push_back(DFSF.LabelReturnAlloca);
-        }
-
-        append_range(Args, drop_begin(CB.args(), FT->getNumParams()));
-
-        CallInst *CustomCI = IRB.CreateCall(CustomF, Args);
-        CustomCI->setCallingConv(CI->getCallingConv());
-        CustomCI->setAttributes(TransformFunctionAttributes(CustomFn,
-            CI->getContext(), CI->getAttributes()));
-
-        // Update the parameter attributes of the custom call instruction to
-        // zero extend the shadow parameters. This is required for targets
-        // which consider PrimitiveShadowTy an illegal type.
-        for (unsigned n = 0; n < FT->getNumParams(); n++) {
-          const unsigned ArgNo = ShadowArgStart + n;
-          if (CustomCI->getArgOperand(ArgNo)->getType() ==
-              DFSF.DFS.PrimitiveShadowTy)
-            CustomCI->addParamAttr(ArgNo, Attribute::ZExt);
-        }
-
-        if (!FT->getReturnType()->isVoidTy()) {
-          LoadInst *LabelLoad = IRB.CreateLoad(DFSF.DFS.PrimitiveShadowTy,
-                                               DFSF.LabelReturnAlloca);
-          DFSF.setShadow(CustomCI, DFSF.expandFromPrimitiveShadow(
-                                       FT->getReturnType(), LabelLoad, &CB));
-        }
-
-        CI->replaceAllUsesWith(CustomCI);
-        CI->eraseFromParent();
-        return;
-      }
-      break;
-    }
-  }
+  IRBuilder<> IRB(&CB);
 
   FunctionType *FT = CB.getFunctionType();
   if (DFSF.DFS.getInstrumentedABI() == DataFlowSanitizer::IA_TLS) {
+    // Stores argument shadows.
     unsigned ArgOffset = 0;
     const DataLayout &DL = getDataLayout();
     for (unsigned I = 0, N = FT->getNumParams(); I != N; ++I) {
@@ -2118,6 +2132,7 @@ void DFSanVisitor::visitCallBase(CallBase &CB) {
     }
 
     if (DFSF.DFS.getInstrumentedABI() == DataFlowSanitizer::IA_TLS) {
+      // Loads the return value shadow.
       IRBuilder<> NextIRB(Next);
       const DataLayout &DL = getDataLayout();
       unsigned Size = DL.getTypeAllocSize(DFSF.DFS.getShadowTy(&CB));


        


More information about the llvm-commits mailing list