@@ -207,6 +207,14 @@ static cl::opt<bool> ClTrackSelectControlFlow(
207
207
" to results." ),
208
208
cl::Hidden, cl::init(true ));
209
209
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
+
210
218
// Controls how to track origins.
211
219
// * 0: do not track origins.
212
220
// * 1: track origins at memory store operations.
@@ -426,6 +434,7 @@ class DataFlowSanitizer {
426
434
FunctionCallee DFSanMaybeStoreOriginFn;
427
435
SmallPtrSet<Value *, 16 > DFSanRuntimeFunctions;
428
436
MDNode *ColdCallWeights;
437
+ MDNode *OriginStoreWeights;
429
438
DFSanABIList ABIList;
430
439
DenseMap<Value *, Function *> UnwrappedFnMap;
431
440
AttrBuilder ReadOnlyNoneAttrs;
@@ -578,8 +587,9 @@ struct DFSanFunction {
578
587
std::pair<Value *, Value *> loadShadowOrigin (Value *ShadowAddr, uint64_t Size ,
579
588
Align InstAlignment,
580
589
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);
583
593
// / Applies PrimitiveShadow to all primitive subtypes of T, returning
584
594
// / the expanded shadow value.
585
595
// /
@@ -633,6 +643,33 @@ struct DFSanFunction {
633
643
// / checks if it is possible to load labels and origins without using the
634
644
// / callback function.
635
645
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 ;
636
673
};
637
674
638
675
class DFSanVisitor : public InstVisitor <DFSanVisitor> {
@@ -837,6 +874,11 @@ static Value *expandFromPrimitiveShadowRecursive(
837
874
llvm_unreachable (" Unexpected shadow type" );
838
875
}
839
876
877
+ bool DFSanFunction::shouldInstrumentWithCall () {
878
+ return ClInstrumentWithCallThreshold >= 0 &&
879
+ NumOriginStores >= ClInstrumentWithCallThreshold;
880
+ }
881
+
840
882
Value *DFSanFunction::expandFromPrimitiveShadow (Type *T, Value *PrimitiveShadow,
841
883
Instruction *Pos) {
842
884
Type *ShadowTy = DFS.getShadowTy (T);
@@ -1009,6 +1051,7 @@ bool DataFlowSanitizer::init(Module &M) {
1009
1051
/* isVarArg=*/ false );
1010
1052
1011
1053
ColdCallWeights = MDBuilder (*Ctx).createBranchWeights (1 , 1000 );
1054
+ OriginStoreWeights = MDBuilder (*Ctx).createBranchWeights (1 , 1000 );
1012
1055
return true ;
1013
1056
}
1014
1057
@@ -2212,6 +2255,99 @@ void DFSanVisitor::visitLoadInst(LoadInst &LI) {
2212
2255
}
2213
2256
}
2214
2257
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
+
2215
2351
void DFSanFunction::storeZeroPrimitiveShadow (Value *Addr, uint64_t Size ,
2216
2352
Align ShadowAlign,
2217
2353
Instruction *Pos) {
@@ -2226,30 +2362,46 @@ void DFSanFunction::storeZeroPrimitiveShadow(Value *Addr, uint64_t Size,
2226
2362
// Do not write origins for 0 shadows because we do not trace origins for
2227
2363
// untainted sinks.
2228
2364
}
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
+
2233
2373
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 ()) {
2236
2376
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
+ }
2238
2386
return ;
2239
2387
}
2240
2388
}
2241
2389
2242
- const Align ShadowAlign (Alignment. value () * DFS. ShadowWidthBytes );
2390
+ const Align ShadowAlign = getShadowAlign (InstAlignment );
2243
2391
if (DFS.isZeroShadow (PrimitiveShadow)) {
2244
2392
storeZeroPrimitiveShadow (Addr, Size , ShadowAlign, Pos);
2245
2393
return ;
2246
2394
}
2247
2395
2248
2396
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
+
2250
2401
const unsigned ShadowVecSize = 128 / DFS.ShadowWidthBits ;
2251
2402
uint64_t Offset = 0 ;
2252
- if (Size >= ShadowVecSize) {
2403
+ uint64_t LeftSize = Size ;
2404
+ if (LeftSize >= ShadowVecSize) {
2253
2405
auto *ShadowVecTy =
2254
2406
FixedVectorType::get (DFS.PrimitiveShadowTy , ShadowVecSize);
2255
2407
Value *ShadowVec = UndefValue::get (ShadowVecTy);
@@ -2264,18 +2416,23 @@ void DFSanFunction::storePrimitiveShadow(Value *Addr, uint64_t Size,
2264
2416
Value *CurShadowVecAddr =
2265
2417
IRB.CreateConstGEP1_32 (ShadowVecTy, ShadowVecAddr, Offset);
2266
2418
IRB.CreateAlignedStore (ShadowVec, CurShadowVecAddr, ShadowAlign);
2267
- Size -= ShadowVecSize;
2419
+ LeftSize -= ShadowVecSize;
2268
2420
++Offset;
2269
- } while (Size >= ShadowVecSize);
2421
+ } while (LeftSize >= ShadowVecSize);
2270
2422
Offset *= ShadowVecSize;
2271
2423
}
2272
- while (Size > 0 ) {
2424
+ while (LeftSize > 0 ) {
2273
2425
Value *CurShadowAddr =
2274
2426
IRB.CreateConstGEP1_32 (DFS.PrimitiveShadowTy , ShadowAddr, Offset);
2275
2427
IRB.CreateAlignedStore (PrimitiveShadow, CurShadowAddr, ShadowAlign);
2276
- --Size ;
2428
+ --LeftSize ;
2277
2429
++Offset;
2278
2430
}
2431
+
2432
+ if (ShouldTrackOrigins) {
2433
+ storeOrigin (Pos, Addr, Size , PrimitiveShadow, Origin, OriginAddr,
2434
+ InstAlignment);
2435
+ }
2279
2436
}
2280
2437
2281
2438
static AtomicOrdering addReleaseOrdering (AtomicOrdering AO) {
@@ -2310,19 +2467,36 @@ void DFSanVisitor::visitStoreInst(StoreInst &SI) {
2310
2467
if (SI.isAtomic ())
2311
2468
SI.setOrdering (addReleaseOrdering (SI.getOrdering ()));
2312
2469
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;
2314
2474
2315
2475
Value *Shadow =
2316
2476
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
+
2317
2483
Value *PrimitiveShadow;
2318
2484
if (ClCombinePointerLabelsOnStore) {
2319
2485
Value *PtrShadow = DFSF.getShadow (SI.getPointerOperand ());
2486
+ if (ShouldTrackOrigins) {
2487
+ Shadows.push_back (PtrShadow);
2488
+ Origins.push_back (DFSF.getOrigin (SI.getPointerOperand ()));
2489
+ }
2320
2490
PrimitiveShadow = DFSF.combineShadows (Shadow, PtrShadow, &SI);
2321
2491
} else {
2322
2492
PrimitiveShadow = DFSF.collapseToPrimitiveShadow (Shadow, &SI);
2323
2493
}
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);
2326
2500
if (ClEventCallbacks) {
2327
2501
IRBuilder<> IRB (&SI);
2328
2502
Value *Addr8 = IRB.CreateBitCast (SI.getPointerOperand (), DFSF.DFS .Int8Ptr );
0 commit comments