[dfsan] Instrument origin variable and function definitions

[stephan.yichao.zhao]

Summary

This is a part of https://reviews.llvm.org/D95835.

[morehouse]

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

482

Can this be constant instead of a function?

llvm/test/Instrumentation/DataFlowSanitizer/basic.ll

12

IIRC, even with --check-prefix, the CHECK lines are still verified. So don't we get a failure for the line above this one (track_origins = 0)?

[stephan.yichao.zhao]

update

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

482

Thank you. Replaced by kNumOfElementsInArgOrgTLS;

llvm/test/Instrumentation/DataFlowSanitizer/basic.ll

12

Hm, I found FileCheck ignores CHECK if --check-prefix is used.
But the problem is that we may want to rename the last one with prefix CHECK_NO_ORIGIN.

Updated.

[gbalats]

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

335

Is this tied to ShadowWidthBits? With fast8, how would this change?

llvm/test/Instrumentation/DataFlowSanitizer/basic.ll

1–2

You could use --check-prefixes to merge these.

[morehouse]

llvm/test/Instrumentation/DataFlowSanitizer/basic.ll

1–2

12

Interesting, I guess I've been doing it wrong then. Thanks for informing!

[stephan.yichao.zhao]

replaced by --check-prefixes

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

335

Origin tracking always uses 4-byte.
Every 4 user byte shares one 4-byte origin. So its memory overhead is always 1x.

Origin is like a hash value, 32bit gives us enough hash space.
It would not be changed in 8-bit mode.

llvm/test/Instrumentation/DataFlowSanitizer/basic.ll

1–2

Done. Thank you.

1–2

Thank you. The two comments match exactly. :)

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

 
 // 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));
 
+// Controls how to track origins.
+// * 0: do not track origins.
+// * 1: track origins at memory store operations.
+// * 2: TODO: track origins at memory store operations and callsites.
+static cl::opt<int> ClTrackOrigins("dfsan-track-origins",
+                                   cl::desc("Track origins of labels"),
+                                   cl::Hidden, cl::init(0));
+
 static StringRef GetGlobalTypeString(const GlobalValue &G) {
   // Types of GlobalVariables are always pointer types.
   Type *GType = G.getValueType();
   // For now we support excluding struct types only.
   if (StructType *SGType = dyn_cast<StructType>(GType)) {
     if (!SGType->isLiteral())
       return SGType->getName();
   }
return AttributeList::get(
       CallSiteAttrs.getRetAttributes(),
       llvm::makeArrayRef(ArgumentAttributes));
 }
 
 class DataFlowSanitizer {
   friend struct DFSanFunction;
   friend class DFSanVisitor;
 
-  enum { ShadowWidthBits = 16, ShadowWidthBytes = ShadowWidthBits / 8 };
+  enum {
+    ShadowWidthBits = 16,
+    ShadowWidthBytes = ShadowWidthBits / 8,
+    OriginWidthBits = 32,
+    OriginWidthBytes = OriginWidthBits / 8
+  };
 
   /// Which ABI should be used for instrumented functions?
   enum InstrumentedABI {
     /// Argument and return value labels are passed through additional
     /// arguments and by modifying the return type.
     IA_Args,
 
     /// Argument and return value labels are passed through TLS variables
enum WrapperKind {
     /// extra pointer argument to return the shadow.  This allows the wrapped
     /// form of the function type to be expressed in C.
     WK_Custom
   };
 
   Module *Mod;
   LLVMContext *Ctx;
   Type *Int8Ptr;
+  IntegerType *OriginTy;
+  PointerType *OriginPtrTy;
   /// The shadow type for all primitive types and vector types.
   IntegerType *PrimitiveShadowTy;
   PointerType *PrimitiveShadowPtrTy;
   IntegerType *IntptrTy;
   ConstantInt *ZeroPrimitiveShadow;
   ConstantInt *ShadowPtrMask;
   ConstantInt *ShadowPtrMul;
   Constant *ArgTLS;
+  ArrayType *ArgOriginTLSTy;
+  Constant *ArgOriginTLS;
   Constant *RetvalTLS;
+  Constant *RetvalOriginTLS;
   Constant *ExternalShadowMask;
   FunctionType *DFSanUnionFnTy;
   FunctionType *DFSanUnionLoadFnTy;
+  FunctionType *DFSanLoadLabelAndOriginFnTy;
   FunctionType *DFSanUnimplementedFnTy;
   FunctionType *DFSanSetLabelFnTy;
   FunctionType *DFSanNonzeroLabelFnTy;
   FunctionType *DFSanVarargWrapperFnTy;
   FunctionType *DFSanCmpCallbackFnTy;
   FunctionType *DFSanLoadStoreCallbackFnTy;
   FunctionType *DFSanMemTransferCallbackFnTy;
+  FunctionType *DFSanChainOriginFnTy;
+  FunctionType *DFSanMemOriginTransferFnTy;
+  FunctionType *DFSanMaybeStoreOriginFnTy;
   FunctionCallee DFSanUnionFn;
   FunctionCallee DFSanCheckedUnionFn;
   FunctionCallee DFSanUnionLoadFn;
   FunctionCallee DFSanUnionLoadFast16LabelsFn;
+  FunctionCallee DFSanLoadLabelAndOriginFn;
   FunctionCallee DFSanUnimplementedFn;
   FunctionCallee DFSanSetLabelFn;
   FunctionCallee DFSanNonzeroLabelFn;
   FunctionCallee DFSanVarargWrapperFn;
   FunctionCallee DFSanLoadCallbackFn;
   FunctionCallee DFSanStoreCallbackFn;
   FunctionCallee DFSanMemTransferCallbackFn;
   FunctionCallee DFSanCmpCallbackFn;
-  SmallPtrSet<Value *, 12> DFSanRuntimeFunctions;
+  FunctionCallee DFSanChainOriginFn;
+  FunctionCallee DFSanMemOriginTransferFn;
+  FunctionCallee DFSanMaybeStoreOriginFn;
+  SmallPtrSet<Value *, 16> DFSanRuntimeFunctions;
   MDNode *ColdCallWeights;
   DFSanABIList ABIList;
   DenseMap<Value *, Function *> UnwrappedFnMap;
   AttrBuilder ReadOnlyNoneAttrs;
   bool DFSanRuntimeShadowMask = false;
 
   Value *getShadowAddress(Value *Addr, Instruction *Pos);
   bool isInstrumented(const Function *F);
   bool isInstrumented(const GlobalAlias *GA);
   FunctionType *getArgsFunctionType(FunctionType *T);
   FunctionType *getTrampolineFunctionType(FunctionType *T);
   TransformedFunction getCustomFunctionType(FunctionType *T);
   InstrumentedABI getInstrumentedABI();
   WrapperKind getWrapperKind(Function *F);
   void addGlobalNamePrefix(GlobalValue *GV);
   Function *buildWrapperFunction(Function *F, StringRef NewFName,
                                  GlobalValue::LinkageTypes NewFLink,
                                  FunctionType *NewFT);
   Constant *getOrBuildTrampolineFunction(FunctionType *FT, StringRef FName);
   void initializeCallbackFunctions(Module &M);
   void initializeRuntimeFunctions(Module &M);
 
   bool init(Module &M);
 
+  /// Returns whether the pass tracks origins. Support only fast16 mode in TLS
+  /// ABI mode.
+  bool shouldTrackOrigins();
+
   /// Returns whether the pass tracks labels for struct fields and array
   /// indices. Support only fast16 mode in TLS ABI mode.
   bool shouldTrackFieldsAndIndices();
 
   /// Returns a zero constant with the shadow type of OrigTy.
   ///
   /// getZeroShadow({T1,T2,...}) = {getZeroShadow(T1),getZeroShadow(T2,...}
   /// getZeroShadow([n x T]) = [n x getZeroShadow(T)]
class DataFlowSanitizer {
   /// getShadowTy(other type) = i16
   ///
   /// Note that a shadow type is always i16 when shouldTrackFieldsAndIndices
   /// returns false.
   Type *getShadowTy(Type *OrigTy);
   /// Returns the shadow type of of V's type.
   Type *getShadowTy(Value *V);
 
+  const uint64_t kNumOfElementsInArgOrgTLS = kArgTLSSize / OriginWidthBytes;
+
 public:
   DataFlowSanitizer(const std::vector<std::string> &ABIListFiles);
 
   bool runImpl(Module &M);
 };
 
 struct DFSanFunction {
   DataFlowSanitizer &DFS;
if (!isa<ArrayType>(T) && !isa<StructType>(T)) {
     if (const ConstantInt *CI = dyn_cast<ConstantInt>(V))
       return CI->isZero();
     return false;
   }
 
   return isa<ConstantAggregateZero>(V);
 }
 
+bool DataFlowSanitizer::shouldTrackOrigins() {
+  return ClTrackOrigins && getInstrumentedABI() == DataFlowSanitizer::IA_TLS &&
+         ClFast16Labels;
+}
+
 bool DataFlowSanitizer::shouldTrackFieldsAndIndices() {
   return getInstrumentedABI() == DataFlowSanitizer::IA_TLS && ClFast16Labels;
 }
 
 Constant *DataFlowSanitizer::getZeroShadow(Type *OrigTy) {
   if (!shouldTrackFieldsAndIndices())
     return ZeroPrimitiveShadow;
 
bool DataFlowSanitizer::init(Module &M) {
   bool IsAArch64 = TargetTriple.getArch() == Triple::aarch64 ||
                    TargetTriple.getArch() == Triple::aarch64_be;
 
   const DataLayout &DL = M.getDataLayout();
 
   Mod = &M;
   Ctx = &M.getContext();
   Int8Ptr = Type::getInt8PtrTy(*Ctx);
+  OriginTy = IntegerType::get(*Ctx, OriginWidthBits);
+  OriginPtrTy = PointerType::getUnqual(OriginTy);
   PrimitiveShadowTy = IntegerType::get(*Ctx, ShadowWidthBits);
   PrimitiveShadowPtrTy = PointerType::getUnqual(PrimitiveShadowTy);
   IntptrTy = DL.getIntPtrType(*Ctx);
   ZeroPrimitiveShadow = ConstantInt::getSigned(PrimitiveShadowTy, 0);
   ShadowPtrMul = ConstantInt::getSigned(IntptrTy, ShadowWidthBytes);
   if (IsX86_64)
     ShadowPtrMask = ConstantInt::getSigned(IntptrTy, ~0x700000000000LL);
   else if (IsMIPS64)
     ShadowPtrMask = ConstantInt::getSigned(IntptrTy, ~0xF000000000LL);
   // AArch64 supports multiple VMAs and the shadow mask is set at runtime.
   else if (IsAArch64)
     DFSanRuntimeShadowMask = true;
   else
     report_fatal_error("unsupported triple");
 
   Type *DFSanUnionArgs[2] = {PrimitiveShadowTy, PrimitiveShadowTy};
   DFSanUnionFnTy =
       FunctionType::get(PrimitiveShadowTy, DFSanUnionArgs, /*isVarArg=*/false);
   Type *DFSanUnionLoadArgs[2] = {PrimitiveShadowPtrTy, IntptrTy};
   DFSanUnionLoadFnTy = FunctionType::get(PrimitiveShadowTy, DFSanUnionLoadArgs,
                                          /*isVarArg=*/false);
+  Type *DFSanLoadLabelAndOriginArgs[2] = {Int8Ptr, IntptrTy};
+  DFSanLoadLabelAndOriginFnTy =
+      FunctionType::get(IntegerType::get(*Ctx, 64), DFSanLoadLabelAndOriginArgs,
+                        /*isVarArg=*/false);
   DFSanUnimplementedFnTy = FunctionType::get(
       Type::getVoidTy(*Ctx), Type::getInt8PtrTy(*Ctx), /*isVarArg=*/false);
   Type *DFSanSetLabelArgs[3] = {PrimitiveShadowTy, Type::getInt8PtrTy(*Ctx),
                                 IntptrTy};
   DFSanSetLabelFnTy = FunctionType::get(Type::getVoidTy(*Ctx),
                                         DFSanSetLabelArgs, /*isVarArg=*/false);
   DFSanNonzeroLabelFnTy =
       FunctionType::get(Type::getVoidTy(*Ctx), None, /*isVarArg=*/false);
   DFSanVarargWrapperFnTy = FunctionType::get(
       Type::getVoidTy(*Ctx), Type::getInt8PtrTy(*Ctx), /*isVarArg=*/false);
   DFSanCmpCallbackFnTy =
       FunctionType::get(Type::getVoidTy(*Ctx), PrimitiveShadowTy,
                         /*isVarArg=*/false);
+  DFSanChainOriginFnTy =
+      FunctionType::get(OriginTy, OriginTy, /*isVarArg=*/false);
+  Type *DFSanMaybeStoreOriginArgs[4] = {IntegerType::get(*Ctx, ShadowWidthBits),
+                                        Int8Ptr, IntptrTy, OriginTy};
+  DFSanMaybeStoreOriginFnTy = FunctionType::get(
+      Type::getVoidTy(*Ctx), DFSanMaybeStoreOriginArgs, /*isVarArg=*/false);
+  Type *DFSanMemOriginTransferArgs[3] = {Int8Ptr, Int8Ptr, IntptrTy};
+  DFSanMemOriginTransferFnTy = FunctionType::get(
+      Type::getVoidTy(*Ctx), DFSanMemOriginTransferArgs, /*isVarArg=*/false);
   Type *DFSanLoadStoreCallbackArgs[2] = {PrimitiveShadowTy, Int8Ptr};
   DFSanLoadStoreCallbackFnTy =
       FunctionType::get(Type::getVoidTy(*Ctx), DFSanLoadStoreCallbackArgs,
                         /*isVarArg=*/false);
   Type *DFSanMemTransferCallbackArgs[2] = {PrimitiveShadowPtrTy, IntptrTy};
   DFSanMemTransferCallbackFnTy =
       FunctionType::get(Type::getVoidTy(*Ctx), DFSanMemTransferCallbackArgs,
                         /*isVarArg=*/false);
AL = AL.addAttribute(M.getContext(), AttributeList::FunctionIndex,
                          Attribute::NoUnwind);
     AL = AL.addAttribute(M.getContext(), AttributeList::FunctionIndex,
                          Attribute::ReadOnly);
     AL = AL.addAttribute(M.getContext(), AttributeList::ReturnIndex,
                          Attribute::ZExt);
     DFSanUnionLoadFast16LabelsFn = Mod->getOrInsertFunction(
         "__dfsan_union_load_fast16labels", DFSanUnionLoadFnTy, AL);
   }
+  {
+    AttributeList AL;
+    AL = AL.addAttribute(M.getContext(), AttributeList::FunctionIndex,
+                         Attribute::NoUnwind);
+    AL = AL.addAttribute(M.getContext(), AttributeList::FunctionIndex,
+                         Attribute::ReadOnly);
+    AL = AL.addAttribute(M.getContext(), AttributeList::ReturnIndex,
+                         Attribute::ZExt);
+    DFSanLoadLabelAndOriginFn = Mod->getOrInsertFunction(
+        "__dfsan_load_label_and_origin", DFSanLoadLabelAndOriginFnTy, AL);
+  }
   DFSanUnimplementedFn =
       Mod->getOrInsertFunction("__dfsan_unimplemented", DFSanUnimplementedFnTy);
   {
     AttributeList AL;
     AL = AL.addParamAttribute(M.getContext(), 0, Attribute::ZExt);
     DFSanSetLabelFn =
         Mod->getOrInsertFunction("__dfsan_set_label", DFSanSetLabelFnTy, AL);
   }
   DFSanNonzeroLabelFn =
       Mod->getOrInsertFunction("__dfsan_nonzero_label", DFSanNonzeroLabelFnTy);
   DFSanVarargWrapperFn = Mod->getOrInsertFunction("__dfsan_vararg_wrapper",
                                                   DFSanVarargWrapperFnTy);
+  {
+    AttributeList AL;
+    AL = AL.addParamAttribute(M.getContext(), 0, Attribute::ZExt);
+    AL = AL.addAttribute(M.getContext(), AttributeList::ReturnIndex,
+                         Attribute::ZExt);
+    DFSanChainOriginFn = Mod->getOrInsertFunction("__dfsan_chain_origin",
+                                                  DFSanChainOriginFnTy, AL);
+  }
+  DFSanMemOriginTransferFn = Mod->getOrInsertFunction(
+      "__dfsan_mem_origin_transfer", DFSanMemOriginTransferFnTy);
+
+  {
+    AttributeList AL;
+    AL = AL.addParamAttribute(M.getContext(), 0, Attribute::ZExt);
+    AL = AL.addParamAttribute(M.getContext(), 3, Attribute::ZExt);
+    DFSanMaybeStoreOriginFn = Mod->getOrInsertFunction(
+        "__dfsan_maybe_store_origin", DFSanMaybeStoreOriginFnTy, AL);
+  }
 
   DFSanRuntimeFunctions.insert(DFSanUnionFn.getCallee()->stripPointerCasts());
   DFSanRuntimeFunctions.insert(
       DFSanCheckedUnionFn.getCallee()->stripPointerCasts());
   DFSanRuntimeFunctions.insert(
       DFSanUnionLoadFn.getCallee()->stripPointerCasts());
   DFSanRuntimeFunctions.insert(
       DFSanUnionLoadFast16LabelsFn.getCallee()->stripPointerCasts());
   DFSanRuntimeFunctions.insert(
+      DFSanLoadLabelAndOriginFn.getCallee()->stripPointerCasts());
+  DFSanRuntimeFunctions.insert(
       DFSanUnimplementedFn.getCallee()->stripPointerCasts());
   DFSanRuntimeFunctions.insert(
       DFSanSetLabelFn.getCallee()->stripPointerCasts());
   DFSanRuntimeFunctions.insert(
       DFSanNonzeroLabelFn.getCallee()->stripPointerCasts());
   DFSanRuntimeFunctions.insert(
       DFSanVarargWrapperFn.getCallee()->stripPointerCasts());
   DFSanRuntimeFunctions.insert(
       DFSanLoadCallbackFn.getCallee()->stripPointerCasts());
   DFSanRuntimeFunctions.insert(
       DFSanStoreCallbackFn.getCallee()->stripPointerCasts());
   DFSanRuntimeFunctions.insert(
       DFSanMemTransferCallbackFn.getCallee()->stripPointerCasts());
   DFSanRuntimeFunctions.insert(
       DFSanCmpCallbackFn.getCallee()->stripPointerCasts());
+  DFSanRuntimeFunctions.insert(
+      DFSanChainOriginFn.getCallee()->stripPointerCasts());
+  DFSanRuntimeFunctions.insert(
+      DFSanMemOriginTransferFn.getCallee()->stripPointerCasts());
+  DFSanRuntimeFunctions.insert(
+      DFSanMaybeStoreOriginFn.getCallee()->stripPointerCasts());
 }
 
 // Initializes event callback functions and declare them in the module
 void DataFlowSanitizer::initializeCallbackFunctions(Module &M) {
   DFSanLoadCallbackFn = Mod->getOrInsertFunction("__dfsan_load_callback",
                                                  DFSanLoadStoreCallbackFnTy);
   DFSanStoreCallbackFn = Mod->getOrInsertFunction("__dfsan_store_callback",
                                                   DFSanLoadStoreCallbackFnTy);
bool DataFlowSanitizer::runImpl(Module &M) {
 
   // These globals must be kept in sync with the ones in dfsan.cpp.
   ArgTLS = getOrInsertGlobal(
       "__dfsan_arg_tls",
       ArrayType::get(Type::getInt64Ty(*Ctx), kArgTLSSize / 8));
   RetvalTLS = getOrInsertGlobal(
       "__dfsan_retval_tls",
       ArrayType::get(Type::getInt64Ty(*Ctx), kRetvalTLSSize / 8));
+  ArgOriginTLSTy = ArrayType::get(OriginTy, kNumOfElementsInArgOrgTLS);
+  ArgOriginTLS = getOrInsertGlobal("__dfsan_arg_origin_tls", ArgOriginTLSTy);
+  RetvalOriginTLS = getOrInsertGlobal("__dfsan_retval_origin_tls", OriginTy);
+
+  (void)Mod->getOrInsertGlobal("__dfsan_track_origins", OriginTy, [&] {
+    Changed = true;
+    return new GlobalVariable(
+        M, OriginTy, true, GlobalValue::WeakODRLinkage,
+        ConstantInt::getSigned(OriginTy, shouldTrackOrigins()),
+        "__dfsan_track_origins");
+  });
 
   ExternalShadowMask =
       Mod->getOrInsertGlobal(kDFSanExternShadowPtrMask, IntptrTy);
 
   initializeCallbackFunctions(M);
   initializeRuntimeFunctions(M);
 
   std::vector<Function *> FnsToInstrument;

llvm/test/Instrumentation/DataFlowSanitizer/basic.ll

+; RUN: opt < %s -dfsan -S | FileCheck %s --check-prefix=CHECK --check-prefix=CHECK_NO_ORIGIN
+; RUN: opt < %s -dfsan -dfsan-track-origins=1 -dfsan-fast-16-labels=true -S | FileCheck %s --check-prefix=CHECK --check-prefix=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"
+
+; CHECK: @__dfsan_arg_tls = external thread_local(initialexec) global [100 x i64]
+; CHECK: @__dfsan_retval_tls = external thread_local(initialexec) global [100 x i64]
+; CHECK: @__dfsan_arg_origin_tls = external thread_local(initialexec) global [200 x i32]
+; CHECK: @__dfsan_retval_origin_tls = external thread_local(initialexec) global i32
+; CHECK_NO_ORIGIN: @__dfsan_track_origins = weak_odr constant i32 0
+; CHECK_ORIGIN: @__dfsan_track_origins = weak_odr constant i32 1
+; CHECK: @__dfsan_shadow_ptr_mask = external global i64
+; CHECK: declare void @__dfsan_load_callback(i16, i8*)
+; CHECK: declare void @__dfsan_store_callback(i16, i8*)
+; CHECK: declare void @__dfsan_mem_transfer_callback(i16*, i64)
+; CHECK: declare void @__dfsan_cmp_callback(i16)
+
+; CHECK: ; Function Attrs: nounwind readnone
+; CHECK-NEXT: declare zeroext i16 @__dfsan_union(i16 zeroext, i16 zeroext) #0
+
+; CHECK: ; Function Attrs: nounwind readnone
+; CHECK-NEXT: declare zeroext i16 @dfsan_union(i16 zeroext, i16 zeroext) #0
+
+; CHECK: ; Function Attrs: nounwind readonly
+; CHECK-NEXT: declare zeroext i16 @__dfsan_union_load(i16*, i64) #1
+
+; CHECK: ; Function Attrs: nounwind readonly
+; CHECK-NEXT: declare zeroext i16 @__dfsan_union_load_fast16labels(i16*, i64) #1
+
+; CHECK: ; Function Attrs: nounwind readonly
+; CHECK-NEXT: declare zeroext i64 @__dfsan_load_label_and_origin(i8*, i64) #1
+
+; CHECK: declare void @__dfsan_unimplemented(i8*)
+; CHECK: declare void @__dfsan_set_label(i16 zeroext, i8*, i64)
+; CHECK: declare void @__dfsan_nonzero_label()
+; CHECK: declare void @__dfsan_vararg_wrapper(i8*)
+; CHECK: declare zeroext i32 @__dfsan_chain_origin(i32 zeroext)
+; CHECK: declare void @__dfsan_mem_origin_transfer(i8*, i8*, i64)
+; CHECK: declare void @__dfsan_maybe_store_origin(i16 zeroext, i8*, i64, i32 zeroext)
+
+define void @foo() {
+  ret void
+}