[dfsan] Tracking origins at phi nodes

[stephan.yichao.zhao]

Summary

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

[morehouse]

Can we add an end-to-end test as well?

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

1572–1580

Nit: range for loop is cleaner if we don't need to advance iterators ourselves.

2713

This is unrelated and shouldn't be in this diff...

3141–3142

Nit: range for loop is cleaner

[stephan.yichao.zhao]

used range loops

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

2713

rebased after D98192.

[morehouse]

Can we add an end-to-end test to compiler-rt?

[stephan.yichao.zhao]

added a compiler-rt test

[morehouse]

LGTM

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

struct DFSanFunction {
   bool IsNativeABI;
   AllocaInst *LabelReturnAlloca = nullptr;
   AllocaInst *OriginReturnAlloca = nullptr;
   DenseMap<Value *, Value *> ValShadowMap;
   DenseMap<Value *, Value *> ValOriginMap;
   DenseMap<AllocaInst *, AllocaInst *> AllocaShadowMap;
   DenseMap<AllocaInst *, AllocaInst *> AllocaOriginMap;
 
-  std::vector<std::pair<PHINode *, PHINode *>> PHIFixups;
+  struct PHIFixupElement {
+    PHINode *Phi;
+    PHINode *ShadowPhi;
+    PHINode *OriginPhi;
+  };
+  std::vector<PHIFixupElement> PHIFixups;
+
   DenseSet<Instruction *> SkipInsts;
   std::vector<Value *> NonZeroChecks;
   bool AvoidNewBlocks;
 
   struct CachedShadow {
     BasicBlock *Block; // The block where Shadow is defined.
     Value *Shadow;
   };
for (BasicBlock *BB : BBList) {
         Inst = Next;
       }
     }
 
     // We will not necessarily be able to compute the shadow for every phi node
     // until we have visited every block.  Therefore, the code that handles phi
     // nodes adds them to the PHIFixups list so that they can be properly
     // handled here.
-    for (auto PHIFixup : DFSF.PHIFixups) {
-      PHINode *PN, *ShadowPN;
-      std::tie(PN, ShadowPN) = PHIFixup;
-      for (unsigned Val = 0, N = PN->getNumIncomingValues(); Val < N; ++Val) {
-        ShadowPN->setIncomingValue(Val,
-                                   DFSF.getShadow(PN->getIncomingValue(Val)));
+    for (DFSanFunction::PHIFixupElement &P : DFSF.PHIFixups) {
+      for (unsigned Val = 0, N = P.Phi->getNumIncomingValues(); Val != N;
+           ++Val) {
+        P.ShadowPhi->setIncomingValue(
+            Val, DFSF.getShadow(P.Phi->getIncomingValue(Val)));
+        if (P.OriginPhi)
+          P.OriginPhi->setIncomingValue(
+              Val, DFSF.getOrigin(P.Phi->getIncomingValue(Val)));
       }
     }
 
     // -dfsan-debug-nonzero-labels will split the CFG in all kinds of crazy
     // places (i.e. instructions in basic blocks we haven't even begun visiting
     // yet).  To make our life easier, do this work in a pass after the main
     // instrumentation.
     if (ClDebugNonzeroLabels) {
       for (Value *V : DFSF.NonZeroChecks) {
void DFSanVisitor::visitMemTransferInst(MemTransferInst &I) {
   // need to move origins before moving shadows.
   if (DFSF.DFS.shouldTrackOrigins()) {
     IRB.CreateCall(
         DFSF.DFS.DFSanMemOriginTransferFn,
         {IRB.CreatePointerCast(I.getArgOperand(0), IRB.getInt8PtrTy()),
          IRB.CreatePointerCast(I.getArgOperand(1), IRB.getInt8PtrTy()),
          IRB.CreateIntCast(I.getArgOperand(2), DFSF.DFS.IntptrTy, false)});
   }
 
   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(),
                                                     DFSF.DFS.ShadowWidthBytes));
   Type *Int8Ptr = Type::getInt8PtrTy(*DFSF.DFS.Ctx);
   Value *DestShadow = IRB.CreateBitCast(RawDestShadow, Int8Ptr);
   SrcShadow = IRB.CreateBitCast(SrcShadow, Int8Ptr);
void DFSanVisitor::visitPHINode(PHINode &PN) {
   PHINode *ShadowPN =
       PHINode::Create(ShadowTy, PN.getNumIncomingValues(), "", &PN);
 
   // Give the shadow phi node valid predecessors to fool SplitEdge into working.
   Value *UndefShadow = UndefValue::get(ShadowTy);
   for (BasicBlock *BB : PN.blocks())
     ShadowPN->addIncoming(UndefShadow, BB);
 
-  DFSF.PHIFixups.push_back(std::make_pair(&PN, ShadowPN));
   DFSF.setShadow(&PN, ShadowPN);
+
+  PHINode *OriginPN = nullptr;
+  if (DFSF.DFS.shouldTrackOrigins()) {
+    OriginPN =
+        PHINode::Create(DFSF.DFS.OriginTy, PN.getNumIncomingValues(), "", &PN);
+    Value *UndefOrigin = UndefValue::get(DFSF.DFS.OriginTy);
+    for (BasicBlock *BB : PN.blocks())
+      OriginPN->addIncoming(UndefOrigin, BB);
+    DFSF.setOrigin(&PN, OriginPN);
+  }
+
+  DFSF.PHIFixups.push_back({&PN, ShadowPN, OriginPN});
 }
 
 namespace {
 class DataFlowSanitizerLegacyPass : public ModulePass {
 private:
   std::vector<std::string> ABIListFiles;
 
 public:

llvm/test/Instrumentation/DataFlowSanitizer/origin_cached_shadows.ll

+; RUN: opt < %s -dfsan -dfsan-track-origins=1 -dfsan-fast-16-labels=true -S | FileCheck %s --check-prefix=CHECK
+;
+; %15 and %17 have the same key in shadow cache. They should not reuse the same
+; shadow because their blocks do not dominate each other. Origin tracking
+; splt blocks. This test ensures DT is updated correctly, and cached shadows
+; are not mis-used.
+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 [[TLS_ARR:\[100 x i64\]]]
+; CHECK: @__dfsan_shadow_width_bits = weak_odr constant i32 [[#SBITS:]]
+; CHECK: @__dfsan_shadow_width_bytes = weak_odr constant i32 [[#SBYTES:]]
+
+define void @cached_shadows(double %0) {
+  ; CHECK: @"dfs$cached_shadows"
+  ; CHECK:  [[AO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 0), align 4
+  ; CHECK:  [[AS:%.*]] = load i[[#SBITS]], i[[#SBITS]]* bitcast ([[TLS_ARR]]* @__dfsan_arg_tls to i[[#SBITS]]*), align [[#SBYTES]]
+  ; CHECK: [[L1:[0-9]+]]:
+  ; CHECK:  {{.*}} = phi i[[#SBITS]]
+  ; CHECK:  {{.*}} = phi i32
+  ; CHECK:  {{.*}} = phi double [ 3.000000e+00
+  ; CHECK:  [[S_L1:%.*]] = phi i[[#SBITS]] [ 0, %[[L0:[0-9]+]] ], [ [[S_L7:%.*]], %[[L7:[0-9]+]] ]
+  ; CHECK:  [[O_L1:%.*]] = phi i32 [ 0, %[[L0]] ], [ [[O_L7:%.*]], %[[L7]] ]
+  ; CHECK:  [[V_L1:%.*]] = phi double [ 4.000000e+00, %[[L0]] ], [ [[V_L7:%.*]], %[[L7]] ]
+  ; CHECK:  br i1 {{%[0-9]+}}, label %[[L2:[0-9]+]], label %[[L4:[0-9]+]]
+  ; CHECK: [[L2]]:
+  ; CHECK:  br i1 {{%[0-9]+}}, label %[[L3:[0-9]+]], label %[[L7]]
+  ; CHECK: [[L3]]:
+  ; CHECK:  [[S_L3:%.*]] = or i[[#SBITS]]
+  ; CHECK:  [[AS_NE_L3:%.*]] = icmp ne i[[#SBITS]] [[AS]], 0
+  ; CHECK:  [[O_L3:%.*]] = select i1 [[AS_NE_L3]], i32 %2, i32 [[O_L1]]
+  ; CHECK:  [[V_L3:%.*]] = fsub double [[V_L1]], %0
+  ; CHECK:  br label %[[L7]]
+  ; CHECK: [[L4]]:
+  ; CHECK:  br i1 %_dfscmp, label %[[L5:[0-9]+]], label %[[L6:[0-9]+]]
+  ; CHECK: [[L5]]:
+  ; CHECK:  br label %[[L6]]
+  ; CHECK: [[L6]]:
+  ; CHECK:  [[S_L6:%.*]] = or i[[#SBITS]]
+  ; CHECK:  [[AS_NE_L6:%.*]] = icmp ne i[[#SBITS]] [[AS]], 0
+  ; CHECK:  [[O_L6:%.*]] = select i1 [[AS_NE_L6]], i32 [[AO]], i32 [[O_L1]]
+  ; CHECK:  [[V_L6:%.*]] = fadd double %24, %0
+  ; CHECK:  br label %[[L7]]
+  ; CHECK: [[L7]]:
+  ; CHECK:  [[S_L7]] = phi i[[#SBITS]] [ [[S_L3]], %[[L3]] ], [ [[S_L1]], %[[L2]] ], [ [[S_L6]], %[[L6]] ]
+  ; CHECK:  [[O_L7]] = phi i32 [ [[O_L3]], %[[L3]] ], [ [[O_L1]], %[[L2]] ], [ [[O_L6]], %[[L6]] ]
+  ; CHECK:  [[V_L7]] = phi double [ [[V_L3]], %[[L3]] ], [ [[V_L1]], %[[L2]] ], [ [[V_L6]], %[[L6]] ]
+  ; CHECK:  br i1 {{%[0-9]+}}, label %[[L1]], label %[[L8:[0-9]+]]
+  ; CHECK: [[L8]]:
+  
+  %2 = alloca double, align 8
+  %3 = alloca double, align 8
+  %4 = bitcast double* %2 to i8*
+  store volatile double 1.000000e+00, double* %2, align 8
+  %5 = bitcast double* %3 to i8*
+  store volatile double 2.000000e+00, double* %3, align 8
+  br label %6
+
+6:                                                ; preds = %18, %1
+  %7 = phi double [ 3.000000e+00, %1 ], [ %19, %18 ]
+  %8 = phi double [ 4.000000e+00, %1 ], [ %20, %18 ]
+  %9 = load volatile double, double* %3, align 8
+  %10 = fcmp une double %9, 0.000000e+00
+  %11 = load volatile double, double* %3, align 8
+  br i1 %10, label %12, label %16
+
+12:                                               ; preds = %6
+  %13 = fcmp une double %11, 0.000000e+00
+  br i1 %13, label %14, label %18
+
+14:                                               ; preds = %12
+  %15 = fsub double %8, %0
+  br label %18
+
+16:                                               ; preds = %6
+  store volatile double %11, double* %2, align 8
+  %17 = fadd double %8, %0
+  br label %18
+
+18:                                               ; preds = %16, %14, %12
+  %19 = phi double [ %8, %14 ], [ %7, %12 ], [ %8, %16 ]
+  %20 = phi double [ %15, %14 ], [ %8, %12 ], [ %17, %16 ]
+  %21 = fcmp olt double %19, 9.900000e+01
+  br i1 %21, label %6, label %22
+
+22:                                               ; preds = %18
+  ret void
+}
 No newline at end of file

llvm/test/Instrumentation/DataFlowSanitizer/origin_phi.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"
+
+; CHECK: @__dfsan_arg_tls = external thread_local(initialexec) global [[TLS_ARR:\[100 x i64\]]]
+; CHECK: @__dfsan_shadow_width_bits = weak_odr constant i32 [[#SBITS:]]
+; CHECK: @__dfsan_shadow_width_bytes = weak_odr constant i32 [[#SBYTES:]]
+
+define i32 @phiop(i32 %a, i32 %b, i1 %c) {
+  ; CHECK: @"dfs$phiop"
+  ; CHECK: entry:
+  ; CHECK: [[BO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 1), align 4
+  ; CHECK: [[AO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 0), align 4
+  ; CHECK: [[BS:%.*]] = load i[[#SBITS]], i[[#SBITS]]* inttoptr (i64 add (i64 ptrtoint ([[TLS_ARR]]* @__dfsan_arg_tls to i64), i64 2) to i[[#SBITS]]*), align [[#SBYTES]]
+  ; CHECK: [[AS:%.*]] = load i[[#SBITS]], i[[#SBITS]]* bitcast ([[TLS_ARR]]* @__dfsan_arg_tls to i[[#SBITS]]*), align [[#SBYTES]]
+  ; CHECK: br i1 %c, label %next, label %done
+  ; CHECK: next:
+  ; CHECK: br i1 %c, label %T, label %F
+  ; CHECK: T:
+  ; CHECK: [[BS_NE:%.*]] = icmp ne i[[#SBITS]] [[BS]], 0
+  ; CHECK: [[BAO_T:%.*]] = select i1 [[BS_NE]], i32 [[BO]], i32 [[AO]]
+  ; CHECK: br label %done
+  ; CHECK: F:
+  ; CHECK: [[AS_NE:%.*]] = icmp ne i[[#SBITS]] [[AS]], 0
+  ; CHECK: [[BAO_F:%.*]] = select i1 [[AS_NE]], i32 [[AO]], i32 [[BO]]
+  ; CHECK: br label %done
+  ; CHECK: done:
+  ; CHECK: [[PO:%.*]] = phi i32 [ [[BAO_T]], %T ], [ [[BAO_F]], %F ], [ [[AO]], %entry ]
+  ; CHECK: store i32 [[PO]], i32* @__dfsan_retval_origin_tls, align 4
+
+entry:
+  br i1 %c, label %next, label %done
+next:  
+  br i1 %c, label %T, label %F 
+T:
+  %sum = add i32 %a, %b 
+  br label %done
+F:
+  %diff = sub i32 %b, %a 
+  br label %done
+done:
+  %r = phi i32 [%sum, %T], [%diff, %F], [%a, %entry]
+  ret i32 %r
+}
 No newline at end of file