DataFlowSanitizer: Instrumentation for memset.

Differential Revision: http://llvm-reviews.chandlerc.com/D1395

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@188412 91177308-0d34-0410-b5e6-96231b3b80d8

2 files changed, 33 insertions, 1 deletions

diff --git a/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp b/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp
index dd01d83..b645ab8 100644
--- a/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp
+++ b/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp
@@ -159,9 +159,11 @@ class DataFlowSanitizer : public ModulePass {
   FunctionType *DFSanUnionFnTy;
   FunctionType *DFSanUnionLoadFnTy;
   FunctionType *DFSanUnimplementedFnTy;
+  FunctionType *DFSanSetLabelFnTy;
   Constant *DFSanUnionFn;
   Constant *DFSanUnionLoadFn;
   Constant *DFSanUnimplementedFn;
+  Constant *DFSanSetLabelFn;
   MDNode *ColdCallWeights;
   OwningPtr<SpecialCaseList> ABIList;
   DenseMap<Value *, Function *> UnwrappedFnMap;
@@ -235,6 +237,7 @@ class DFSanVisitor : public InstVisitor<DFSanVisitor> {
   void visitInsertValueInst(InsertValueInst &I);
   void visitAllocaInst(AllocaInst &I);
   void visitSelectInst(SelectInst &I);
+  void visitMemSetInst(MemSetInst &I);
   void visitMemTransferInst(MemTransferInst &I);
 };
 
@@ -305,6 +308,9 @@ bool DataFlowSanitizer::doInitialization(Module &M) {
       FunctionType::get(ShadowTy, DFSanUnionLoadArgs, /*isVarArg=*/ false);
   DFSanUnimplementedFnTy = FunctionType::get(
       Type::getVoidTy(*Ctx), Type::getInt8PtrTy(*Ctx), /*isVarArg=*/false);
+  Type *DFSanSetLabelArgs[3] = { ShadowTy, Type::getInt8PtrTy(*Ctx), IntptrTy };
+  DFSanSetLabelFnTy = FunctionType::get(Type::getVoidTy(*Ctx),
+                                        DFSanSetLabelArgs, /*isVarArg=*/false);
 
   if (GetArgTLSPtr) {
     Type *ArgTLSTy = ArrayType::get(ShadowTy, 64);
@@ -378,6 +384,11 @@ bool DataFlowSanitizer::runOnModule(Module &M) {
   }
   DFSanUnimplementedFn =
       Mod->getOrInsertFunction("__dfsan_unimplemented", DFSanUnimplementedFnTy);
+  DFSanSetLabelFn =
+      Mod->getOrInsertFunction("__dfsan_set_label", DFSanSetLabelFnTy);
+  if (Function *F = dyn_cast<Function>(DFSanSetLabelFn)) {
+    F->addAttribute(1, Attribute::ZExt);
+  }
 
   std::vector<Function *> FnsToInstrument;
   llvm::SmallPtrSet<Function *, 2> FnsWithNativeABI;
@@ -385,7 +396,8 @@ bool DataFlowSanitizer::runOnModule(Module &M) {
     if (!i->isIntrinsic() &&
         i != DFSanUnionFn &&
         i != DFSanUnionLoadFn &&
-        i != DFSanUnimplementedFn)
+        i != DFSanUnimplementedFn &&
+        i != DFSanSetLabelFn)
       FnsToInstrument.push_back(&*i);
   }
 
@@ -947,6 +959,15 @@ void DFSanVisitor::visitSelectInst(SelectInst &I) {
   }
 }
 
+void DFSanVisitor::visitMemSetInst(MemSetInst &I) {
+  IRBuilder<> IRB(&I);
+  Value *ValShadow = DFSF.getShadow(I.getValue());
+  IRB.CreateCall3(
+      DFSF.DFS.DFSanSetLabelFn, ValShadow,
+      IRB.CreateBitCast(I.getDest(), Type::getInt8PtrTy(*DFSF.DFS.Ctx)),
+      IRB.CreateZExtOrTrunc(I.getLength(), DFSF.DFS.IntptrTy));
+}
+
 void DFSanVisitor::visitMemTransferInst(MemTransferInst &I) {
   IRBuilder<> IRB(&I);
   Value *DestShadow = DFSF.DFS.getShadowAddress(I.getDest(), &I);
diff --git a/test/Instrumentation/DataFlowSanitizer/memset.ll b/test/Instrumentation/DataFlowSanitizer/memset.ll
new file mode 100644
index 0000000..2cc25db
--- /dev/null
+++ b/test/Instrumentation/DataFlowSanitizer/memset.ll
@@ -0,0 +1,11 @@
+; RUN: opt < %s -dfsan -dfsan-args-abi -S | FileCheck %s
+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"
+
+declare void @llvm.memset.p0i8.i64(i8* nocapture, i8, i64, i32, i1)
+
+define void @ms(i8* %p, i8 %v) {
+  ; CHECK-LABEL: @ms(i8*, i8, i16, i16)
+  ; CHECK: call void @__dfsan_set_label(i16 %3, i8* %0, i64 1)
+  call void @llvm.memset.p0i8.i64(i8* %p, i8 %v, i64 1, i32 1, i1 1)
+  ret void
+}