[DFSan] Add __dfsan_load_callback.

Summary:
When -dfsan-event-callbacks is specified, insert a call to
__dfsan_load_callback() on every load.

Reviewers: vitalybuka, pcc, kcc

Reviewed By: vitalybuka, kcc

Subscribers: hiraditya, #sanitizers, llvm-commits, eugenis, kcc

Tags: #sanitizers, #llvm

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

Author: morehouse

compiler-rt/test/dfsan/event_callbacks.c

@@ -1,5 +1,5 @@
-// RUN: %clang_dfsan -fno-sanitize=dataflow -fPIE -DCALLBACKS -c %s -o %t-callbacks.o
-// RUN: %clang_dfsan -mllvm -dfsan-event-callbacks %s %t-callbacks.o -o %t
+// RUN: %clang_dfsan -fno-sanitize=dataflow -O2 -fPIE -DCALLBACKS -c %s -o %t-callbacks.o
+// RUN: %clang_dfsan -O2 -mllvm -dfsan-event-callbacks %s %t-callbacks.o -o %t
 // RUN: %run %t 2>&1 | FileCheck %s
 
 // Tests that callbacks are inserted for store events when
@@ -35,12 +35,16 @@ void __dfsan_store_callback(dfsan_label Label) {
     assert(0);
   }
 
-  // CHECK: Label 1 stored to memory
-  // CHECK: Label 2 stored to memory
-  // CHECK: Label 3 stored to memory
   fprintf(stderr, "Label %u stored to memory\n", Label);
 }
 
+void __dfsan_load_callback(dfsan_label Label) {
+  if (!Label)
+    return;
+
+  fprintf(stderr, "Label %u loaded from memory\n", Label);
+}
+
 #else
 // Compile this code with DFSan and -dfsan-event-callbacks to insert the
 // callbacks.
@@ -57,10 +61,25 @@ int main(void) {
   dfsan_set_label(LabelJ, &J, sizeof(J));
   LabelIJ = dfsan_union(LabelI, LabelJ);
 
+  // CHECK: Label 1 stored to memory
   volatile int Sink = I;
+
+  // CHECK: Label 1 loaded from memory
+  assert(Sink == 1);
+
+  // CHECK: Label 2 stored to memory
   Sink = J;
+
+  // CHECK: Label 2 loaded from memory
+  assert(Sink == 2);
+
+  // CHECK: Label 2 loaded from memory
+  // CHECK: Label 3 stored to memory
   Sink += I;
 
+  // CHECK: Label 3 loaded from memory
+  assert(Sink == 3);
+
   return 0;
 }
 

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

@@ -163,10 +163,11 @@ static cl::opt<bool> ClDebugNonzeroLabels(
     cl::Hidden);
 
 // Experimental feature that inserts callbacks for certain data events.
-// Currently callbacks are only inserted for stores.
+// Currently callbacks are only inserted for loads and stores.
 //
 // If this flag is set to true, the user must provide definitions for the
 // following callback functions:
+//   void __dfsan_load_callback(dfsan_label Label);
 //   void __dfsan_store_callback(dfsan_label Label);
 static cl::opt<bool> ClEventCallbacks(
     "dfsan-event-callbacks",
@@ -356,14 +357,15 @@ class DataFlowSanitizer : public ModulePass {
   FunctionType *DFSanSetLabelFnTy;
   FunctionType *DFSanNonzeroLabelFnTy;
   FunctionType *DFSanVarargWrapperFnTy;
-  FunctionType *DFSanStoreCallbackFnTy;
+  FunctionType *DFSanLoadStoreCallbackFnTy;
   FunctionCallee DFSanUnionFn;
   FunctionCallee DFSanCheckedUnionFn;
   FunctionCallee DFSanUnionLoadFn;
   FunctionCallee DFSanUnimplementedFn;
   FunctionCallee DFSanSetLabelFn;
   FunctionCallee DFSanNonzeroLabelFn;
   FunctionCallee DFSanVarargWrapperFn;
+  FunctionCallee DFSanLoadCallbackFn;
   FunctionCallee DFSanStoreCallbackFn;
   MDNode *ColdCallWeights;
   DFSanABIList ABIList;
@@ -596,7 +598,7 @@ bool DataFlowSanitizer::doInitialization(Module &M) {
       Type::getVoidTy(*Ctx), None, /*isVarArg=*/false);
   DFSanVarargWrapperFnTy = FunctionType::get(
       Type::getVoidTy(*Ctx), Type::getInt8PtrTy(*Ctx), /*isVarArg=*/false);
-  DFSanStoreCallbackFnTy =
+  DFSanLoadStoreCallbackFnTy =
       FunctionType::get(Type::getVoidTy(*Ctx), ShadowTy, /*isVarArg=*/false);
 
   if (GetArgTLSPtr) {
@@ -798,8 +800,10 @@ bool DataFlowSanitizer::runOnModule(Module &M) {
   DFSanVarargWrapperFn = Mod->getOrInsertFunction("__dfsan_vararg_wrapper",
                                                   DFSanVarargWrapperFnTy);
 
+  DFSanLoadCallbackFn = Mod->getOrInsertFunction("__dfsan_load_callback",
+                                                 DFSanLoadStoreCallbackFnTy);
   DFSanStoreCallbackFn = Mod->getOrInsertFunction("__dfsan_store_callback",
-                                                  DFSanStoreCallbackFnTy);
+                                                  DFSanLoadStoreCallbackFnTy);
 
   std::vector<Function *> FnsToInstrument;
   SmallPtrSet<Function *, 2> FnsWithNativeABI;
@@ -812,6 +816,7 @@ bool DataFlowSanitizer::runOnModule(Module &M) {
         &i != DFSanSetLabelFn.getCallee()->stripPointerCasts() &&
         &i != DFSanNonzeroLabelFn.getCallee()->stripPointerCasts() &&
         &i != DFSanVarargWrapperFn.getCallee()->stripPointerCasts() &&
+        &i != DFSanLoadCallbackFn.getCallee()->stripPointerCasts() &&
         &i != DFSanStoreCallbackFn.getCallee()->stripPointerCasts())
       FnsToInstrument.push_back(&i);
   }
@@ -1344,6 +1349,10 @@ void DFSanVisitor::visitLoadInst(LoadInst &LI) {
     DFSF.NonZeroChecks.push_back(Shadow);
 
   DFSF.setShadow(&LI, Shadow);
+  if (ClEventCallbacks) {
+    IRBuilder<> IRB(&LI);
+    IRB.CreateCall(DFSF.DFS.DFSanLoadCallbackFn, Shadow);
+  }
 }
 
 void DFSanFunction::storeShadow(Value *Addr, uint64_t Size, Align Alignment,