DataFlowSanitizer.cpp revision a77d9f726a7e3c51f04d1d74d091ae1a87d63544 
1//===-- DataFlowSanitizer.cpp - dynamic data flow analysis ----------------===//
2//
3//                     The LLVM Compiler Infrastructure
4//
5// This file is distributed under the University of Illinois Open Source
6// License. See LICENSE.TXT for details.
7//
8//===----------------------------------------------------------------------===//
9/// \file
10/// This file is a part of DataFlowSanitizer, a generalised dynamic data flow
11/// analysis.
12///
13/// Unlike other Sanitizer tools, this tool is not designed to detect a specific
14/// class of bugs on its own.  Instead, it provides a generic dynamic data flow
15/// analysis framework to be used by clients to help detect application-specific
16/// issues within their own code.
17///
18/// The analysis is based on automatic propagation of data flow labels (also
19/// known as taint labels) through a program as it performs computation.  Each
20/// byte of application memory is backed by two bytes of shadow memory which
21/// hold the label.  On Linux/x86_64, memory is laid out as follows:
22///
23/// +--------------------+ 0x800000000000 (top of memory)
24/// | application memory |
25/// +--------------------+ 0x700000008000 (kAppAddr)
26/// |                    |
27/// |       unused       |
28/// |                    |
29/// +--------------------+ 0x200200000000 (kUnusedAddr)
30/// |    union table     |
31/// +--------------------+ 0x200000000000 (kUnionTableAddr)
32/// |   shadow memory    |
33/// +--------------------+ 0x000000010000 (kShadowAddr)
34/// | reserved by kernel |
35/// +--------------------+ 0x000000000000
36///
37/// To derive a shadow memory address from an application memory address,
38/// bits 44-46 are cleared to bring the address into the range
39/// [0x000000008000,0x100000000000).  Then the address is shifted left by 1 to
40/// account for the double byte representation of shadow labels and move the
41/// address into the shadow memory range.  See the function
42/// DataFlowSanitizer::getShadowAddress below.
43///
44/// For more information, please refer to the design document:
45/// http://clang.llvm.org/docs/DataFlowSanitizerDesign.html
46
47#include "llvm/Transforms/Instrumentation.h"
48#include "llvm/ADT/DenseMap.h"
49#include "llvm/ADT/DenseSet.h"
50#include "llvm/ADT/DepthFirstIterator.h"
51#include "llvm/Analysis/ValueTracking.h"
52#include "llvm/IR/InlineAsm.h"
53#include "llvm/IR/IRBuilder.h"
54#include "llvm/IR/LLVMContext.h"
55#include "llvm/IR/MDBuilder.h"
56#include "llvm/IR/Type.h"
57#include "llvm/IR/Value.h"
58#include "llvm/InstVisitor.h"
59#include "llvm/Pass.h"
60#include "llvm/Support/CommandLine.h"
61#include "llvm/Transforms/Utils/BasicBlockUtils.h"
62#include "llvm/Transforms/Utils/Local.h"
63#include "llvm/Transforms/Utils/SpecialCaseList.h"
64#include <iterator>
65
66using namespace llvm;
67
68// The -dfsan-preserve-alignment flag controls whether this pass assumes that
69// alignment requirements provided by the input IR are correct.  For example,
70// if the input IR contains a load with alignment 8, this flag will cause
71// the shadow load to have alignment 16.  This flag is disabled by default as
72// we have unfortunately encountered too much code (including Clang itself;
73// see PR14291) which performs misaligned access.
74static cl::opt<bool> ClPreserveAlignment(
75    "dfsan-preserve-alignment",
76    cl::desc("respect alignment requirements provided by input IR"), cl::Hidden,
77    cl::init(false));
78
79// The ABI list file controls how shadow parameters are passed.  The pass treats
80// every function labelled "uninstrumented" in the ABI list file as conforming
81// to the "native" (i.e. unsanitized) ABI.  Unless the ABI list contains
82// additional annotations for those functions, a call to one of those functions
83// will produce a warning message, as the labelling behaviour of the function is
84// unknown.  The other supported annotations are "functional" and "discard",
85// which are described below under DataFlowSanitizer::WrapperKind.
86static cl::opt<std::string> ClABIListFile(
87    "dfsan-abilist",
88    cl::desc("File listing native ABI functions and how the pass treats them"),
89    cl::Hidden);
90
91// Controls whether the pass uses IA_Args or IA_TLS as the ABI for instrumented
92// functions (see DataFlowSanitizer::InstrumentedABI below).
93static cl::opt<bool> ClArgsABI(
94    "dfsan-args-abi",
95    cl::desc("Use the argument ABI rather than the TLS ABI"),
96    cl::Hidden);
97
98static cl::opt<bool> ClDebugNonzeroLabels(
99    "dfsan-debug-nonzero-labels",
100    cl::desc("Insert calls to __dfsan_nonzero_label on observing a parameter, "
101             "load or return with a nonzero label"),
102    cl::Hidden);
103
104namespace {
105
106class DataFlowSanitizer : public ModulePass {
107  friend struct DFSanFunction;
108  friend class DFSanVisitor;
109
110  enum {
111    ShadowWidth = 16
112  };
113
114  /// Which ABI should be used for instrumented functions?
115  enum InstrumentedABI {
116    /// Argument and return value labels are passed through additional
117    /// arguments and by modifying the return type.
118    IA_Args,
119
120    /// Argument and return value labels are passed through TLS variables
121    /// __dfsan_arg_tls and __dfsan_retval_tls.
122    IA_TLS
123  };
124
125  /// How should calls to uninstrumented functions be handled?
126  enum WrapperKind {
127    /// This function is present in an uninstrumented form but we don't know
128    /// how it should be handled.  Print a warning and call the function anyway.
129    /// Don't label the return value.
130    WK_Warning,
131
132    /// This function does not write to (user-accessible) memory, and its return
133    /// value is unlabelled.
134    WK_Discard,
135
136    /// This function does not write to (user-accessible) memory, and the label
137    /// of its return value is the union of the label of its arguments.
138    WK_Functional,
139
140    /// Instead of calling the function, a custom wrapper __dfsw_F is called,
141    /// where F is the name of the function.  This function may wrap the
142    /// original function or provide its own implementation.  This is similar to
143    /// the IA_Args ABI, except that IA_Args uses a struct return type to
144    /// pass the return value shadow in a register, while WK_Custom uses an
145    /// extra pointer argument to return the shadow.  This allows the wrapped
146    /// form of the function type to be expressed in C.
147    WK_Custom
148  };
149
150  DataLayout *DL;
151  Module *Mod;
152  LLVMContext *Ctx;
153  IntegerType *ShadowTy;
154  PointerType *ShadowPtrTy;
155  IntegerType *IntptrTy;
156  ConstantInt *ZeroShadow;
157  ConstantInt *ShadowPtrMask;
158  ConstantInt *ShadowPtrMul;
159  Constant *ArgTLS;
160  Constant *RetvalTLS;
161  void *(*GetArgTLSPtr)();
162  void *(*GetRetvalTLSPtr)();
163  Constant *GetArgTLS;
164  Constant *GetRetvalTLS;
165  FunctionType *DFSanUnionFnTy;
166  FunctionType *DFSanUnionLoadFnTy;
167  FunctionType *DFSanUnimplementedFnTy;
168  FunctionType *DFSanSetLabelFnTy;
169  FunctionType *DFSanNonzeroLabelFnTy;
170  Constant *DFSanUnionFn;
171  Constant *DFSanUnionLoadFn;
172  Constant *DFSanUnimplementedFn;
173  Constant *DFSanSetLabelFn;
174  Constant *DFSanNonzeroLabelFn;
175  MDNode *ColdCallWeights;
176  OwningPtr<SpecialCaseList> ABIList;
177  DenseMap<Value *, Function *> UnwrappedFnMap;
178  AttributeSet ReadOnlyNoneAttrs;
179
180  Value *getShadowAddress(Value *Addr, Instruction *Pos);
181  Value *combineShadows(Value *V1, Value *V2, Instruction *Pos);
182  bool isInstrumented(Function *F);
183  FunctionType *getArgsFunctionType(FunctionType *T);
184  FunctionType *getCustomFunctionType(FunctionType *T);
185  InstrumentedABI getInstrumentedABI();
186  WrapperKind getWrapperKind(Function *F);
187
188 public:
189  DataFlowSanitizer(StringRef ABIListFile = StringRef(),
190                    void *(*getArgTLS)() = 0, void *(*getRetValTLS)() = 0);
191  static char ID;
192  bool doInitialization(Module &M);
193  bool runOnModule(Module &M);
194};
195
196struct DFSanFunction {
197  DataFlowSanitizer &DFS;
198  Function *F;
199  DataFlowSanitizer::InstrumentedABI IA;
200  bool IsNativeABI;
201  Value *ArgTLSPtr;
202  Value *RetvalTLSPtr;
203  AllocaInst *LabelReturnAlloca;
204  DenseMap<Value *, Value *> ValShadowMap;
205  DenseMap<AllocaInst *, AllocaInst *> AllocaShadowMap;
206  std::vector<std::pair<PHINode *, PHINode *> > PHIFixups;
207  DenseSet<Instruction *> SkipInsts;
208  DenseSet<Value *> NonZeroChecks;
209
210  DFSanFunction(DataFlowSanitizer &DFS, Function *F, bool IsNativeABI)
211      : DFS(DFS), F(F), IA(DFS.getInstrumentedABI()),
212        IsNativeABI(IsNativeABI), ArgTLSPtr(0), RetvalTLSPtr(0),
213        LabelReturnAlloca(0) {}
214  Value *getArgTLSPtr();
215  Value *getArgTLS(unsigned Index, Instruction *Pos);
216  Value *getRetvalTLS();
217  Value *getShadow(Value *V);
218  void setShadow(Instruction *I, Value *Shadow);
219  Value *combineOperandShadows(Instruction *Inst);
220  Value *loadShadow(Value *ShadowAddr, uint64_t Size, uint64_t Align,
221                    Instruction *Pos);
222  void storeShadow(Value *Addr, uint64_t Size, uint64_t Align, Value *Shadow,
223                   Instruction *Pos);
224};
225
226class DFSanVisitor : public InstVisitor<DFSanVisitor> {
227 public:
228  DFSanFunction &DFSF;
229  DFSanVisitor(DFSanFunction &DFSF) : DFSF(DFSF) {}
230
231  void visitOperandShadowInst(Instruction &I);
232
233  void visitBinaryOperator(BinaryOperator &BO);
234  void visitCastInst(CastInst &CI);
235  void visitCmpInst(CmpInst &CI);
236  void visitGetElementPtrInst(GetElementPtrInst &GEPI);
237  void visitLoadInst(LoadInst &LI);
238  void visitStoreInst(StoreInst &SI);
239  void visitReturnInst(ReturnInst &RI);
240  void visitCallSite(CallSite CS);
241  void visitPHINode(PHINode &PN);
242  void visitExtractElementInst(ExtractElementInst &I);
243  void visitInsertElementInst(InsertElementInst &I);
244  void visitShuffleVectorInst(ShuffleVectorInst &I);
245  void visitExtractValueInst(ExtractValueInst &I);
246  void visitInsertValueInst(InsertValueInst &I);
247  void visitAllocaInst(AllocaInst &I);
248  void visitSelectInst(SelectInst &I);
249  void visitMemSetInst(MemSetInst &I);
250  void visitMemTransferInst(MemTransferInst &I);
251};
252
253}
254
255char DataFlowSanitizer::ID;
256INITIALIZE_PASS(DataFlowSanitizer, "dfsan",
257                "DataFlowSanitizer: dynamic data flow analysis.", false, false)
258
259ModulePass *llvm::createDataFlowSanitizerPass(StringRef ABIListFile,
260                                              void *(*getArgTLS)(),
261                                              void *(*getRetValTLS)()) {
262  return new DataFlowSanitizer(ABIListFile, getArgTLS, getRetValTLS);
263}
264
265DataFlowSanitizer::DataFlowSanitizer(StringRef ABIListFile,
266                                     void *(*getArgTLS)(),
267                                     void *(*getRetValTLS)())
268    : ModulePass(ID), GetArgTLSPtr(getArgTLS), GetRetvalTLSPtr(getRetValTLS),
269      ABIList(SpecialCaseList::createOrDie(ABIListFile.empty() ? ClABIListFile
270                                                               : ABIListFile)) {
271}
272
273FunctionType *DataFlowSanitizer::getArgsFunctionType(FunctionType *T) {
274  llvm::SmallVector<Type *, 4> ArgTypes;
275  std::copy(T->param_begin(), T->param_end(), std::back_inserter(ArgTypes));
276  for (unsigned i = 0, e = T->getNumParams(); i != e; ++i)
277    ArgTypes.push_back(ShadowTy);
278  if (T->isVarArg())
279    ArgTypes.push_back(ShadowPtrTy);
280  Type *RetType = T->getReturnType();
281  if (!RetType->isVoidTy())
282    RetType = StructType::get(RetType, ShadowTy, (Type *)0);
283  return FunctionType::get(RetType, ArgTypes, T->isVarArg());
284}
285
286FunctionType *DataFlowSanitizer::getCustomFunctionType(FunctionType *T) {
287  assert(!T->isVarArg());
288  llvm::SmallVector<Type *, 4> ArgTypes;
289  std::copy(T->param_begin(), T->param_end(), std::back_inserter(ArgTypes));
290  for (unsigned i = 0, e = T->getNumParams(); i != e; ++i)
291    ArgTypes.push_back(ShadowTy);
292  Type *RetType = T->getReturnType();
293  if (!RetType->isVoidTy())
294    ArgTypes.push_back(ShadowPtrTy);
295  return FunctionType::get(T->getReturnType(), ArgTypes, false);
296}
297
298bool DataFlowSanitizer::doInitialization(Module &M) {
299  DL = getAnalysisIfAvailable<DataLayout>();
300  if (!DL)
301    return false;
302
303  Mod = &M;
304  Ctx = &M.getContext();
305  ShadowTy = IntegerType::get(*Ctx, ShadowWidth);
306  ShadowPtrTy = PointerType::getUnqual(ShadowTy);
307  IntptrTy = DL->getIntPtrType(*Ctx);
308  ZeroShadow = ConstantInt::getSigned(ShadowTy, 0);
309  ShadowPtrMask = ConstantInt::getSigned(IntptrTy, ~0x700000000000LL);
310  ShadowPtrMul = ConstantInt::getSigned(IntptrTy, ShadowWidth / 8);
311
312  Type *DFSanUnionArgs[2] = { ShadowTy, ShadowTy };
313  DFSanUnionFnTy =
314      FunctionType::get(ShadowTy, DFSanUnionArgs, /*isVarArg=*/ false);
315  Type *DFSanUnionLoadArgs[2] = { ShadowPtrTy, IntptrTy };
316  DFSanUnionLoadFnTy =
317      FunctionType::get(ShadowTy, DFSanUnionLoadArgs, /*isVarArg=*/ false);
318  DFSanUnimplementedFnTy = FunctionType::get(
319      Type::getVoidTy(*Ctx), Type::getInt8PtrTy(*Ctx), /*isVarArg=*/false);
320  Type *DFSanSetLabelArgs[3] = { ShadowTy, Type::getInt8PtrTy(*Ctx), IntptrTy };
321  DFSanSetLabelFnTy = FunctionType::get(Type::getVoidTy(*Ctx),
322                                        DFSanSetLabelArgs, /*isVarArg=*/false);
323  DFSanNonzeroLabelFnTy = FunctionType::get(
324      Type::getVoidTy(*Ctx), ArrayRef<Type *>(), /*isVarArg=*/false);
325
326  if (GetArgTLSPtr) {
327    Type *ArgTLSTy = ArrayType::get(ShadowTy, 64);
328    ArgTLS = 0;
329    GetArgTLS = ConstantExpr::getIntToPtr(
330        ConstantInt::get(IntptrTy, uintptr_t(GetArgTLSPtr)),
331        PointerType::getUnqual(
332            FunctionType::get(PointerType::getUnqual(ArgTLSTy), (Type *)0)));
333  }
334  if (GetRetvalTLSPtr) {
335    RetvalTLS = 0;
336    GetRetvalTLS = ConstantExpr::getIntToPtr(
337        ConstantInt::get(IntptrTy, uintptr_t(GetRetvalTLSPtr)),
338        PointerType::getUnqual(
339            FunctionType::get(PointerType::getUnqual(ShadowTy), (Type *)0)));
340  }
341
342  ColdCallWeights = MDBuilder(*Ctx).createBranchWeights(1, 1000);
343  return true;
344}
345
346bool DataFlowSanitizer::isInstrumented(Function *F) {
347  return !ABIList->isIn(*F, "uninstrumented");
348}
349
350DataFlowSanitizer::InstrumentedABI DataFlowSanitizer::getInstrumentedABI() {
351  return ClArgsABI ? IA_Args : IA_TLS;
352}
353
354DataFlowSanitizer::WrapperKind DataFlowSanitizer::getWrapperKind(Function *F) {
355  if (ABIList->isIn(*F, "functional"))
356    return WK_Functional;
357  if (ABIList->isIn(*F, "discard"))
358    return WK_Discard;
359  if (ABIList->isIn(*F, "custom"))
360    return WK_Custom;
361
362  return WK_Warning;
363}
364
365bool DataFlowSanitizer::runOnModule(Module &M) {
366  if (!DL)
367    return false;
368
369  if (ABIList->isIn(M, "skip"))
370    return false;
371
372  if (!GetArgTLSPtr) {
373    Type *ArgTLSTy = ArrayType::get(ShadowTy, 64);
374    ArgTLS = Mod->getOrInsertGlobal("__dfsan_arg_tls", ArgTLSTy);
375    if (GlobalVariable *G = dyn_cast<GlobalVariable>(ArgTLS))
376      G->setThreadLocalMode(GlobalVariable::InitialExecTLSModel);
377  }
378  if (!GetRetvalTLSPtr) {
379    RetvalTLS = Mod->getOrInsertGlobal("__dfsan_retval_tls", ShadowTy);
380    if (GlobalVariable *G = dyn_cast<GlobalVariable>(RetvalTLS))
381      G->setThreadLocalMode(GlobalVariable::InitialExecTLSModel);
382  }
383
384  DFSanUnionFn = Mod->getOrInsertFunction("__dfsan_union", DFSanUnionFnTy);
385  if (Function *F = dyn_cast<Function>(DFSanUnionFn)) {
386    F->addAttribute(AttributeSet::FunctionIndex, Attribute::ReadNone);
387    F->addAttribute(AttributeSet::ReturnIndex, Attribute::ZExt);
388    F->addAttribute(1, Attribute::ZExt);
389    F->addAttribute(2, Attribute::ZExt);
390  }
391  DFSanUnionLoadFn =
392      Mod->getOrInsertFunction("__dfsan_union_load", DFSanUnionLoadFnTy);
393  if (Function *F = dyn_cast<Function>(DFSanUnionLoadFn)) {
394    F->addAttribute(AttributeSet::ReturnIndex, Attribute::ZExt);
395  }
396  DFSanUnimplementedFn =
397      Mod->getOrInsertFunction("__dfsan_unimplemented", DFSanUnimplementedFnTy);
398  DFSanSetLabelFn =
399      Mod->getOrInsertFunction("__dfsan_set_label", DFSanSetLabelFnTy);
400  if (Function *F = dyn_cast<Function>(DFSanSetLabelFn)) {
401    F->addAttribute(1, Attribute::ZExt);
402  }
403  DFSanNonzeroLabelFn =
404      Mod->getOrInsertFunction("__dfsan_nonzero_label", DFSanNonzeroLabelFnTy);
405
406  std::vector<Function *> FnsToInstrument;
407  llvm::SmallPtrSet<Function *, 2> FnsWithNativeABI;
408  for (Module::iterator i = M.begin(), e = M.end(); i != e; ++i) {
409    if (!i->isIntrinsic() &&
410        i != DFSanUnionFn &&
411        i != DFSanUnionLoadFn &&
412        i != DFSanUnimplementedFn &&
413        i != DFSanSetLabelFn &&
414        i != DFSanNonzeroLabelFn)
415      FnsToInstrument.push_back(&*i);
416  }
417
418  AttrBuilder B;
419  B.addAttribute(Attribute::ReadOnly).addAttribute(Attribute::ReadNone);
420  ReadOnlyNoneAttrs = AttributeSet::get(*Ctx, AttributeSet::FunctionIndex, B);
421
422  // First, change the ABI of every function in the module.  ABI-listed
423  // functions keep their original ABI and get a wrapper function.
424  for (std::vector<Function *>::iterator i = FnsToInstrument.begin(),
425                                         e = FnsToInstrument.end();
426       i != e; ++i) {
427    Function &F = **i;
428    FunctionType *FT = F.getFunctionType();
429
430    if (FT->getNumParams() == 0 && !FT->isVarArg() &&
431        FT->getReturnType()->isVoidTy())
432      continue;
433
434    if (isInstrumented(&F)) {
435      if (getInstrumentedABI() == IA_Args) {
436        FunctionType *NewFT = getArgsFunctionType(FT);
437        Function *NewF = Function::Create(NewFT, F.getLinkage(), "", &M);
438        NewF->copyAttributesFrom(&F);
439        NewF->removeAttributes(
440            AttributeSet::ReturnIndex,
441            AttributeFuncs::typeIncompatible(NewFT->getReturnType(),
442                                             AttributeSet::ReturnIndex));
443        for (Function::arg_iterator FArg = F.arg_begin(),
444                                    NewFArg = NewF->arg_begin(),
445                                    FArgEnd = F.arg_end();
446             FArg != FArgEnd; ++FArg, ++NewFArg) {
447          FArg->replaceAllUsesWith(NewFArg);
448        }
449        NewF->getBasicBlockList().splice(NewF->begin(), F.getBasicBlockList());
450
451        for (Function::use_iterator ui = F.use_begin(), ue = F.use_end();
452             ui != ue;) {
453          BlockAddress *BA = dyn_cast<BlockAddress>(ui.getUse().getUser());
454          ++ui;
455          if (BA) {
456            BA->replaceAllUsesWith(
457                BlockAddress::get(NewF, BA->getBasicBlock()));
458            delete BA;
459          }
460        }
461        F.replaceAllUsesWith(
462            ConstantExpr::getBitCast(NewF, PointerType::getUnqual(FT)));
463        NewF->takeName(&F);
464        F.eraseFromParent();
465        *i = NewF;
466      }
467               // Hopefully, nobody will try to indirectly call a vararg
468               // function... yet.
469    } else if (FT->isVarArg()) {
470      UnwrappedFnMap[&F] = &F;
471      *i = 0;
472    } else {
473      // Build a wrapper function for F.  The wrapper simply calls F, and is
474      // added to FnsToInstrument so that any instrumentation according to its
475      // WrapperKind is done in the second pass below.
476      FunctionType *NewFT = getInstrumentedABI() == IA_Args
477                                ? getArgsFunctionType(FT)
478                                : FT;
479      Function *NewF =
480          Function::Create(NewFT, GlobalValue::LinkOnceODRLinkage,
481                           std::string("dfsw$") + F.getName(), &M);
482      NewF->copyAttributesFrom(&F);
483      NewF->removeAttributes(
484              AttributeSet::ReturnIndex,
485              AttributeFuncs::typeIncompatible(NewFT->getReturnType(),
486                                               AttributeSet::ReturnIndex));
487      if (getInstrumentedABI() == IA_TLS)
488        NewF->removeAttributes(AttributeSet::FunctionIndex,
489                               ReadOnlyNoneAttrs);
490
491      BasicBlock *BB = BasicBlock::Create(*Ctx, "entry", NewF);
492      std::vector<Value *> Args;
493      unsigned n = FT->getNumParams();
494      for (Function::arg_iterator ai = NewF->arg_begin(); n != 0; ++ai, --n)
495        Args.push_back(&*ai);
496      CallInst *CI = CallInst::Create(&F, Args, "", BB);
497      if (FT->getReturnType()->isVoidTy())
498        ReturnInst::Create(*Ctx, BB);
499      else
500        ReturnInst::Create(*Ctx, CI, BB);
501
502      Value *WrappedFnCst =
503          ConstantExpr::getBitCast(NewF, PointerType::getUnqual(FT));
504      F.replaceAllUsesWith(WrappedFnCst);
505      UnwrappedFnMap[WrappedFnCst] = &F;
506      *i = NewF;
507
508      if (!F.isDeclaration()) {
509        // This function is probably defining an interposition of an
510        // uninstrumented function and hence needs to keep the original ABI.
511        // But any functions it may call need to use the instrumented ABI, so
512        // we instrument it in a mode which preserves the original ABI.
513        FnsWithNativeABI.insert(&F);
514
515        // This code needs to rebuild the iterators, as they may be invalidated
516        // by the push_back, taking care that the new range does not include
517        // any functions added by this code.
518        size_t N = i - FnsToInstrument.begin(),
519               Count = e - FnsToInstrument.begin();
520        FnsToInstrument.push_back(&F);
521        i = FnsToInstrument.begin() + N;
522        e = FnsToInstrument.begin() + Count;
523      }
524    }
525  }
526
527  for (std::vector<Function *>::iterator i = FnsToInstrument.begin(),
528                                         e = FnsToInstrument.end();
529       i != e; ++i) {
530    if (!*i || (*i)->isDeclaration())
531      continue;
532
533    removeUnreachableBlocks(**i);
534
535    DFSanFunction DFSF(*this, *i, FnsWithNativeABI.count(*i));
536
537    // DFSanVisitor may create new basic blocks, which confuses df_iterator.
538    // Build a copy of the list before iterating over it.
539    llvm::SmallVector<BasicBlock *, 4> BBList;
540    std::copy(df_begin(&(*i)->getEntryBlock()), df_end(&(*i)->getEntryBlock()),
541              std::back_inserter(BBList));
542
543    for (llvm::SmallVector<BasicBlock *, 4>::iterator i = BBList.begin(),
544                                                      e = BBList.end();
545         i != e; ++i) {
546      Instruction *Inst = &(*i)->front();
547      while (1) {
548        // DFSanVisitor may split the current basic block, changing the current
549        // instruction's next pointer and moving the next instruction to the
550        // tail block from which we should continue.
551        Instruction *Next = Inst->getNextNode();
552        // DFSanVisitor may delete Inst, so keep track of whether it was a
553        // terminator.
554        bool IsTerminator = isa<TerminatorInst>(Inst);
555        if (!DFSF.SkipInsts.count(Inst))
556          DFSanVisitor(DFSF).visit(Inst);
557        if (IsTerminator)
558          break;
559        Inst = Next;
560      }
561    }
562
563    // We will not necessarily be able to compute the shadow for every phi node
564    // until we have visited every block.  Therefore, the code that handles phi
565    // nodes adds them to the PHIFixups list so that they can be properly
566    // handled here.
567    for (std::vector<std::pair<PHINode *, PHINode *> >::iterator
568             i = DFSF.PHIFixups.begin(),
569             e = DFSF.PHIFixups.end();
570         i != e; ++i) {
571      for (unsigned val = 0, n = i->first->getNumIncomingValues(); val != n;
572           ++val) {
573        i->second->setIncomingValue(
574            val, DFSF.getShadow(i->first->getIncomingValue(val)));
575      }
576    }
577
578    // -dfsan-debug-nonzero-labels will split the CFG in all kinds of crazy
579    // places (i.e. instructions in basic blocks we haven't even begun visiting
580    // yet).  To make our life easier, do this work in a pass after the main
581    // instrumentation.
582    if (ClDebugNonzeroLabels) {
583      for (DenseSet<Value *>::iterator i = DFSF.NonZeroChecks.begin(),
584                                       e = DFSF.NonZeroChecks.end();
585           i != e; ++i) {
586        Instruction *Pos;
587        if (Instruction *I = dyn_cast<Instruction>(*i))
588          Pos = I->getNextNode();
589        else
590          Pos = DFSF.F->getEntryBlock().begin();
591        while (isa<PHINode>(Pos) || isa<AllocaInst>(Pos))
592          Pos = Pos->getNextNode();
593        IRBuilder<> IRB(Pos);
594        Instruction *NeInst = cast<Instruction>(
595            IRB.CreateICmpNE(*i, DFSF.DFS.ZeroShadow));
596        BranchInst *BI = cast<BranchInst>(SplitBlockAndInsertIfThen(
597            NeInst, /*Unreachable=*/ false, ColdCallWeights));
598        IRBuilder<> ThenIRB(BI);
599        ThenIRB.CreateCall(DFSF.DFS.DFSanNonzeroLabelFn);
600      }
601    }
602  }
603
604  return false;
605}
606
607Value *DFSanFunction::getArgTLSPtr() {
608  if (ArgTLSPtr)
609    return ArgTLSPtr;
610  if (DFS.ArgTLS)
611    return ArgTLSPtr = DFS.ArgTLS;
612
613  IRBuilder<> IRB(F->getEntryBlock().begin());
614  return ArgTLSPtr = IRB.CreateCall(DFS.GetArgTLS);
615}
616
617Value *DFSanFunction::getRetvalTLS() {
618  if (RetvalTLSPtr)
619    return RetvalTLSPtr;
620  if (DFS.RetvalTLS)
621    return RetvalTLSPtr = DFS.RetvalTLS;
622
623  IRBuilder<> IRB(F->getEntryBlock().begin());
624  return RetvalTLSPtr = IRB.CreateCall(DFS.GetRetvalTLS);
625}
626
627Value *DFSanFunction::getArgTLS(unsigned Idx, Instruction *Pos) {
628  IRBuilder<> IRB(Pos);
629  return IRB.CreateConstGEP2_64(getArgTLSPtr(), 0, Idx);
630}
631
632Value *DFSanFunction::getShadow(Value *V) {
633  if (!isa<Argument>(V) && !isa<Instruction>(V))
634    return DFS.ZeroShadow;
635  Value *&Shadow = ValShadowMap[V];
636  if (!Shadow) {
637    if (Argument *A = dyn_cast<Argument>(V)) {
638      if (IsNativeABI)
639        return DFS.ZeroShadow;
640      switch (IA) {
641      case DataFlowSanitizer::IA_TLS: {
642        Value *ArgTLSPtr = getArgTLSPtr();
643        Instruction *ArgTLSPos =
644            DFS.ArgTLS ? &*F->getEntryBlock().begin()
645                       : cast<Instruction>(ArgTLSPtr)->getNextNode();
646        IRBuilder<> IRB(ArgTLSPos);
647        Shadow = IRB.CreateLoad(getArgTLS(A->getArgNo(), ArgTLSPos));
648        break;
649      }
650      case DataFlowSanitizer::IA_Args: {
651        unsigned ArgIdx = A->getArgNo() + F->getArgumentList().size() / 2;
652        Function::arg_iterator i = F->arg_begin();
653        while (ArgIdx--)
654          ++i;
655        Shadow = i;
656        assert(Shadow->getType() == DFS.ShadowTy);
657        break;
658      }
659      }
660      NonZeroChecks.insert(Shadow);
661    } else {
662      Shadow = DFS.ZeroShadow;
663    }
664  }
665  return Shadow;
666}
667
668void DFSanFunction::setShadow(Instruction *I, Value *Shadow) {
669  assert(!ValShadowMap.count(I));
670  assert(Shadow->getType() == DFS.ShadowTy);
671  ValShadowMap[I] = Shadow;
672}
673
674Value *DataFlowSanitizer::getShadowAddress(Value *Addr, Instruction *Pos) {
675  assert(Addr != RetvalTLS && "Reinstrumenting?");
676  IRBuilder<> IRB(Pos);
677  return IRB.CreateIntToPtr(
678      IRB.CreateMul(
679          IRB.CreateAnd(IRB.CreatePtrToInt(Addr, IntptrTy), ShadowPtrMask),
680          ShadowPtrMul),
681      ShadowPtrTy);
682}
683
684// Generates IR to compute the union of the two given shadows, inserting it
685// before Pos.  Returns the computed union Value.
686Value *DataFlowSanitizer::combineShadows(Value *V1, Value *V2,
687                                         Instruction *Pos) {
688  if (V1 == ZeroShadow)
689    return V2;
690  if (V2 == ZeroShadow)
691    return V1;
692  if (V1 == V2)
693    return V1;
694  IRBuilder<> IRB(Pos);
695  BasicBlock *Head = Pos->getParent();
696  Value *Ne = IRB.CreateICmpNE(V1, V2);
697  Instruction *NeInst = dyn_cast<Instruction>(Ne);
698  if (NeInst) {
699    BranchInst *BI = cast<BranchInst>(SplitBlockAndInsertIfThen(
700        NeInst, /*Unreachable=*/ false, ColdCallWeights));
701    IRBuilder<> ThenIRB(BI);
702    CallInst *Call = ThenIRB.CreateCall2(DFSanUnionFn, V1, V2);
703    Call->addAttribute(AttributeSet::ReturnIndex, Attribute::ZExt);
704    Call->addAttribute(1, Attribute::ZExt);
705    Call->addAttribute(2, Attribute::ZExt);
706
707    BasicBlock *Tail = BI->getSuccessor(0);
708    PHINode *Phi = PHINode::Create(ShadowTy, 2, "", Tail->begin());
709    Phi->addIncoming(Call, Call->getParent());
710    Phi->addIncoming(ZeroShadow, Head);
711    Pos = Phi;
712    return Phi;
713  } else {
714    assert(0 && "todo");
715    return 0;
716  }
717}
718
719// A convenience function which folds the shadows of each of the operands
720// of the provided instruction Inst, inserting the IR before Inst.  Returns
721// the computed union Value.
722Value *DFSanFunction::combineOperandShadows(Instruction *Inst) {
723  if (Inst->getNumOperands() == 0)
724    return DFS.ZeroShadow;
725
726  Value *Shadow = getShadow(Inst->getOperand(0));
727  for (unsigned i = 1, n = Inst->getNumOperands(); i != n; ++i) {
728    Shadow = DFS.combineShadows(Shadow, getShadow(Inst->getOperand(i)), Inst);
729  }
730  return Shadow;
731}
732
733void DFSanVisitor::visitOperandShadowInst(Instruction &I) {
734  Value *CombinedShadow = DFSF.combineOperandShadows(&I);
735  DFSF.setShadow(&I, CombinedShadow);
736}
737
738// Generates IR to load shadow corresponding to bytes [Addr, Addr+Size), where
739// Addr has alignment Align, and take the union of each of those shadows.
740Value *DFSanFunction::loadShadow(Value *Addr, uint64_t Size, uint64_t Align,
741                                 Instruction *Pos) {
742  if (AllocaInst *AI = dyn_cast<AllocaInst>(Addr)) {
743    llvm::DenseMap<AllocaInst *, AllocaInst *>::iterator i =
744        AllocaShadowMap.find(AI);
745    if (i != AllocaShadowMap.end()) {
746      IRBuilder<> IRB(Pos);
747      return IRB.CreateLoad(i->second);
748    }
749  }
750
751  uint64_t ShadowAlign = Align * DFS.ShadowWidth / 8;
752  SmallVector<Value *, 2> Objs;
753  GetUnderlyingObjects(Addr, Objs, DFS.DL);
754  bool AllConstants = true;
755  for (SmallVector<Value *, 2>::iterator i = Objs.begin(), e = Objs.end();
756       i != e; ++i) {
757    if (isa<Function>(*i) || isa<BlockAddress>(*i))
758      continue;
759    if (isa<GlobalVariable>(*i) && cast<GlobalVariable>(*i)->isConstant())
760      continue;
761
762    AllConstants = false;
763    break;
764  }
765  if (AllConstants)
766    return DFS.ZeroShadow;
767
768  Value *ShadowAddr = DFS.getShadowAddress(Addr, Pos);
769  switch (Size) {
770  case 0:
771    return DFS.ZeroShadow;
772  case 1: {
773    LoadInst *LI = new LoadInst(ShadowAddr, "", Pos);
774    LI->setAlignment(ShadowAlign);
775    return LI;
776  }
777  case 2: {
778    IRBuilder<> IRB(Pos);
779    Value *ShadowAddr1 =
780        IRB.CreateGEP(ShadowAddr, ConstantInt::get(DFS.IntptrTy, 1));
781    return DFS.combineShadows(IRB.CreateAlignedLoad(ShadowAddr, ShadowAlign),
782                              IRB.CreateAlignedLoad(ShadowAddr1, ShadowAlign),
783                              Pos);
784  }
785  }
786  if (Size % (64 / DFS.ShadowWidth) == 0) {
787    // Fast path for the common case where each byte has identical shadow: load
788    // shadow 64 bits at a time, fall out to a __dfsan_union_load call if any
789    // shadow is non-equal.
790    BasicBlock *FallbackBB = BasicBlock::Create(*DFS.Ctx, "", F);
791    IRBuilder<> FallbackIRB(FallbackBB);
792    CallInst *FallbackCall = FallbackIRB.CreateCall2(
793        DFS.DFSanUnionLoadFn, ShadowAddr, ConstantInt::get(DFS.IntptrTy, Size));
794    FallbackCall->addAttribute(AttributeSet::ReturnIndex, Attribute::ZExt);
795
796    // Compare each of the shadows stored in the loaded 64 bits to each other,
797    // by computing (WideShadow rotl ShadowWidth) == WideShadow.
798    IRBuilder<> IRB(Pos);
799    Value *WideAddr =
800        IRB.CreateBitCast(ShadowAddr, Type::getInt64PtrTy(*DFS.Ctx));
801    Value *WideShadow = IRB.CreateAlignedLoad(WideAddr, ShadowAlign);
802    Value *TruncShadow = IRB.CreateTrunc(WideShadow, DFS.ShadowTy);
803    Value *ShlShadow = IRB.CreateShl(WideShadow, DFS.ShadowWidth);
804    Value *ShrShadow = IRB.CreateLShr(WideShadow, 64 - DFS.ShadowWidth);
805    Value *RotShadow = IRB.CreateOr(ShlShadow, ShrShadow);
806    Value *ShadowsEq = IRB.CreateICmpEQ(WideShadow, RotShadow);
807
808    BasicBlock *Head = Pos->getParent();
809    BasicBlock *Tail = Head->splitBasicBlock(Pos);
810    // In the following code LastBr will refer to the previous basic block's
811    // conditional branch instruction, whose true successor is fixed up to point
812    // to the next block during the loop below or to the tail after the final
813    // iteration.
814    BranchInst *LastBr = BranchInst::Create(FallbackBB, FallbackBB, ShadowsEq);
815    ReplaceInstWithInst(Head->getTerminator(), LastBr);
816
817    for (uint64_t Ofs = 64 / DFS.ShadowWidth; Ofs != Size;
818         Ofs += 64 / DFS.ShadowWidth) {
819      BasicBlock *NextBB = BasicBlock::Create(*DFS.Ctx, "", F);
820      IRBuilder<> NextIRB(NextBB);
821      WideAddr = NextIRB.CreateGEP(WideAddr, ConstantInt::get(DFS.IntptrTy, 1));
822      Value *NextWideShadow = NextIRB.CreateAlignedLoad(WideAddr, ShadowAlign);
823      ShadowsEq = NextIRB.CreateICmpEQ(WideShadow, NextWideShadow);
824      LastBr->setSuccessor(0, NextBB);
825      LastBr = NextIRB.CreateCondBr(ShadowsEq, FallbackBB, FallbackBB);
826    }
827
828    LastBr->setSuccessor(0, Tail);
829    FallbackIRB.CreateBr(Tail);
830    PHINode *Shadow = PHINode::Create(DFS.ShadowTy, 2, "", &Tail->front());
831    Shadow->addIncoming(FallbackCall, FallbackBB);
832    Shadow->addIncoming(TruncShadow, LastBr->getParent());
833    return Shadow;
834  }
835
836  IRBuilder<> IRB(Pos);
837  CallInst *FallbackCall = IRB.CreateCall2(
838      DFS.DFSanUnionLoadFn, ShadowAddr, ConstantInt::get(DFS.IntptrTy, Size));
839  FallbackCall->addAttribute(AttributeSet::ReturnIndex, Attribute::ZExt);
840  return FallbackCall;
841}
842
843void DFSanVisitor::visitLoadInst(LoadInst &LI) {
844  uint64_t Size = DFSF.DFS.DL->getTypeStoreSize(LI.getType());
845  uint64_t Align;
846  if (ClPreserveAlignment) {
847    Align = LI.getAlignment();
848    if (Align == 0)
849      Align = DFSF.DFS.DL->getABITypeAlignment(LI.getType());
850  } else {
851    Align = 1;
852  }
853  IRBuilder<> IRB(&LI);
854  Value *LoadedShadow =
855      DFSF.loadShadow(LI.getPointerOperand(), Size, Align, &LI);
856  Value *PtrShadow = DFSF.getShadow(LI.getPointerOperand());
857  Value *CombinedShadow = DFSF.DFS.combineShadows(LoadedShadow, PtrShadow, &LI);
858  if (CombinedShadow != DFSF.DFS.ZeroShadow)
859    DFSF.NonZeroChecks.insert(CombinedShadow);
860
861  DFSF.setShadow(&LI, CombinedShadow);
862}
863
864void DFSanFunction::storeShadow(Value *Addr, uint64_t Size, uint64_t Align,
865                                Value *Shadow, Instruction *Pos) {
866  if (AllocaInst *AI = dyn_cast<AllocaInst>(Addr)) {
867    llvm::DenseMap<AllocaInst *, AllocaInst *>::iterator i =
868        AllocaShadowMap.find(AI);
869    if (i != AllocaShadowMap.end()) {
870      IRBuilder<> IRB(Pos);
871      IRB.CreateStore(Shadow, i->second);
872      return;
873    }
874  }
875
876  uint64_t ShadowAlign = Align * DFS.ShadowWidth / 8;
877  IRBuilder<> IRB(Pos);
878  Value *ShadowAddr = DFS.getShadowAddress(Addr, Pos);
879  if (Shadow == DFS.ZeroShadow) {
880    IntegerType *ShadowTy = IntegerType::get(*DFS.Ctx, Size * DFS.ShadowWidth);
881    Value *ExtZeroShadow = ConstantInt::get(ShadowTy, 0);
882    Value *ExtShadowAddr =
883        IRB.CreateBitCast(ShadowAddr, PointerType::getUnqual(ShadowTy));
884    IRB.CreateAlignedStore(ExtZeroShadow, ExtShadowAddr, ShadowAlign);
885    return;
886  }
887
888  const unsigned ShadowVecSize = 128 / DFS.ShadowWidth;
889  uint64_t Offset = 0;
890  if (Size >= ShadowVecSize) {
891    VectorType *ShadowVecTy = VectorType::get(DFS.ShadowTy, ShadowVecSize);
892    Value *ShadowVec = UndefValue::get(ShadowVecTy);
893    for (unsigned i = 0; i != ShadowVecSize; ++i) {
894      ShadowVec = IRB.CreateInsertElement(
895          ShadowVec, Shadow, ConstantInt::get(Type::getInt32Ty(*DFS.Ctx), i));
896    }
897    Value *ShadowVecAddr =
898        IRB.CreateBitCast(ShadowAddr, PointerType::getUnqual(ShadowVecTy));
899    do {
900      Value *CurShadowVecAddr = IRB.CreateConstGEP1_32(ShadowVecAddr, Offset);
901      IRB.CreateAlignedStore(ShadowVec, CurShadowVecAddr, ShadowAlign);
902      Size -= ShadowVecSize;
903      ++Offset;
904    } while (Size >= ShadowVecSize);
905    Offset *= ShadowVecSize;
906  }
907  while (Size > 0) {
908    Value *CurShadowAddr = IRB.CreateConstGEP1_32(ShadowAddr, Offset);
909    IRB.CreateAlignedStore(Shadow, CurShadowAddr, ShadowAlign);
910    --Size;
911    ++Offset;
912  }
913}
914
915void DFSanVisitor::visitStoreInst(StoreInst &SI) {
916  uint64_t Size =
917      DFSF.DFS.DL->getTypeStoreSize(SI.getValueOperand()->getType());
918  uint64_t Align;
919  if (ClPreserveAlignment) {
920    Align = SI.getAlignment();
921    if (Align == 0)
922      Align = DFSF.DFS.DL->getABITypeAlignment(SI.getValueOperand()->getType());
923  } else {
924    Align = 1;
925  }
926  DFSF.storeShadow(SI.getPointerOperand(), Size, Align,
927                   DFSF.getShadow(SI.getValueOperand()), &SI);
928}
929
930void DFSanVisitor::visitBinaryOperator(BinaryOperator &BO) {
931  visitOperandShadowInst(BO);
932}
933
934void DFSanVisitor::visitCastInst(CastInst &CI) { visitOperandShadowInst(CI); }
935
936void DFSanVisitor::visitCmpInst(CmpInst &CI) { visitOperandShadowInst(CI); }
937
938void DFSanVisitor::visitGetElementPtrInst(GetElementPtrInst &GEPI) {
939  visitOperandShadowInst(GEPI);
940}
941
942void DFSanVisitor::visitExtractElementInst(ExtractElementInst &I) {
943  visitOperandShadowInst(I);
944}
945
946void DFSanVisitor::visitInsertElementInst(InsertElementInst &I) {
947  visitOperandShadowInst(I);
948}
949
950void DFSanVisitor::visitShuffleVectorInst(ShuffleVectorInst &I) {
951  visitOperandShadowInst(I);
952}
953
954void DFSanVisitor::visitExtractValueInst(ExtractValueInst &I) {
955  visitOperandShadowInst(I);
956}
957
958void DFSanVisitor::visitInsertValueInst(InsertValueInst &I) {
959  visitOperandShadowInst(I);
960}
961
962void DFSanVisitor::visitAllocaInst(AllocaInst &I) {
963  bool AllLoadsStores = true;
964  for (Instruction::use_iterator i = I.use_begin(), e = I.use_end(); i != e;
965       ++i) {
966    if (isa<LoadInst>(*i))
967      continue;
968
969    if (StoreInst *SI = dyn_cast<StoreInst>(*i)) {
970      if (SI->getPointerOperand() == &I)
971        continue;
972    }
973
974    AllLoadsStores = false;
975    break;
976  }
977  if (AllLoadsStores) {
978    IRBuilder<> IRB(&I);
979    DFSF.AllocaShadowMap[&I] = IRB.CreateAlloca(DFSF.DFS.ShadowTy);
980  }
981  DFSF.setShadow(&I, DFSF.DFS.ZeroShadow);
982}
983
984void DFSanVisitor::visitSelectInst(SelectInst &I) {
985  Value *CondShadow = DFSF.getShadow(I.getCondition());
986  Value *TrueShadow = DFSF.getShadow(I.getTrueValue());
987  Value *FalseShadow = DFSF.getShadow(I.getFalseValue());
988
989  if (isa<VectorType>(I.getCondition()->getType())) {
990    DFSF.setShadow(
991        &I, DFSF.DFS.combineShadows(
992                CondShadow,
993                DFSF.DFS.combineShadows(TrueShadow, FalseShadow, &I), &I));
994  } else {
995    Value *ShadowSel;
996    if (TrueShadow == FalseShadow) {
997      ShadowSel = TrueShadow;
998    } else {
999      ShadowSel =
1000          SelectInst::Create(I.getCondition(), TrueShadow, FalseShadow, "", &I);
1001    }
1002    DFSF.setShadow(&I, DFSF.DFS.combineShadows(CondShadow, ShadowSel, &I));
1003  }
1004}
1005
1006void DFSanVisitor::visitMemSetInst(MemSetInst &I) {
1007  IRBuilder<> IRB(&I);
1008  Value *ValShadow = DFSF.getShadow(I.getValue());
1009  IRB.CreateCall3(
1010      DFSF.DFS.DFSanSetLabelFn, ValShadow,
1011      IRB.CreateBitCast(I.getDest(), Type::getInt8PtrTy(*DFSF.DFS.Ctx)),
1012      IRB.CreateZExtOrTrunc(I.getLength(), DFSF.DFS.IntptrTy));
1013}
1014
1015void DFSanVisitor::visitMemTransferInst(MemTransferInst &I) {
1016  IRBuilder<> IRB(&I);
1017  Value *DestShadow = DFSF.DFS.getShadowAddress(I.getDest(), &I);
1018  Value *SrcShadow = DFSF.DFS.getShadowAddress(I.getSource(), &I);
1019  Value *LenShadow = IRB.CreateMul(
1020      I.getLength(),
1021      ConstantInt::get(I.getLength()->getType(), DFSF.DFS.ShadowWidth / 8));
1022  Value *AlignShadow;
1023  if (ClPreserveAlignment) {
1024    AlignShadow = IRB.CreateMul(I.getAlignmentCst(),
1025                                ConstantInt::get(I.getAlignmentCst()->getType(),
1026                                                 DFSF.DFS.ShadowWidth / 8));
1027  } else {
1028    AlignShadow = ConstantInt::get(I.getAlignmentCst()->getType(),
1029                                   DFSF.DFS.ShadowWidth / 8);
1030  }
1031  Type *Int8Ptr = Type::getInt8PtrTy(*DFSF.DFS.Ctx);
1032  DestShadow = IRB.CreateBitCast(DestShadow, Int8Ptr);
1033  SrcShadow = IRB.CreateBitCast(SrcShadow, Int8Ptr);
1034  IRB.CreateCall5(I.getCalledValue(), DestShadow, SrcShadow, LenShadow,
1035                  AlignShadow, I.getVolatileCst());
1036}
1037
1038void DFSanVisitor::visitReturnInst(ReturnInst &RI) {
1039  if (!DFSF.IsNativeABI && RI.getReturnValue()) {
1040    switch (DFSF.IA) {
1041    case DataFlowSanitizer::IA_TLS: {
1042      Value *S = DFSF.getShadow(RI.getReturnValue());
1043      IRBuilder<> IRB(&RI);
1044      IRB.CreateStore(S, DFSF.getRetvalTLS());
1045      break;
1046    }
1047    case DataFlowSanitizer::IA_Args: {
1048      IRBuilder<> IRB(&RI);
1049      Type *RT = DFSF.F->getFunctionType()->getReturnType();
1050      Value *InsVal =
1051          IRB.CreateInsertValue(UndefValue::get(RT), RI.getReturnValue(), 0);
1052      Value *InsShadow =
1053          IRB.CreateInsertValue(InsVal, DFSF.getShadow(RI.getReturnValue()), 1);
1054      RI.setOperand(0, InsShadow);
1055      break;
1056    }
1057    }
1058  }
1059}
1060
1061void DFSanVisitor::visitCallSite(CallSite CS) {
1062  Function *F = CS.getCalledFunction();
1063  if ((F && F->isIntrinsic()) || isa<InlineAsm>(CS.getCalledValue())) {
1064    visitOperandShadowInst(*CS.getInstruction());
1065    return;
1066  }
1067
1068  IRBuilder<> IRB(CS.getInstruction());
1069
1070  DenseMap<Value *, Function *>::iterator i =
1071      DFSF.DFS.UnwrappedFnMap.find(CS.getCalledValue());
1072  if (i != DFSF.DFS.UnwrappedFnMap.end()) {
1073    Function *F = i->second;
1074    switch (DFSF.DFS.getWrapperKind(F)) {
1075    case DataFlowSanitizer::WK_Warning: {
1076      CS.setCalledFunction(F);
1077      IRB.CreateCall(DFSF.DFS.DFSanUnimplementedFn,
1078                     IRB.CreateGlobalStringPtr(F->getName()));
1079      DFSF.setShadow(CS.getInstruction(), DFSF.DFS.ZeroShadow);
1080      return;
1081    }
1082    case DataFlowSanitizer::WK_Discard: {
1083      CS.setCalledFunction(F);
1084      DFSF.setShadow(CS.getInstruction(), DFSF.DFS.ZeroShadow);
1085      return;
1086    }
1087    case DataFlowSanitizer::WK_Functional: {
1088      CS.setCalledFunction(F);
1089      visitOperandShadowInst(*CS.getInstruction());
1090      return;
1091    }
1092    case DataFlowSanitizer::WK_Custom: {
1093      // Don't try to handle invokes of custom functions, it's too complicated.
1094      // Instead, invoke the dfsw$ wrapper, which will in turn call the __dfsw_
1095      // wrapper.
1096      if (CallInst *CI = dyn_cast<CallInst>(CS.getInstruction())) {
1097        FunctionType *FT = F->getFunctionType();
1098        FunctionType *CustomFT = DFSF.DFS.getCustomFunctionType(FT);
1099        std::string CustomFName = "__dfsw_";
1100        CustomFName += F->getName();
1101        Constant *CustomF =
1102            DFSF.DFS.Mod->getOrInsertFunction(CustomFName, CustomFT);
1103        if (Function *CustomFn = dyn_cast<Function>(CustomF)) {
1104          CustomFn->copyAttributesFrom(F);
1105
1106          // Custom functions returning non-void will write to the return label.
1107          if (!FT->getReturnType()->isVoidTy()) {
1108            CustomFn->removeAttributes(AttributeSet::FunctionIndex,
1109                                       DFSF.DFS.ReadOnlyNoneAttrs);
1110          }
1111        }
1112
1113        std::vector<Value *> Args;
1114
1115        CallSite::arg_iterator i = CS.arg_begin();
1116        for (unsigned n = FT->getNumParams(); n != 0; ++i, --n)
1117          Args.push_back(*i);
1118
1119        i = CS.arg_begin();
1120        for (unsigned n = FT->getNumParams(); n != 0; ++i, --n)
1121          Args.push_back(DFSF.getShadow(*i));
1122
1123        if (!FT->getReturnType()->isVoidTy()) {
1124          if (!DFSF.LabelReturnAlloca) {
1125            DFSF.LabelReturnAlloca =
1126                new AllocaInst(DFSF.DFS.ShadowTy, "labelreturn",
1127                               DFSF.F->getEntryBlock().begin());
1128          }
1129          Args.push_back(DFSF.LabelReturnAlloca);
1130        }
1131
1132        CallInst *CustomCI = IRB.CreateCall(CustomF, Args);
1133        CustomCI->setCallingConv(CI->getCallingConv());
1134        CustomCI->setAttributes(CI->getAttributes());
1135
1136        if (!FT->getReturnType()->isVoidTy()) {
1137          LoadInst *LabelLoad = IRB.CreateLoad(DFSF.LabelReturnAlloca);
1138          DFSF.setShadow(CustomCI, LabelLoad);
1139        }
1140
1141        CI->replaceAllUsesWith(CustomCI);
1142        CI->eraseFromParent();
1143        return;
1144      }
1145      break;
1146    }
1147    }
1148  }
1149
1150  FunctionType *FT = cast<FunctionType>(
1151      CS.getCalledValue()->getType()->getPointerElementType());
1152  if (DFSF.DFS.getInstrumentedABI() == DataFlowSanitizer::IA_TLS) {
1153    for (unsigned i = 0, n = FT->getNumParams(); i != n; ++i) {
1154      IRB.CreateStore(DFSF.getShadow(CS.getArgument(i)),
1155                      DFSF.getArgTLS(i, CS.getInstruction()));
1156    }
1157  }
1158
1159  Instruction *Next = 0;
1160  if (!CS.getType()->isVoidTy()) {
1161    if (InvokeInst *II = dyn_cast<InvokeInst>(CS.getInstruction())) {
1162      if (II->getNormalDest()->getSinglePredecessor()) {
1163        Next = II->getNormalDest()->begin();
1164      } else {
1165        BasicBlock *NewBB =
1166            SplitEdge(II->getParent(), II->getNormalDest(), &DFSF.DFS);
1167        Next = NewBB->begin();
1168      }
1169    } else {
1170      Next = CS->getNextNode();
1171    }
1172
1173    if (DFSF.DFS.getInstrumentedABI() == DataFlowSanitizer::IA_TLS) {
1174      IRBuilder<> NextIRB(Next);
1175      LoadInst *LI = NextIRB.CreateLoad(DFSF.getRetvalTLS());
1176      DFSF.SkipInsts.insert(LI);
1177      DFSF.setShadow(CS.getInstruction(), LI);
1178      DFSF.NonZeroChecks.insert(LI);
1179    }
1180  }
1181
1182  // Do all instrumentation for IA_Args down here to defer tampering with the
1183  // CFG in a way that SplitEdge may be able to detect.
1184  if (DFSF.DFS.getInstrumentedABI() == DataFlowSanitizer::IA_Args) {
1185    FunctionType *NewFT = DFSF.DFS.getArgsFunctionType(FT);
1186    Value *Func =
1187        IRB.CreateBitCast(CS.getCalledValue(), PointerType::getUnqual(NewFT));
1188    std::vector<Value *> Args;
1189
1190    CallSite::arg_iterator i = CS.arg_begin(), e = CS.arg_end();
1191    for (unsigned n = FT->getNumParams(); n != 0; ++i, --n)
1192      Args.push_back(*i);
1193
1194    i = CS.arg_begin();
1195    for (unsigned n = FT->getNumParams(); n != 0; ++i, --n)
1196      Args.push_back(DFSF.getShadow(*i));
1197
1198    if (FT->isVarArg()) {
1199      unsigned VarArgSize = CS.arg_size() - FT->getNumParams();
1200      ArrayType *VarArgArrayTy = ArrayType::get(DFSF.DFS.ShadowTy, VarArgSize);
1201      AllocaInst *VarArgShadow =
1202          new AllocaInst(VarArgArrayTy, "", DFSF.F->getEntryBlock().begin());
1203      Args.push_back(IRB.CreateConstGEP2_32(VarArgShadow, 0, 0));
1204      for (unsigned n = 0; i != e; ++i, ++n) {
1205        IRB.CreateStore(DFSF.getShadow(*i),
1206                        IRB.CreateConstGEP2_32(VarArgShadow, 0, n));
1207        Args.push_back(*i);
1208      }
1209    }
1210
1211    CallSite NewCS;
1212    if (InvokeInst *II = dyn_cast<InvokeInst>(CS.getInstruction())) {
1213      NewCS = IRB.CreateInvoke(Func, II->getNormalDest(), II->getUnwindDest(),
1214                               Args);
1215    } else {
1216      NewCS = IRB.CreateCall(Func, Args);
1217    }
1218    NewCS.setCallingConv(CS.getCallingConv());
1219    NewCS.setAttributes(CS.getAttributes().removeAttributes(
1220        *DFSF.DFS.Ctx, AttributeSet::ReturnIndex,
1221        AttributeFuncs::typeIncompatible(NewCS.getInstruction()->getType(),
1222                                         AttributeSet::ReturnIndex)));
1223
1224    if (Next) {
1225      ExtractValueInst *ExVal =
1226          ExtractValueInst::Create(NewCS.getInstruction(), 0, "", Next);
1227      DFSF.SkipInsts.insert(ExVal);
1228      ExtractValueInst *ExShadow =
1229          ExtractValueInst::Create(NewCS.getInstruction(), 1, "", Next);
1230      DFSF.SkipInsts.insert(ExShadow);
1231      DFSF.setShadow(ExVal, ExShadow);
1232      DFSF.NonZeroChecks.insert(ExShadow);
1233
1234      CS.getInstruction()->replaceAllUsesWith(ExVal);
1235    }
1236
1237    CS.getInstruction()->eraseFromParent();
1238  }
1239}
1240
1241void DFSanVisitor::visitPHINode(PHINode &PN) {
1242  PHINode *ShadowPN =
1243      PHINode::Create(DFSF.DFS.ShadowTy, PN.getNumIncomingValues(), "", &PN);
1244
1245  // Give the shadow phi node valid predecessors to fool SplitEdge into working.
1246  Value *UndefShadow = UndefValue::get(DFSF.DFS.ShadowTy);
1247  for (PHINode::block_iterator i = PN.block_begin(), e = PN.block_end(); i != e;
1248       ++i) {
1249    ShadowPN->addIncoming(UndefShadow, *i);
1250  }
1251
1252  DFSF.PHIFixups.push_back(std::make_pair(&PN, ShadowPN));
1253  DFSF.setShadow(&PN, ShadowPN);
1254}
1255