Skip to content

Commit db7fe6c

Browse files
[dfsan] Propagate origin tracking at store
This is a part of https://reviews.llvm.org/D95835. Reviewed By: morehouse, gbalats Differential Revision: https://reviews.llvm.org/D97789
1 parent f204804 commit db7fe6c

File tree

4 files changed

+549
-170
lines changed

4 files changed

+549
-170
lines changed

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

+193-19
Original file line numberDiff line numberDiff line change
@@ -207,6 +207,14 @@ static cl::opt<bool> ClTrackSelectControlFlow(
207207
"to results."),
208208
cl::Hidden, cl::init(true));
209209

210+
// TODO: This default value follows MSan. DFSan may use a different value.
211+
static cl::opt<int> ClInstrumentWithCallThreshold(
212+
"dfsan-instrument-with-call-threshold",
213+
cl::desc("If the function being instrumented requires more than "
214+
"this number of origin stores, use callbacks instead of "
215+
"inline checks (-1 means never use callbacks)."),
216+
cl::Hidden, cl::init(3500));
217+
210218
// Controls how to track origins.
211219
// * 0: do not track origins.
212220
// * 1: track origins at memory store operations.
@@ -426,6 +434,7 @@ class DataFlowSanitizer {
426434
FunctionCallee DFSanMaybeStoreOriginFn;
427435
SmallPtrSet<Value *, 16> DFSanRuntimeFunctions;
428436
MDNode *ColdCallWeights;
437+
MDNode *OriginStoreWeights;
429438
DFSanABIList ABIList;
430439
DenseMap<Value *, Function *> UnwrappedFnMap;
431440
AttrBuilder ReadOnlyNoneAttrs;
@@ -578,8 +587,9 @@ struct DFSanFunction {
578587
std::pair<Value *, Value *> loadShadowOrigin(Value *ShadowAddr, uint64_t Size,
579588
Align InstAlignment,
580589
Instruction *Pos);
581-
void storePrimitiveShadow(Value *Addr, uint64_t Size, Align Alignment,
582-
Value *PrimitiveShadow, Instruction *Pos);
590+
void storePrimitiveShadowOrigin(Value *Addr, uint64_t Size,
591+
Align InstAlignment, Value *PrimitiveShadow,
592+
Value *Origin, Instruction *Pos);
583593
/// Applies PrimitiveShadow to all primitive subtypes of T, returning
584594
/// the expanded shadow value.
585595
///
@@ -633,6 +643,33 @@ struct DFSanFunction {
633643
/// checks if it is possible to load labels and origins without using the
634644
/// callback function.
635645
bool useCallbackLoadLabelAndOrigin(uint64_t Size, Align InstAlignment);
646+
647+
/// Returns a chain at the current stack with previous origin V.
648+
Value *updateOrigin(Value *V, IRBuilder<> &IRB);
649+
650+
/// Creates an Intptr = Origin | Origin << 32 if Intptr's size is 64. Returns
651+
/// Origin otherwise.
652+
Value *originToIntptr(IRBuilder<> &IRB, Value *Origin);
653+
654+
/// Stores Origin into the address range [StoreOriginAddr, StoreOriginAddr +
655+
/// Size).
656+
void paintOrigin(IRBuilder<> &IRB, Value *Origin, Value *StoreOriginAddr,
657+
uint64_t StoreOriginSize, Align Alignment);
658+
659+
/// Stores Origin in terms of its Shadow value.
660+
/// * Do not write origins for zero shadows because we do not trace origins
661+
/// for untainted sinks.
662+
/// * Use __dfsan_maybe_store_origin if there are too many origin store
663+
/// instrumentations.
664+
void storeOrigin(Instruction *Pos, Value *Addr, uint64_t Size, Value *Shadow,
665+
Value *Origin, Value *StoreOriginAddr, Align InstAlignment);
666+
667+
/// Convert a scalar value to an i1 by comparing with 0.
668+
Value *convertToBool(Value *V, IRBuilder<> &IRB, const Twine &Name = "");
669+
670+
bool shouldInstrumentWithCall();
671+
672+
int NumOriginStores = 0;
636673
};
637674

638675
class DFSanVisitor : public InstVisitor<DFSanVisitor> {
@@ -837,6 +874,11 @@ static Value *expandFromPrimitiveShadowRecursive(
837874
llvm_unreachable("Unexpected shadow type");
838875
}
839876

877+
bool DFSanFunction::shouldInstrumentWithCall() {
878+
return ClInstrumentWithCallThreshold >= 0 &&
879+
NumOriginStores >= ClInstrumentWithCallThreshold;
880+
}
881+
840882
Value *DFSanFunction::expandFromPrimitiveShadow(Type *T, Value *PrimitiveShadow,
841883
Instruction *Pos) {
842884
Type *ShadowTy = DFS.getShadowTy(T);
@@ -1009,6 +1051,7 @@ bool DataFlowSanitizer::init(Module &M) {
10091051
/*isVarArg=*/false);
10101052

10111053
ColdCallWeights = MDBuilder(*Ctx).createBranchWeights(1, 1000);
1054+
OriginStoreWeights = MDBuilder(*Ctx).createBranchWeights(1, 1000);
10121055
return true;
10131056
}
10141057

@@ -2212,6 +2255,99 @@ void DFSanVisitor::visitLoadInst(LoadInst &LI) {
22122255
}
22132256
}
22142257

2258+
Value *DFSanFunction::updateOrigin(Value *V, IRBuilder<> &IRB) {
2259+
if (!DFS.shouldTrackOrigins())
2260+
return V;
2261+
return IRB.CreateCall(DFS.DFSanChainOriginFn, V);
2262+
}
2263+
2264+
Value *DFSanFunction::originToIntptr(IRBuilder<> &IRB, Value *Origin) {
2265+
const unsigned OriginSize = DataFlowSanitizer::OriginWidthBytes;
2266+
const DataLayout &DL = F->getParent()->getDataLayout();
2267+
unsigned IntptrSize = DL.getTypeStoreSize(DFS.IntptrTy);
2268+
if (IntptrSize == OriginSize)
2269+
return Origin;
2270+
assert(IntptrSize == OriginSize * 2);
2271+
Origin = IRB.CreateIntCast(Origin, DFS.IntptrTy, /* isSigned */ false);
2272+
return IRB.CreateOr(Origin, IRB.CreateShl(Origin, OriginSize * 8));
2273+
}
2274+
2275+
void DFSanFunction::paintOrigin(IRBuilder<> &IRB, Value *Origin,
2276+
Value *StoreOriginAddr,
2277+
uint64_t StoreOriginSize, Align Alignment) {
2278+
const unsigned OriginSize = DataFlowSanitizer::OriginWidthBytes;
2279+
const DataLayout &DL = F->getParent()->getDataLayout();
2280+
const Align IntptrAlignment = DL.getABITypeAlign(DFS.IntptrTy);
2281+
unsigned IntptrSize = DL.getTypeStoreSize(DFS.IntptrTy);
2282+
assert(IntptrAlignment >= MinOriginAlignment);
2283+
assert(IntptrSize >= OriginSize);
2284+
2285+
unsigned Ofs = 0;
2286+
Align CurrentAlignment = Alignment;
2287+
if (Alignment >= IntptrAlignment && IntptrSize > OriginSize) {
2288+
Value *IntptrOrigin = originToIntptr(IRB, Origin);
2289+
Value *IntptrStoreOriginPtr = IRB.CreatePointerCast(
2290+
StoreOriginAddr, PointerType::get(DFS.IntptrTy, 0));
2291+
for (unsigned I = 0; I < StoreOriginSize / IntptrSize; ++I) {
2292+
Value *Ptr =
2293+
I ? IRB.CreateConstGEP1_32(DFS.IntptrTy, IntptrStoreOriginPtr, I)
2294+
: IntptrStoreOriginPtr;
2295+
IRB.CreateAlignedStore(IntptrOrigin, Ptr, CurrentAlignment);
2296+
Ofs += IntptrSize / OriginSize;
2297+
CurrentAlignment = IntptrAlignment;
2298+
}
2299+
}
2300+
2301+
for (unsigned I = Ofs; I < (StoreOriginSize + OriginSize - 1) / OriginSize;
2302+
++I) {
2303+
Value *GEP = I ? IRB.CreateConstGEP1_32(DFS.OriginTy, StoreOriginAddr, I)
2304+
: StoreOriginAddr;
2305+
IRB.CreateAlignedStore(Origin, GEP, CurrentAlignment);
2306+
CurrentAlignment = MinOriginAlignment;
2307+
}
2308+
}
2309+
2310+
Value *DFSanFunction::convertToBool(Value *V, IRBuilder<> &IRB,
2311+
const Twine &Name) {
2312+
Type *VTy = V->getType();
2313+
assert(VTy->isIntegerTy());
2314+
if (VTy->getIntegerBitWidth() == 1)
2315+
// Just converting a bool to a bool, so do nothing.
2316+
return V;
2317+
return IRB.CreateICmpNE(V, ConstantInt::get(VTy, 0), Name);
2318+
}
2319+
2320+
void DFSanFunction::storeOrigin(Instruction *Pos, Value *Addr, uint64_t Size,
2321+
Value *Shadow, Value *Origin,
2322+
Value *StoreOriginAddr, Align InstAlignment) {
2323+
// Do not write origins for zero shadows because we do not trace origins for
2324+
// untainted sinks.
2325+
const Align OriginAlignment = getOriginAlign(InstAlignment);
2326+
Value *CollapsedShadow = collapseToPrimitiveShadow(Shadow, Pos);
2327+
IRBuilder<> IRB(Pos);
2328+
if (auto *ConstantShadow = dyn_cast<Constant>(CollapsedShadow)) {
2329+
if (!ConstantShadow->isZeroValue())
2330+
paintOrigin(IRB, updateOrigin(Origin, IRB), StoreOriginAddr, Size,
2331+
OriginAlignment);
2332+
return;
2333+
}
2334+
2335+
if (shouldInstrumentWithCall()) {
2336+
IRB.CreateCall(DFS.DFSanMaybeStoreOriginFn,
2337+
{CollapsedShadow,
2338+
IRB.CreatePointerCast(Addr, IRB.getInt8PtrTy()),
2339+
ConstantInt::get(DFS.IntptrTy, Size), Origin});
2340+
} else {
2341+
Value *Cmp = convertToBool(CollapsedShadow, IRB, "_dfscmp");
2342+
Instruction *CheckTerm = SplitBlockAndInsertIfThen(
2343+
Cmp, &*IRB.GetInsertPoint(), false, DFS.OriginStoreWeights, &DT);
2344+
IRBuilder<> IRBNew(CheckTerm);
2345+
paintOrigin(IRBNew, updateOrigin(Origin, IRBNew), StoreOriginAddr, Size,
2346+
OriginAlignment);
2347+
++NumOriginStores;
2348+
}
2349+
}
2350+
22152351
void DFSanFunction::storeZeroPrimitiveShadow(Value *Addr, uint64_t Size,
22162352
Align ShadowAlign,
22172353
Instruction *Pos) {
@@ -2226,30 +2362,46 @@ void DFSanFunction::storeZeroPrimitiveShadow(Value *Addr, uint64_t Size,
22262362
// Do not write origins for 0 shadows because we do not trace origins for
22272363
// untainted sinks.
22282364
}
2229-
void DFSanFunction::storePrimitiveShadow(Value *Addr, uint64_t Size,
2230-
Align Alignment,
2231-
Value *PrimitiveShadow,
2232-
Instruction *Pos) {
2365+
2366+
void DFSanFunction::storePrimitiveShadowOrigin(Value *Addr, uint64_t Size,
2367+
Align InstAlignment,
2368+
Value *PrimitiveShadow,
2369+
Value *Origin,
2370+
Instruction *Pos) {
2371+
const bool ShouldTrackOrigins = DFS.shouldTrackOrigins() && Origin;
2372+
22332373
if (AllocaInst *AI = dyn_cast<AllocaInst>(Addr)) {
2234-
const auto I = AllocaShadowMap.find(AI);
2235-
if (I != AllocaShadowMap.end()) {
2374+
const auto SI = AllocaShadowMap.find(AI);
2375+
if (SI != AllocaShadowMap.end()) {
22362376
IRBuilder<> IRB(Pos);
2237-
IRB.CreateStore(PrimitiveShadow, I->second);
2377+
IRB.CreateStore(PrimitiveShadow, SI->second);
2378+
2379+
// Do not write origins for 0 shadows because we do not trace origins for
2380+
// untainted sinks.
2381+
if (ShouldTrackOrigins && !DFS.isZeroShadow(PrimitiveShadow)) {
2382+
const auto OI = AllocaOriginMap.find(AI);
2383+
assert(OI != AllocaOriginMap.end() && Origin);
2384+
IRB.CreateStore(Origin, OI->second);
2385+
}
22382386
return;
22392387
}
22402388
}
22412389

2242-
const Align ShadowAlign(Alignment.value() * DFS.ShadowWidthBytes);
2390+
const Align ShadowAlign = getShadowAlign(InstAlignment);
22432391
if (DFS.isZeroShadow(PrimitiveShadow)) {
22442392
storeZeroPrimitiveShadow(Addr, Size, ShadowAlign, Pos);
22452393
return;
22462394
}
22472395

22482396
IRBuilder<> IRB(Pos);
2249-
Value *ShadowAddr = DFS.getShadowAddress(Addr, Pos);
2397+
Value *ShadowAddr, *OriginAddr;
2398+
std::tie(ShadowAddr, OriginAddr) =
2399+
DFS.getShadowOriginAddress(Addr, InstAlignment, Pos);
2400+
22502401
const unsigned ShadowVecSize = 128 / DFS.ShadowWidthBits;
22512402
uint64_t Offset = 0;
2252-
if (Size >= ShadowVecSize) {
2403+
uint64_t LeftSize = Size;
2404+
if (LeftSize >= ShadowVecSize) {
22532405
auto *ShadowVecTy =
22542406
FixedVectorType::get(DFS.PrimitiveShadowTy, ShadowVecSize);
22552407
Value *ShadowVec = UndefValue::get(ShadowVecTy);
@@ -2264,18 +2416,23 @@ void DFSanFunction::storePrimitiveShadow(Value *Addr, uint64_t Size,
22642416
Value *CurShadowVecAddr =
22652417
IRB.CreateConstGEP1_32(ShadowVecTy, ShadowVecAddr, Offset);
22662418
IRB.CreateAlignedStore(ShadowVec, CurShadowVecAddr, ShadowAlign);
2267-
Size -= ShadowVecSize;
2419+
LeftSize -= ShadowVecSize;
22682420
++Offset;
2269-
} while (Size >= ShadowVecSize);
2421+
} while (LeftSize >= ShadowVecSize);
22702422
Offset *= ShadowVecSize;
22712423
}
2272-
while (Size > 0) {
2424+
while (LeftSize > 0) {
22732425
Value *CurShadowAddr =
22742426
IRB.CreateConstGEP1_32(DFS.PrimitiveShadowTy, ShadowAddr, Offset);
22752427
IRB.CreateAlignedStore(PrimitiveShadow, CurShadowAddr, ShadowAlign);
2276-
--Size;
2428+
--LeftSize;
22772429
++Offset;
22782430
}
2431+
2432+
if (ShouldTrackOrigins) {
2433+
storeOrigin(Pos, Addr, Size, PrimitiveShadow, Origin, OriginAddr,
2434+
InstAlignment);
2435+
}
22792436
}
22802437

22812438
static AtomicOrdering addReleaseOrdering(AtomicOrdering AO) {
@@ -2310,19 +2467,36 @@ void DFSanVisitor::visitStoreInst(StoreInst &SI) {
23102467
if (SI.isAtomic())
23112468
SI.setOrdering(addReleaseOrdering(SI.getOrdering()));
23122469

2313-
const Align Alignment = ClPreserveAlignment ? SI.getAlign() : Align(1);
2470+
const bool ShouldTrackOrigins =
2471+
DFSF.DFS.shouldTrackOrigins() && !SI.isAtomic();
2472+
std::vector<Value *> Shadows;
2473+
std::vector<Value *> Origins;
23142474

23152475
Value *Shadow =
23162476
SI.isAtomic() ? DFSF.DFS.getZeroShadow(Val) : DFSF.getShadow(Val);
2477+
2478+
if (ShouldTrackOrigins) {
2479+
Shadows.push_back(Shadow);
2480+
Origins.push_back(DFSF.getOrigin(Val));
2481+
}
2482+
23172483
Value *PrimitiveShadow;
23182484
if (ClCombinePointerLabelsOnStore) {
23192485
Value *PtrShadow = DFSF.getShadow(SI.getPointerOperand());
2486+
if (ShouldTrackOrigins) {
2487+
Shadows.push_back(PtrShadow);
2488+
Origins.push_back(DFSF.getOrigin(SI.getPointerOperand()));
2489+
}
23202490
PrimitiveShadow = DFSF.combineShadows(Shadow, PtrShadow, &SI);
23212491
} else {
23222492
PrimitiveShadow = DFSF.collapseToPrimitiveShadow(Shadow, &SI);
23232493
}
2324-
DFSF.storePrimitiveShadow(SI.getPointerOperand(), Size, Alignment,
2325-
PrimitiveShadow, &SI);
2494+
Value *Origin = nullptr;
2495+
if (ShouldTrackOrigins) {
2496+
Origin = DFSF.combineOrigins(Shadows, Origins, &SI);
2497+
}
2498+
DFSF.storePrimitiveShadowOrigin(SI.getPointerOperand(), Size, SI.getAlign(),
2499+
PrimitiveShadow, Origin, &SI);
23262500
if (ClEventCallbacks) {
23272501
IRBuilder<> IRB(&SI);
23282502
Value *Addr8 = IRB.CreateBitCast(SI.getPointerOperand(), DFSF.DFS.Int8Ptr);

llvm/test/Instrumentation/DataFlowSanitizer/basic.ll

+19-6
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
; RUN: opt < %s -dfsan -S | FileCheck %s --check-prefixes=CHECK,CHECK_NO_ORIGIN
2-
; RUN: opt < %s -dfsan -dfsan-track-origins=1 -dfsan-fast-16-labels=true -S | FileCheck %s --check-prefixes=CHECK,CHECK_ORIGIN
1+
; RUN: opt < %s -dfsan -S | FileCheck %s --check-prefixes=CHECK,CHECK_NO_ORIGIN -DSHADOW_MASK=-123145302310913
2+
; RUN: opt < %s -dfsan -dfsan-track-origins=1 -dfsan-fast-16-labels=true -S | FileCheck %s --check-prefixes=CHECK,CHECK_ORIGIN -DSHADOW_MASK=-123145302310913
33
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"
44
target triple = "x86_64-unknown-linux-gnu"
55

@@ -12,6 +12,23 @@ target triple = "x86_64-unknown-linux-gnu"
1212
; CHECK: @__dfsan_shadow_width_bits = weak_odr constant i32 [[#SBITS:]]
1313
; CHECK: @__dfsan_shadow_width_bytes = weak_odr constant i32 [[#SBYTES:]]
1414
; CHECK: @__dfsan_shadow_ptr_mask = external global i64
15+
16+
define i8 @load(i8* %p) {
17+
; CHECK-LABEL: define i8 @"dfs$load"
18+
; CHECK: and i64 {{.*}}, [[SHADOW_MASK]]
19+
; CHECK: ret i8 %a
20+
%a = load i8, i8* %p
21+
ret i8 %a
22+
}
23+
24+
define void @store(i8* %p) {
25+
; CHECK-LABEL: define void @"dfs$store"
26+
; CHECK: and i64 {{.*}}, [[SHADOW_MASK]]
27+
; CHECK: ret void
28+
store i8 0, i8* %p
29+
ret void
30+
}
31+
1532
; CHECK: declare void @__dfsan_load_callback(i[[#SBITS]], i8*)
1633
; CHECK: declare void @__dfsan_store_callback(i[[#SBITS]], i8*)
1734
; CHECK: declare void @__dfsan_mem_transfer_callback(i[[#SBITS]]*, i64)
@@ -39,7 +56,3 @@ target triple = "x86_64-unknown-linux-gnu"
3956
; CHECK: declare zeroext i32 @__dfsan_chain_origin(i32 zeroext)
4057
; CHECK: declare void @__dfsan_mem_origin_transfer(i8*, i8*, i64)
4158
; CHECK: declare void @__dfsan_maybe_store_origin(i[[#SBITS]] zeroext, i8*, i64, i32 zeroext)
42-
43-
define void @foo() {
44-
ret void
45-
}

0 commit comments

Comments
 (0)