diff --git a/compiler-rt/lib/dfsan/dfsan.cpp b/compiler-rt/lib/dfsan/dfsan.cpp --- a/compiler-rt/lib/dfsan/dfsan.cpp +++ b/compiler-rt/lib/dfsan/dfsan.cpp @@ -597,15 +597,42 @@ } } -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 ReleaseOrClearOrigins(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()) + ReleaseOrClearOrigins(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); @@ -629,13 +656,34 @@ 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); diff --git a/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp --- a/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp +++ b/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp @@ -926,8 +926,8 @@ /*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 = @@ -1141,6 +1141,7 @@ { 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); } @@ -2207,10 +2208,14 @@ 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)), + IRB.CreateZExtOrTrunc(I.getLength(), DFSF.DFS.IntptrTy)}); } void DFSanVisitor::visitMemTransferInst(MemTransferInst &I) { diff --git a/llvm/test/Instrumentation/DataFlowSanitizer/basic.ll b/llvm/test/Instrumentation/DataFlowSanitizer/basic.ll --- a/llvm/test/Instrumentation/DataFlowSanitizer/basic.ll +++ b/llvm/test/Instrumentation/DataFlowSanitizer/basic.ll @@ -31,7 +31,7 @@ ; 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) diff --git a/llvm/test/Instrumentation/DataFlowSanitizer/memset.ll b/llvm/test/Instrumentation/DataFlowSanitizer/memset.ll --- a/llvm/test/Instrumentation/DataFlowSanitizer/memset.ll +++ b/llvm/test/Instrumentation/DataFlowSanitizer/memset.ll @@ -6,7 +6,7 @@ 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 } diff --git a/llvm/test/Instrumentation/DataFlowSanitizer/origin_mem_intrinsic.ll b/llvm/test/Instrumentation/DataFlowSanitizer/origin_mem_intrinsic.ll new file mode 100644 --- /dev/null +++ b/llvm/test/Instrumentation/DataFlowSanitizer/origin_mem_intrinsic.ll @@ -0,0 +1,14 @@ +; 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