[dfsan] Update memset and dfsan_(set|add)_label with origin tracking

[stephan.yichao.zhao]

Summary

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

[morehouse]

compiler-rt/lib/dfsan/dfsan.cpp

601

Actually, this doesn't zero the non-released pages.

[stephan.yichao.zhao]

renamed ClearOrigins to be ReleaseOrClearOrigins

compiler-rt/lib/dfsan/dfsan.cpp

extern "C" SANITIZER_INTERFACE_ATTRIBUTE void __dfsan_maybe_store_origin(
   if (UNLIKELY(s)) {
     GET_CALLER_PC_BP_SP;
     (void)sp;
     GET_STORE_STACK_TRACE_PC_BP(pc, bp);
     SetOrigin(p, size, ChainOrigin(o, &stack));
   }
 }
 
-extern "C" SANITIZER_INTERFACE_ATTRIBUTE void __dfsan_set_label(
-    dfsan_label label, void *addr, uptr size) {
+// Releases the pages within the origin address range, and sets the origin
+// addresses not on the pages to be 0.
+static void ClearOrigins(void *addr, uptr size) {
+  const uptr beg_origin_addr = (uptr)__dfsan::origin_for(addr);
+  const void *end_addr = (void *)((uptr)addr + size);
+  const uptr end_origin_addr = (uptr)__dfsan::origin_for(end_addr);
+  const uptr page_size = GetPageSizeCached();
+  const uptr beg_aligned = RoundUpTo(beg_origin_addr, page_size);
+  const uptr end_aligned = RoundDownTo(end_origin_addr, page_size);
+
+  // dfsan_set_label can be called from the following cases
+  // 1) mapped ranges by new/delete and malloc/free. This case has origin memory
+  // size > 50k, and happens less frequently.
+  // 2) zero-filling internal data structures by utility libraries. This case
+  // has origin memory size < 16k, and happens more often.
+  // Set kNumPagesThreshold to be 4 to avoid releasing small pages.
+  const int kNumPagesThreshold = 4;
+  if (beg_aligned + kNumPagesThreshold * page_size >= end_aligned)
+    return;
+
+  ReleaseMemoryPagesToOS(beg_aligned, end_aligned);
+}
+
+void SetShadow(dfsan_label label, void *addr, uptr size, dfsan_origin origin) {
   const uptr beg_shadow_addr = (uptr)__dfsan::shadow_for(addr);
 
   if (0 != label) {
     WriteShadowIfDifferent(label, beg_shadow_addr, size);
+    if (__dfsan_get_track_origins())
+      SetOrigin(addr, size, origin);
     return;
   }
 
+  if (__dfsan_get_track_origins())
+    ClearOrigins(addr, size);
+
   // If label is 0, releases the pages within the shadow address range, and sets
   // the shadow addresses not on the pages to be 0.
   const void *end_addr = (void *)((uptr)addr + size);
   const uptr end_shadow_addr = (uptr)__dfsan::shadow_for(end_addr);
   const uptr page_size = GetPageSizeCached();
   const uptr beg_aligned = RoundUpTo(beg_shadow_addr, page_size);
   const uptr end_aligned = RoundDownTo(end_shadow_addr, page_size);
 
   // dfsan_set_label can be called from the following cases
   // 1) mapped ranges by new/delete and malloc/free. This case has shadow memory
   // size > 100k, and happens less frequently.
   // 2) zero-filling internal data structures by utility libraries. This case
   // has shadow memory size < 32k, and happens more often.
   // Set kNumPagesThreshold to be 8 to avoid releasing small pages.
   const int kNumPagesThreshold = 8;
   if (beg_aligned + kNumPagesThreshold * page_size >= end_aligned)
     return WriteShadowIfDifferent(label, beg_shadow_addr, size);
 
   WriteShadowIfDifferent(label, beg_shadow_addr, beg_aligned - beg_shadow_addr);
   ReleaseMemoryPagesToOS(beg_aligned, end_aligned);
   WriteShadowIfDifferent(label, end_aligned, end_shadow_addr - end_aligned);
 }
 
+extern "C" SANITIZER_INTERFACE_ATTRIBUTE void __dfsan_set_label(
+    dfsan_label label, dfsan_origin origin, void *addr, uptr size) {
+  SetShadow(label, addr, size, origin);
+}
+
 SANITIZER_INTERFACE_ATTRIBUTE
 void dfsan_set_label(dfsan_label label, void *addr, uptr size) {
-  __dfsan_set_label(label, addr, size);
+  dfsan_origin init_origin = 0;
+  if (label && __dfsan_get_track_origins()) {
+    GET_CALLER_PC_BP;
+    GET_STORE_STACK_TRACE_PC_BP(pc, bp);
+    init_origin = ChainOrigin(0, &stack, true);
+  }
+  SetShadow(label, addr, size, init_origin);
 }
 
 SANITIZER_INTERFACE_ATTRIBUTE
 void dfsan_add_label(dfsan_label label, void *addr, uptr size) {
+  if (0 == label)
+    return;
+
+  if (__dfsan_get_track_origins()) {
+    GET_CALLER_PC_BP;
+    GET_STORE_STACK_TRACE_PC_BP(pc, bp);
+    dfsan_origin init_origin = ChainOrigin(0, &stack, true);
+    SetOrigin(addr, size, init_origin);
+  }
+
   for (dfsan_label *labelp = shadow_for(addr); size != 0; --size, ++labelp)
     if (*labelp != label)
       *labelp = __dfsan_union(*labelp, label);
 }
 
 // Unlike the other dfsan interface functions the behavior of this function
 // depends on the label of one of its arguments.  Hence it is implemented as a
 // custom function.

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

bool DataFlowSanitizer::init(Module &M) {
   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};
+  Type *DFSanSetLabelArgs[4] = {PrimitiveShadowTy, OriginTy,
+                                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,
void DataFlowSanitizer::initializeRuntimeFunctions(Module &M) {
     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);
+    AL = AL.addParamAttribute(M.getContext(), 1, Attribute::ZExt);
     DFSanSetLabelFn =
         Mod->getOrInsertFunction("__dfsan_set_label", DFSanSetLabelFnTy, AL);
   }
   DFSanNonzeroLabelFn =
       Mod->getOrInsertFunction("__dfsan_nonzero_label", DFSanNonzeroLabelFnTy);
   DFSanVarargWrapperFn = Mod->getOrInsertFunction("__dfsan_vararg_wrapper",
                                                   DFSanVarargWrapperFnTy);
   {
if (ShouldTrackOrigins) {
     }
     DFSF.setOrigin(&I, DFSF.combineOrigins(Shadows, Origins, &I));
   }
 }
 
 void DFSanVisitor::visitMemSetInst(MemSetInst &I) {
   IRBuilder<> IRB(&I);
   Value *ValShadow = DFSF.getShadow(I.getValue());
-  IRB.CreateCall(DFSF.DFS.DFSanSetLabelFn,
-                 {ValShadow, IRB.CreateBitCast(I.getDest(), Type::getInt8PtrTy(
-                                                                *DFSF.DFS.Ctx)),
-                  IRB.CreateZExtOrTrunc(I.getLength(), DFSF.DFS.IntptrTy)});
+  Value *ValOrigin = DFSF.DFS.shouldTrackOrigins()
+                         ? DFSF.getOrigin(I.getValue())
+                         : DFSF.DFS.ZeroOrigin;
+  IRB.CreateCall(
+      DFSF.DFS.DFSanSetLabelFn,
+      {ValShadow, ValOrigin,
+       IRB.CreateBitCast(I.getDest(), Type::getInt8PtrTy(*DFSF.DFS.Ctx)),
 }
 
 void DFSanVisitor::visitMemTransferInst(MemTransferInst &I) {
   IRBuilder<> IRB(&I);
   Value *RawDestShadow = DFSF.DFS.getShadowAddress(I.getDest(), &I);
   Value *SrcShadow = DFSF.DFS.getShadowAddress(I.getSource(), &I);
   Value *LenShadow =
       IRB.CreateMul(I.getLength(), ConstantInt::get(I.getLength()->getType(),

llvm/test/Instrumentation/DataFlowSanitizer/basic.ll

 
 ; 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_set_label(i16 zeroext, i32 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
 }

llvm/test/Instrumentation/DataFlowSanitizer/memset.ll

 ; 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"
 target triple = "x86_64-unknown-linux-gnu"
 
 declare void @llvm.memset.p0i8.i64(i8* nocapture, i8, i64, i1)
 
 define void @ms(i8* %p, i8 %v) {
   ; CHECK-LABEL: @"dfs$ms"(i8* %0, i8 %1, i16 %2, i16 %3)
-  ; CHECK: call void @__dfsan_set_label(i16 %3, i8* %0, i64 1)
+  ; CHECK: call void @__dfsan_set_label(i16 %3, i32 0, i8* %0, i64 1)
   call void @llvm.memset.p0i8.i64(i8* %p, i8 %v, i64 1, i1 1)
   ret void
 }

llvm/test/Instrumentation/DataFlowSanitizer/origin_mem_intrinsic.ll

+; RUN: opt < %s -dfsan -dfsan-track-origins=1 -dfsan-fast-16-labels=true -S | FileCheck %s --check-prefix=CHECK
+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"
+
+declare void @llvm.memset.p0i8.i64(i8* nocapture, i8, i64, i1)
+
+define void @memset(i8* %p, i8 %v) {
+  ; CHECK: @"dfs$memset"
+  ; CHECK: [[O:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 1), align 4
+  ; CHECK: [[S:%.*]] = load i16, i16* inttoptr (i64 add (i64 ptrtoint ([100 x i64]* @__dfsan_arg_tls to i64), i64 2) to i16*), align 2
+  ; CHECK: call void @__dfsan_set_label(i16 [[S]], i32 [[O]], i8* %p, i64 1)
+  call void @llvm.memset.p0i8.i64(i8* %p, i8 %v, i64 1, i1 1)
+  ret void
+}
 No newline at end of file