[DFSan] Add callback that allows to track which function tainted data reaches

[browneee]

compiler-rt/test/dfsan/reaches_function.c

5

Can you make this test build and run a second time with origin tracking enabled?
e.g. https://github.com/llvm/llvm-project/blob/91b38c6aaddefabad2a4c959ea3865e356761ed5/compiler-rt/test/dfsan/flush.c#L3

When origin tracking is enabled (use #ifdef), you could also call dfsan_print_origin_id_trace and add filecheck expectations for that output.

21–22

It may be a better approach to:

• print the file,line,function here
• pipe the output into filecheck
• add filecheck expectations to match for the printed lines, put them at each function and use [[@LINE-1]]

The advantages of this would be:

• you can use https://llvm.org/docs/CommandGuide/FileCheck.html#filecheck-pseudo-numeric-variables instead of integer constants, easier to keep test up to date
• you would be testing for function + line pair together
• if it failed you could easily see the file,line,function values
• Can also check origin traces.

[clg]

addressing comments

[clg]

fix test

[browneee]

compiler-rt/test/dfsan/reaches_function.c

53

To clarify, this is triggered by the tainted value is returned from add?

Should we expect the location to be here-1 because it should be the line number of the containing function, or should the line number be down at 60?

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

1061–1065

Should this attempt to use debug information from Instruction &I (the instruction where this occurs), rather than just using the debug loc for the containing function?

[clg]

addressing comments

compiler-rt/test/dfsan/reaches_function.c

53

changed it, see other comment.

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

1061–1065

Sure, we might as well give more fine-grained information on where the access took place.

[browneee]

LGTM.

Test failures look unrelated, but rebasing might make it green again.

Let me know if you'd like me to commit it for you.

[clg]

That would be very much appreciated :)

compiler-rt/include/sanitizer/dfsan_interface.h

 
 /// Signature of the callback argument to dfsan_set_write_callback().
 typedef void (*dfsan_write_callback_t)(int fd, const void *buf, size_t count);
 
 /// Signature of the callback argument to dfsan_set_conditional_callback().
 typedef void (*dfsan_conditional_callback_t)(dfsan_label label,
                                              dfsan_origin origin);
 
+/// Signature of the callback argument to dfsan_set_reaches_function_callback().
+/// The description is intended to hold the name of the variable.
+typedef void (*dfsan_reaches_function_callback_t)(dfsan_label label,
+                                                  dfsan_origin origin,
+                                                  const char *file,
+                                                  unsigned int line,
+                                                  const char *function);
+
 /// Computes the union of \c l1 and \c l2, resulting in a union label.
 dfsan_label dfsan_union(dfsan_label l1, dfsan_label l2);
 
 /// Sets the label for each address in [addr,addr+size) to \c label.
 void dfsan_set_label(dfsan_label label, void *addr, size_t size);
 
 /// Sets the label for each address in [addr,addr+size) to the union of the
 /// current label for that address and \c label.

 
 /// Conditional expressions occur during signal handlers.
 /// Making callbacks that handle signals well is tricky, so when
 /// -dfsan-conditional-callbacks=true, conditional expressions used in signal
 /// handlers will add the labels they see into a global (bitwise-or together).
 /// This function returns all label bits seen in signal handler conditions.
 dfsan_label dfsan_get_labels_in_signal_conditional();
 
+/// Sets a callback to be invoked when tainted data reaches a function.
+/// This could occur at function entry, or at a load instruction.
+/// These callbacks will only be added if -dfsan-reaches-function-callbacks=1.
+void dfsan_set_reaches_function_callback(
+    dfsan_reaches_function_callback_t callback);
+
+/// Making callbacks that handle signals well is tricky, so when
+/// -dfsan-reaches-function-callbacks=true, functions reached in signal
+/// handlers will add the labels they see into a global (bitwise-or together).
+/// This function returns all label bits seen during signal handlers.
+dfsan_label dfsan_get_labels_in_signal_reaches_function();
+
 /// Interceptor hooks.
 /// Whenever a dfsan's custom function is called the corresponding
 /// hook is called it non-zero. The hooks should be defined by the user.
 /// The primary use case is taint-guided fuzzing, where the fuzzer
 /// needs to see the parameters of the function and the labels.
 /// FIXME: implement more hooks.
 void dfsan_weak_hook_memcmp(void *caller_pc, const void *s1, const void *s2,
                             size_t n, dfsan_label s1_label,

compiler-rt/lib/dfsan/dfsan.cpp

extern "C" SANITIZER_INTERFACE_ATTRIBUTE void dfsan_set_conditional_callback(
   __dfsan::conditional_callback = callback;
 }
 
 extern "C" SANITIZER_INTERFACE_ATTRIBUTE dfsan_label
 dfsan_get_labels_in_signal_conditional() {
   return __dfsan::labels_in_signal_conditional;
 }
 
+namespace __dfsan {
+
+typedef void (*dfsan_reaches_function_callback_t)(dfsan_label label,
+                                                  dfsan_origin origin,
+                                                  const char *file,
+                                                  unsigned int line,
+                                                  const char *function);
+static dfsan_reaches_function_callback_t reaches_function_callback = nullptr;
+static dfsan_label labels_in_signal_reaches_function = 0;
+
+static void ReachesFunctionCallback(dfsan_label label, dfsan_origin origin,
+                                    const char *file, unsigned int line,
+                                    const char *function) {
+  if (label == 0) {
+    return;
+  }
+  if (reaches_function_callback == nullptr) {
+    return;
+  }
+
+  // This initial ReachesFunctionCallback handler needs to be in here in dfsan
+  // runtime (rather than being an entirely user implemented hook) so that it
+  // has access to dfsan thread information.
+  DFsanThread *t = GetCurrentThread();
+  // A callback operation which does useful work (like record the flow) will
+  // likely be too long executed in a signal handler.
+  if (t && t->InSignalHandler()) {
+    // Record set of labels used in signal handler for completeness.
+    labels_in_signal_reaches_function |= label;
+    return;
+  }
+
+  reaches_function_callback(label, origin, file, line, function);
+}
+
+}  // namespace __dfsan
+
+extern "C" SANITIZER_INTERFACE_ATTRIBUTE void
+__dfsan_reaches_function_callback_origin(dfsan_label label, dfsan_origin origin,
+                                         const char *file, unsigned int line,
+                                         const char *function) {
+  __dfsan::ReachesFunctionCallback(label, origin, file, line, function);
+}
+
+extern "C" SANITIZER_INTERFACE_ATTRIBUTE void
+__dfsan_reaches_function_callback(dfsan_label label, const char *file,
+                                  unsigned int line, const char *function) {
+  __dfsan::ReachesFunctionCallback(label, 0, file, line, function);
+}
+
+extern "C" SANITIZER_INTERFACE_ATTRIBUTE void
+dfsan_set_reaches_function_callback(
+    __dfsan::dfsan_reaches_function_callback_t callback) {
+  __dfsan::reaches_function_callback = callback;
+}
+
+extern "C" SANITIZER_INTERFACE_ATTRIBUTE dfsan_label
+dfsan_get_labels_in_signal_reaches_function() {
+  return __dfsan::labels_in_signal_reaches_function;
+}
+
 class Decorator : public __sanitizer::SanitizerCommonDecorator {
  public:
   Decorator() : SanitizerCommonDecorator() {}
   const char *Origin() const { return Magenta(); }
 };
 
 namespace {
 
if (start >= maxVirtualAddress)
       continue;
 
     if (!MmapFixedSuperNoReserve(start, size, kMemoryLayout[i].name)) {
       Printf("FATAL: DataFlowSanitizer: failed to clear memory region\n");
       Die();
     }
   }
   __dfsan::labels_in_signal_conditional = 0;
+  __dfsan::labels_in_signal_reaches_function = 0;
 }
 
 // TODO: CheckMemoryLayoutSanity is based on msan.
 // Consider refactoring these into a shared implementation.
 static void CheckMemoryLayoutSanity() {
   uptr prev_end = 0;
   for (unsigned i = 0; i < kMemoryLayoutSize; ++i) {
     uptr start = kMemoryLayout[i].start;

compiler-rt/lib/dfsan/done_abilist.txt

 fun:dfsan_get_init_origin=uninstrumented
 fun:dfsan_get_init_origin=discard
 fun:dfsan_get_track_origins=uninstrumented
 fun:dfsan_get_track_origins=discard
 fun:dfsan_set_conditional_callback=uninstrumented
 fun:dfsan_set_conditional_callback=discard
 fun:dfsan_get_labels_in_signal_conditional=uninstrumented
 fun:dfsan_get_labels_in_signal_conditional=discard
+fun:dfsan_set_reaches_function_callback=uninstrumented
+fun:dfsan_set_reaches_function_callback=discard
+fun:dfsan_get_labels_in_signal_reaches_function=uninstrumented
+fun:dfsan_get_labels_in_signal_reaches_function=discard
+fun:dfsan_reaches_function_callback=uninstrumented
+fun:dfsan_reaches_function_callback=discard
 
 ###############################################################################
 # glibc
 ###############################################################################
 # Functions of memory allocators
 fun:__libc_memalign=discard
 fun:aligned_alloc=discard
 fun:calloc=discard

compiler-rt/test/dfsan/Inputs/flags_abilist.txt

 fun:f=uninstrumented
 
 fun:function_to_force_zero=force_zero_labels
 
 fun:main=uninstrumented
 fun:main=discard
 
 fun:dfsan_set_label=uninstrumented
 fun:dfsan_set_label=discard
 
 fun:my_dfsan_conditional_callback=uninstrumented
 fun:my_dfsan_conditional_callback=discard
 
 fun:dfsan_set_conditional_callback=uninstrumented
 fun:dfsan_set_conditional_callback=discard
+
+fun:my_dfsan_reaches_function_callback=uninstrumented
+fun:my_dfsan_reaches_function_callback=discard
+
+fun:dfsan_set_reaches_function_callback=uninstrumented
+fun:dfsan_set_reaches_function_callback=discard

compiler-rt/test/dfsan/reaches_function.c

+// RUN: %clang_dfsan -fno-sanitize=dataflow -O2 -fPIE -DCALLBACKS -c %s -o %t-callbacks.o
+// RUN: %clang_dfsan -gmlt -fsanitize-ignorelist=%S/Inputs/flags_abilist.txt -O2 -mllvm -dfsan-reaches-function-callbacks=1 %s %t-callbacks.o -o %t
+// RUN: %run %t 2>&1 | FileCheck %s
+
+// RUN: %clang_dfsan -fno-sanitize=dataflow -O2 -fPIE -DCALLBACKS -DORIGIN_TRACKING -c %s -o %t-callbacks.o
+// RUN: %clang_dfsan -gmlt -fsanitize-ignorelist=%S/Inputs/flags_abilist.txt -O2 -mllvm -dfsan-reaches-function-callbacks=1 -mllvm -dfsan-track-origins=2 %s %t-callbacks.o -o %t
+// RUN: %run %t 2>&1 | FileCheck --check-prefix=CHECK-ORIGIN-TRACKING %s
+
+// REQUIRES: x86_64-target-arch
+
+// Tests that callbacks are inserted for reached functions when
+// -dfsan-reaches-function-callbacks is specified.
+
+#include <assert.h>
+#include <stdio.h>
+#include <string.h>
+#include <sanitizer/dfsan_interface.h>
+
+#ifdef CALLBACKS
+// Compile this code without DFSan to avoid recursive instrumentation.
+
+void my_dfsan_reaches_function_callback(dfsan_label label, dfsan_origin origin,
+                                        const char *file, unsigned int line,
+                                        const char *function) {
+#ifdef ORIGIN_TRACKING
+  dfsan_print_origin_id_trace(origin);
+#else
+  printf("%s:%d %s\n", file, line, function);
+#endif
+}
+
+#else
+
+__attribute__((noinline)) uint64_t add(uint64_t *a, uint64_t *b) {
+
+  return *a + *b;
+  // CHECK: {{.*}}compiler-rt/test/dfsan/reaches_function.c:[[# @LINE - 1]] add.dfsan
+  // CHECK-ORIGIN-TRACKING: Origin value: 0x10000002, Taint value was stored to memory at
+  // CHECK-ORIGIN-TRACKING: #0 {{.*}} in add.dfsan {{.*}}compiler-rt/test/dfsan/reaches_function.c:[[# @LINE - 3]]:{{.*}}
+  // CHECK-ORIGIN-TRACKING: Origin value: 0x1, Taint value was created at
+  // CHECK-ORIGIN-TRACKING: #0 {{.*}} in main {{.*}}compiler-rt/test/dfsan/reaches_function.c:{{.*}}
+}
+
+extern void my_dfsan_reaches_function_callback(dfsan_label label,
+                                               dfsan_origin origin,
+                                               const char *file,
+                                               unsigned int line,
+                                               const char *function);
+
+int main(int argc, char *argv[]) {
+
+  dfsan_set_reaches_function_callback(my_dfsan_reaches_function_callback);
+
+  uint64_t a = 0;
+  uint64_t b = 0;
+
+  dfsan_set_label(8, &a, sizeof(a));
+  uint64_t c = add(&a, &b);
+  // CHECK: {{.*}}compiler-rt/test/dfsan/reaches_function.c:[[# @LINE - 1]] main
+  // CHECK-ORIGIN-TRACKING: Origin value: 0x10000002, Taint value was stored to memory at
+  // CHECK-ORIGIN-TRACKING: #0 {{.*}} in add.dfsan {{.*}}compiler-rt/test/dfsan/reaches_function.c:{{.*}}
+  // CHECK-ORIGIN-TRACKING: Origin value: 0x1, Taint value was created at
+  // CHECK-ORIGIN-TRACKING: #0 {{.*}} in main {{.*}}compiler-rt/test/dfsan/reaches_function.c:[[# @LINE - 6]]:{{.*}}
+  return c;
+}
+
+#endif // #ifdef CALLBACKS

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

 // Experimental feature that inserts callbacks for conditionals, including:
 // conditional branch, switch, select.
 // This must be true for dfsan_set_conditional_callback() to have effect.
 static cl::opt<bool> ClConditionalCallbacks(
     "dfsan-conditional-callbacks",
     cl::desc("Insert calls to callback functions on conditionals."), cl::Hidden,
     cl::init(false));
 
+// Experimental feature that inserts callbacks for data reaching a function,
+// either via function arguments and loads.
+// This must be true for dfsan_set_reaches_function_callback() to have effect.
+static cl::opt<bool> ClReachesFunctionCallbacks(
+    "dfsan-reaches-function-callbacks",
+    cl::desc("Insert calls to callback functions on data reaching a function."),
+    cl::Hidden, cl::init(false));
+
 // Controls whether the pass tracks the control flow of select instructions.
 static cl::opt<bool> ClTrackSelectControlFlow(
     "dfsan-track-select-control-flow",
     cl::desc("Propagate labels from condition values of select instructions "
              "to results."),
     cl::Hidden, cl::init(true));
 
 // TODO: This default value follows MSan. DFSan may use a different value.
class DataFlowSanitizer {
   FunctionType *DFSanLoadLabelAndOriginFnTy;
   FunctionType *DFSanUnimplementedFnTy;
   FunctionType *DFSanWrapperExternWeakNullFnTy;
   FunctionType *DFSanSetLabelFnTy;
   FunctionType *DFSanNonzeroLabelFnTy;
   FunctionType *DFSanVarargWrapperFnTy;
   FunctionType *DFSanConditionalCallbackFnTy;
   FunctionType *DFSanConditionalCallbackOriginFnTy;
+  FunctionType *DFSanReachesFunctionCallbackFnTy;
+  FunctionType *DFSanReachesFunctionCallbackOriginFnTy;
   FunctionType *DFSanCmpCallbackFnTy;
   FunctionType *DFSanLoadStoreCallbackFnTy;
   FunctionType *DFSanMemTransferCallbackFnTy;
   FunctionType *DFSanChainOriginFnTy;
   FunctionType *DFSanChainOriginIfTaintedFnTy;
   FunctionType *DFSanMemOriginTransferFnTy;
   FunctionType *DFSanMemShadowOriginTransferFnTy;
   FunctionType *DFSanMemShadowOriginConditionalExchangeFnTy;
   FunctionType *DFSanMaybeStoreOriginFnTy;
   FunctionCallee DFSanUnionLoadFn;
   FunctionCallee DFSanLoadLabelAndOriginFn;
   FunctionCallee DFSanUnimplementedFn;
   FunctionCallee DFSanWrapperExternWeakNullFn;
   FunctionCallee DFSanSetLabelFn;
   FunctionCallee DFSanNonzeroLabelFn;
   FunctionCallee DFSanVarargWrapperFn;
   FunctionCallee DFSanLoadCallbackFn;
   FunctionCallee DFSanStoreCallbackFn;
   FunctionCallee DFSanMemTransferCallbackFn;
   FunctionCallee DFSanConditionalCallbackFn;
   FunctionCallee DFSanConditionalCallbackOriginFn;
+  FunctionCallee DFSanReachesFunctionCallbackFn;
+  FunctionCallee DFSanReachesFunctionCallbackOriginFn;
   FunctionCallee DFSanCmpCallbackFn;
   FunctionCallee DFSanChainOriginFn;
   FunctionCallee DFSanChainOriginIfTaintedFn;
   FunctionCallee DFSanMemOriginTransferFn;
   FunctionCallee DFSanMemShadowOriginTransferFn;
   FunctionCallee DFSanMemShadowOriginConditionalExchangeFn;
   FunctionCallee DFSanMaybeStoreOriginFn;
   SmallPtrSet<Value *, 16> DFSanRuntimeFunctions;
void storeZeroPrimitiveShadow(Value *Addr, uint64_t Size, Align ShadowAlign,
                                 Instruction *Pos);
 
   Align getShadowAlign(Align InstAlignment);
 
   // If ClConditionalCallbacks is enabled, insert a callback after a given
   // branch instruction using the given conditional expression.
   void addConditionalCallbacksIfEnabled(Instruction &I, Value *Condition);
 
+  // If ClReachesFunctionCallbacks is enabled, insert a callback for each
+  // argument and load instruction.
+  void addReachesFunctionCallbacksIfEnabled(IRBuilder<> &IRB, Instruction &I,
+                                            Value *Data);
+
   bool isLookupTableConstant(Value *P);
 
 private:
   /// Collapses the shadow with aggregate type into a single primitive shadow
   /// value.
   template <class AggregateType>
   Value *collapseAggregateShadow(AggregateType *AT, Value *Shadow,
                                  IRBuilder<> &IRB);
if (DFS.shouldTrackOrigins()) {
     Value *CondOrigin = getOrigin(Condition);
     IRB.CreateCall(DFS.DFSanConditionalCallbackOriginFn,
                    {CondShadow, CondOrigin});
   } else {
     IRB.CreateCall(DFS.DFSanConditionalCallbackFn, {CondShadow});
   }
 }
 
+void DFSanFunction::addReachesFunctionCallbacksIfEnabled(IRBuilder<> &IRB,
+                                                         Instruction &I,
+                                                         Value *Data) {
+  if (!ClReachesFunctionCallbacks) {
+    return;
+  }
+  const DebugLoc &dbgloc = I.getDebugLoc();
+  Value *DataShadow = collapseToPrimitiveShadow(getShadow(Data), IRB);
+  ConstantInt *CILine;
+  llvm::Value *FilePathPtr;
+
+  if (dbgloc.get() == nullptr) {
+    CILine = llvm::ConstantInt::get(I.getContext(), llvm::APInt(32, 0, false));
+    FilePathPtr = IRB.CreateGlobalStringPtr(
+        I.getFunction()->getParent()->getSourceFileName());
+  } else {
+    CILine = llvm::ConstantInt::get(I.getContext(),
+                                    llvm::APInt(32, dbgloc.getLine(), false));
+    FilePathPtr =
+        IRB.CreateGlobalStringPtr(dbgloc->getFilename());
+  }
+
+  llvm::Value *FunctionNamePtr =
+      IRB.CreateGlobalStringPtr(I.getFunction()->getName());
+
+  CallInst *CB;
+  std::vector<Value *> args;
+
+  if (DFS.shouldTrackOrigins()) {
+    Value *DataOrigin = getOrigin(Data);
+    args = { DataShadow, DataOrigin, FilePathPtr, CILine, FunctionNamePtr };
+    CB = IRB.CreateCall(DFS.DFSanReachesFunctionCallbackOriginFn, args);
+  } else {
+    args = { DataShadow, FilePathPtr, CILine, FunctionNamePtr };
+    CB = IRB.CreateCall(DFS.DFSanReachesFunctionCallbackFn, args);
+  }
+  CB->setDebugLoc(dbgloc);
+}
+
 Type *DataFlowSanitizer::getShadowTy(Type *OrigTy) {
   if (!OrigTy->isSized())
     return PrimitiveShadowTy;
   if (isa<IntegerType>(OrigTy))
     return PrimitiveShadowTy;
   if (isa<VectorType>(OrigTy))
     return PrimitiveShadowTy;
   if (ArrayType *AT = dyn_cast<ArrayType>(OrigTy))
DFSanVarargWrapperFnTy = FunctionType::get(
       Type::getVoidTy(*Ctx), Type::getInt8PtrTy(*Ctx), /*isVarArg=*/false);
   DFSanConditionalCallbackFnTy =
       FunctionType::get(Type::getVoidTy(*Ctx), PrimitiveShadowTy,
                         /*isVarArg=*/false);
   Type *DFSanConditionalCallbackOriginArgs[2] = {PrimitiveShadowTy, OriginTy};
   DFSanConditionalCallbackOriginFnTy = FunctionType::get(
       Type::getVoidTy(*Ctx), DFSanConditionalCallbackOriginArgs,
       /*isVarArg=*/false);
+  Type *DFSanReachesFunctionCallbackArgs[4] = {PrimitiveShadowTy, Int8Ptr,
+                                               OriginTy, Int8Ptr};
+  DFSanReachesFunctionCallbackFnTy =
+      FunctionType::get(Type::getVoidTy(*Ctx), DFSanReachesFunctionCallbackArgs,
+                        /*isVarArg=*/false);
+  Type *DFSanReachesFunctionCallbackOriginArgs[5] = {
+      PrimitiveShadowTy, OriginTy, Int8Ptr, OriginTy, Int8Ptr};
+  DFSanReachesFunctionCallbackOriginFnTy = FunctionType::get(
+      Type::getVoidTy(*Ctx), DFSanReachesFunctionCallbackOriginArgs,
+      /*isVarArg=*/false);
   DFSanCmpCallbackFnTy =
       FunctionType::get(Type::getVoidTy(*Ctx), PrimitiveShadowTy,
                         /*isVarArg=*/false);
   DFSanChainOriginFnTy =
       FunctionType::get(OriginTy, OriginTy, /*isVarArg=*/false);
   Type *DFSanChainOriginIfTaintedArgs[2] = {PrimitiveShadowTy, OriginTy};
   DFSanChainOriginIfTaintedFnTy = FunctionType::get(
       OriginTy, DFSanChainOriginIfTaintedArgs, /*isVarArg=*/false);
DFSanRuntimeFunctions.insert(
       DFSanStoreCallbackFn.getCallee()->stripPointerCasts());
   DFSanRuntimeFunctions.insert(
       DFSanMemTransferCallbackFn.getCallee()->stripPointerCasts());
   DFSanRuntimeFunctions.insert(
       DFSanConditionalCallbackFn.getCallee()->stripPointerCasts());
   DFSanRuntimeFunctions.insert(
       DFSanConditionalCallbackOriginFn.getCallee()->stripPointerCasts());
   DFSanRuntimeFunctions.insert(
+      DFSanReachesFunctionCallbackFn.getCallee()->stripPointerCasts());
+  DFSanRuntimeFunctions.insert(
+      DFSanReachesFunctionCallbackOriginFn.getCallee()->stripPointerCasts());
+  DFSanRuntimeFunctions.insert(
       DFSanCmpCallbackFn.getCallee()->stripPointerCasts());
   DFSanRuntimeFunctions.insert(
       DFSanChainOriginFn.getCallee()->stripPointerCasts());
   DFSanRuntimeFunctions.insert(
       DFSanChainOriginIfTaintedFn.getCallee()->stripPointerCasts());
   DFSanRuntimeFunctions.insert(
       DFSanMemOriginTransferFn.getCallee()->stripPointerCasts());
   DFSanRuntimeFunctions.insert(
void DataFlowSanitizer::initializeCallbackFunctions(Module &M) {
   DFSanCmpCallbackFn =
       Mod->getOrInsertFunction("__dfsan_cmp_callback", DFSanCmpCallbackFnTy);
 
   DFSanConditionalCallbackFn = Mod->getOrInsertFunction(
       "__dfsan_conditional_callback", DFSanConditionalCallbackFnTy);
   DFSanConditionalCallbackOriginFn =
       Mod->getOrInsertFunction("__dfsan_conditional_callback_origin",
                                DFSanConditionalCallbackOriginFnTy);
+  DFSanReachesFunctionCallbackFn = Mod->getOrInsertFunction(
+      "__dfsan_reaches_function_callback", DFSanReachesFunctionCallbackFnTy);
+  DFSanReachesFunctionCallbackOriginFn =
+      Mod->getOrInsertFunction("__dfsan_reaches_function_callback_origin",
+                               DFSanReachesFunctionCallbackOriginFnTy);
 }
 
 void DataFlowSanitizer::injectMetadataGlobals(Module &M) {
   // These variables can be used:
   // - by the runtime (to discover what the shadow width was, during
   //   compilation)
   // - in testing (to avoid hardcoding the shadow width and type but instead
   //   extract them by pattern matching)
for (Function *F : FnsToInstrument) {
     if (!F || F->isDeclaration())
       continue;
 
     removeUnreachableBlocks(*F);
 
     DFSanFunction DFSF(*this, F, FnsWithNativeABI.count(F),
                        FnsWithForceZeroLabel.count(F), GetTLI(*F));
 
+    if (ClReachesFunctionCallbacks) {
+      // Add callback for arguments reaching this function.
+      for (auto &FArg : F->args()) {
+        Instruction *Next = &F->getEntryBlock().front();
+        Value *FArgShadow = DFSF.getShadow(&FArg);
+        if (isZeroShadow(FArgShadow))
+          continue;
+        if (Instruction *FArgShadowInst = dyn_cast<Instruction>(FArgShadow)) {
+          Next = FArgShadowInst->getNextNode();
+        }
+        if (shouldTrackOrigins()) {
+          if (Instruction *Origin =
+                  dyn_cast<Instruction>(DFSF.getOrigin(&FArg))) {
+            // Ensure IRB insertion point is after loads for shadow and origin.
+            Instruction *OriginNext = Origin->getNextNode();
+            if (Next->comesBefore(OriginNext)) {
+              Next = OriginNext;
+            }
+          }
+        }
+        IRBuilder<> IRB(Next);
+        DFSF.addReachesFunctionCallbacksIfEnabled(IRB, *Next, &FArg);
+      }
+    }
+
     // DFSanVisitor may create new basic blocks, which confuses df_iterator.
     // Build a copy of the list before iterating over it.
     SmallVector<BasicBlock *, 4> BBList(depth_first(&F->getEntryBlock()));
 
     for (BasicBlock *BB : BBList) {
       Instruction *Inst = &BB->front();
       while (true) {
         // DFSanVisitor may split the current basic block, changing the current
void DFSanVisitor::visitLoadInst(LoadInst &LI) {
   // When an application load is atomic, increase atomic ordering between
   // atomic application loads and stores to ensure happen-before order; load
   // shadow data after application data; store zero shadow data before
   // application data. This ensure shadow loads return either labels of the
   // initial application data or zeros.
   if (LI.isAtomic())
     LI.setOrdering(addAcquireOrdering(LI.getOrdering()));
 
+  Instruction *AfterLi = LI.getNextNode();
   Instruction *Pos = LI.isAtomic() ? LI.getNextNode() : &LI;
   std::vector<Value *> Shadows;
   std::vector<Value *> Origins;
   Value *PrimitiveShadow, *Origin;
   std::tie(PrimitiveShadow, Origin) =
       DFSF.loadShadowOrigin(LI.getPointerOperand(), Size, LI.getAlign(), Pos);
   const bool ShouldTrackOrigins = DFSF.DFS.shouldTrackOrigins();
   if (ShouldTrackOrigins) {
if (ShouldTrackOrigins) {
     DFSF.setOrigin(&LI, DFSF.combineOrigins(Shadows, Origins, Pos));
   }
 
   if (ClEventCallbacks) {
     IRBuilder<> IRB(Pos);
     Value *Addr8 = IRB.CreateBitCast(LI.getPointerOperand(), DFSF.DFS.Int8Ptr);
     IRB.CreateCall(DFSF.DFS.DFSanLoadCallbackFn, {PrimitiveShadow, Addr8});
   }
+
+  IRBuilder<> IRB(AfterLi);
+  DFSF.addReachesFunctionCallbacksIfEnabled(IRB, LI, &LI);
 }
 
 Value *DFSanFunction::updateOriginIfTainted(Value *Shadow, Value *Origin,
                                             IRBuilder<> &IRB) {
   assert(DFS.shouldTrackOrigins());
   return IRB.CreateCall(DFS.DFSanChainOriginIfTaintedFn, {Shadow, Origin});
 }
 
if (!CB.getType()->isVoidTy()) {
     }
 
     if (ShouldTrackOrigins) {
       LoadInst *LI = NextIRB.CreateLoad(DFSF.DFS.OriginTy,
                                         DFSF.getRetvalOriginTLS(), "_dfsret_o");
       DFSF.SkipInsts.insert(LI);
       DFSF.setOrigin(&CB, LI);
     }
+
+    DFSF.addReachesFunctionCallbacksIfEnabled(NextIRB, CB, &CB);
   }
 }
 
 void DFSanVisitor::visitPHINode(PHINode &PN) {
   Type *ShadowTy = DFSF.DFS.getShadowTy(&PN);
   PHINode *ShadowPN =
       PHINode::Create(ShadowTy, PN.getNumIncomingValues(), "", &PN);
 

llvm/test/Instrumentation/DataFlowSanitizer/reaches_function.ll

+; RUN: opt < %s -passes=dfsan -dfsan-reaches-function-callbacks=1 -S | FileCheck %s
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-pc-linux-gnu"
+
+declare i32 @f()
+
+define void @load(i32) {
+  ; CHECK-LABEL: define void @load.dfsan
+  ; CHECK: call{{.*}}@__dfsan_reaches_function_callback
+  %i = alloca i32
+  store i32 %0, ptr %i
+  ret void
+}
+
+define void @store(ptr) {
+  ; CHECK-LABEL: define void @store.dfsan
+  ; CHECK: call{{.*}}@__dfsan_reaches_function_callback
+  %load = load i32, ptr %0
+  ret void
+}
+
+define void @call() {
+  ; CHECK-LABEL: define void @call.dfsan
+  ; CHECK: call{{.*}}@__dfsan_reaches_function_callback
+  %ret = call i32 @f()
+  ret void
+}
+
+; CHECK-LABEL: @__dfsan_reaches_function_callback(i8, ptr, i32, ptr)