Instrumented local variables and implemented runtime.

src/IceASanInstrumentation.cpp

 //===- subzero/src/IceASanInstrumentation.cpp - ASan ------------*- C++ -*-===//
 //
 //                        The Subzero Code Generator
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 ///
 /// \file
 /// \brief Implements the AddressSanitizer instrumentation class.
 ///
 //===----------------------------------------------------------------------===//
 
 #include "IceASanInstrumentation.h"
 
 #include "IceBuildDefs.h"
 #include "IceCfg.h"
 #include "IceCfgNode.h"
 #include "IceGlobalInits.h"
 #include "IceInst.h"
 #include "IceTargetLowering.h"
 #include "IceTypes.h"
 
 #include <sstream>
 #include <unordered_map>
+#include <vector>
 
 namespace Ice {
 
 namespace {
+
 constexpr SizeT RzSize = 32;
 const std::string RzPrefix = "__$rz";
 const llvm::NaClBitcodeRecord::RecordVector RzContents =
     llvm::NaClBitcodeRecord::RecordVector(RzSize, 'R');
 
 // TODO(tlively): Handle all allocation functions
 // In order to instrument the code correctly, the .pexe must not have had its
 // symbols stripped.
 using string_map = std::unordered_map<std::string, std::string>;
 const string_map FuncSubstitutions = {{"malloc", "__asan_malloc"},
                                       {"free", "__asan_free"}};
 
 } // end of anonymous namespace
 
+ICE_TLS_DEFINE_FIELD(std::vector<InstCall *> *, ASanInstrumentation,
+                     LocalDtors);
+
 // Create redzones around all global variables, ensuring that the initializer
 // types of the redzones and their associated globals match so that they are
 // laid out together in memory.
 void ASanInstrumentation::instrumentGlobals(VariableDeclarationList &Globals) {
   if (DidInsertRedZones)
     return;
 
   VariableDeclarationList NewGlobals;
   // Global holding pointers to all redzones
   auto *RzArray = VariableDeclaration::create(&NewGlobals);
@@ skipping 64 matching lines @@
   Rz->setIsConstant(Global->getIsConstant());
   RzArray->addInitializer(VariableDeclaration::RelocInitializer::create(
       List, Rz, RelocOffsetArray(0)));
   ++RzArraySize;
   return Rz;
 }
 
 // Check for an alloca signaling the presence of local variables and add a
 // redzone if it is found
 void ASanInstrumentation::instrumentFuncStart(LoweringContext &Context) {
-  auto *FirstAlloca = llvm::dyn_cast<InstAlloca>(Context.getCur());
-  if (FirstAlloca == nullptr)
-    return;
+  if (ICE_TLS_GET_FIELD(LocalDtors) == nullptr)
+    ICE_TLS_SET_FIELD(LocalDtors, new std::vector<InstCall *>());
 
-  constexpr SizeT Alignment = 4;
-  InstAlloca *RzAlloca = createLocalRz(Context, RzSize, Alignment);
+  Cfg *Func = Context.getNode()->getCfg();
+  bool HasLocals = false;
+  LoweringContext C;
+  C.init(Context.getNode());
+  std::vector<Inst *> Initializations;
+  Constant *InitFunc =
+      Ctx->getConstantExternSym(Ctx->getGlobalString("__asan_poison"));
+  Constant *DestroyFunc =
+      Ctx->getConstantExternSym(Ctx->getGlobalString("__asan_unpoison"));
 
-  // insert before the current instruction
-  InstList::iterator Next = Context.getNext();
-  Context.setInsertPoint(Context.getCur());
-  Context.insert(RzAlloca);
-  Context.setNext(Next);
-}
+  InstAlloca *Cur;
+  ConstantInteger32 *VarSizeOp;
+  while (
+      (Cur = llvm::dyn_cast<InstAlloca>(iteratorToInst(C.getCur()))) &&
+      (VarSizeOp = llvm::dyn_cast<ConstantInteger32>(Cur->getSizeInBytes()))) {
+    HasLocals = true;
 
-void ASanInstrumentation::instrumentAlloca(LoweringContext &Context,
-                                           InstAlloca *Instr) {
-  auto *VarSizeOp = llvm::dyn_cast<ConstantInteger32>(Instr->getSizeInBytes());
-  SizeT VarSize = (VarSizeOp == nullptr) ? RzSize : VarSizeOp->getValue();
-  SizeT Padding = Utils::OffsetToAlignment(VarSize, RzSize);
-  constexpr SizeT Alignment = 1;
-  InstAlloca *Rz = createLocalRz(Context, RzSize + Padding, Alignment);
-  Context.insert(Rz);
-}
+    // create the new alloca that includes a redzone
+    SizeT VarSize = VarSizeOp->getValue();
+    Variable *Dest = Cur->getDest();
+    SizeT RzPadding = RzSize + Utils::OffsetToAlignment(VarSize, RzSize);
+    auto *ByteCount =
+        ConstantInteger32::create(Ctx, IceType_i32, VarSize + RzPadding);
+    constexpr SizeT Alignment = 8;
+    auto *NewVar = InstAlloca::create(Func, Dest, ByteCount, Alignment);
 
-InstAlloca *ASanInstrumentation::createLocalRz(LoweringContext &Context,
-                                               SizeT Size, SizeT Alignment) {
-  Cfg *Func = Context.getNode()->getCfg();
-  Variable *Rz = Func->makeVariable(IceType_i32);
-  Rz->setName(Func, nextRzName());
-  auto *ByteCount = ConstantInteger32::create(Ctx, IceType_i32, Size);
-  auto *RzAlloca = InstAlloca::create(Func, Rz, ByteCount, Alignment);
-  return RzAlloca;
+    // calculate the redzone offset
+    Variable *RzLocVar = Func->makeVariable(IceType_i32);
+    RzLocVar->setName(Func, nextRzName());
+    auto *Offset = ConstantInteger32::create(Ctx, IceType_i32, VarSize);
+    auto *RzLoc = InstArithmetic::create(Func, InstArithmetic::Add, RzLocVar,
+                                         Dest, Offset);
+
+    // instructions to poison and unpoison the redzone
+    constexpr SizeT NumArgs = 2;
+    constexpr Variable *Void = nullptr;
+    constexpr bool NoTailcall = false;
+    auto *Init = InstCall::create(Func, NumArgs, Void, InitFunc, NoTailcall);
+    auto *Destroy =
+        InstCall::create(Func, NumArgs, Void, DestroyFunc, NoTailcall);
+    Init->addArg(RzLocVar);
+    Destroy->addArg(RzLocVar);
+    auto *RzSizeConst = ConstantInteger32::create(Ctx, IceType_i32, RzPadding);
+    Init->addArg(RzSizeConst);
+    Destroy->addArg(RzSizeConst);
+
+    Cur->setDeleted();
+    C.insert(NewVar);
+    ICE_TLS_GET_FIELD(LocalDtors)->emplace_back(Destroy);
+    Initializations.emplace_back(RzLoc);
+    Initializations.emplace_back(Init);
+
+    C.advanceCur();
+    C.advanceNext();
+  }
+
+  C.setInsertPoint(C.getCur());
+
+  // add the leftmost redzone
+  if (HasLocals) {
+    Variable *LastRz = Func->makeVariable(IceType_i32);
+    LastRz->setName(Func, nextRzName());
+    auto *ByteCount = ConstantInteger32::create(Ctx, IceType_i32, RzSize);
+    constexpr SizeT Alignment = 8;
+    auto *RzAlloca = InstAlloca::create(Func, LastRz, ByteCount, Alignment);
+
+    constexpr SizeT NumArgs = 2;
+    constexpr Variable *Void = nullptr;
+    constexpr bool NoTailcall = false;
+    auto *Init = InstCall::create(Func, NumArgs, Void, InitFunc, NoTailcall);
+    auto *Destroy =
+        InstCall::create(Func, NumArgs, Void, DestroyFunc, NoTailcall);
+    Init->addArg(LastRz);
+    Destroy->addArg(LastRz);
+    Init->addArg(RzAlloca->getSizeInBytes());
+    Destroy->addArg(RzAlloca->getSizeInBytes());
+
+    ICE_TLS_GET_FIELD(LocalDtors)->emplace_back(Destroy);
+    C.insert(RzAlloca);
+    C.insert(Init);
+  }
+
+  // insert initializers for the redzones
+  for (Inst *Init : Initializations) {
+    C.insert(Init);
+  }
 }
 
 void ASanInstrumentation::instrumentCall(LoweringContext &Context,
                                          InstCall *Instr) {
   auto *CallTarget =
       llvm::dyn_cast<ConstantRelocatable>(Instr->getCallTarget());
   if (CallTarget == nullptr)
     return;
 
   std::string TargetName = CallTarget->getName().toStringOrEmpty();
@@ skipping 36 matching lines @@
                                 AccessCheck, NoTailCall);
   Call->addArg(Op);
   Call->addArg(ConstantInteger32::create(Ctx, IceType_i32, Size));
   // play games to insert the call before the access instruction
   InstList::iterator Next = Context.getNext();
   Context.setInsertPoint(Context.getCur());
   Context.insert(Call);
   Context.setNext(Next);
 }
 
+void ASanInstrumentation::instrumentRet(LoweringContext &Context, InstRet *) {
+  InstList::iterator Next = Context.getNext();
+  Context.setInsertPoint(Context.getCur());
+  for (InstCall *RzUnpoison : *ICE_TLS_GET_FIELD(LocalDtors)) {
+    Context.insert(RzUnpoison);
+  }
+  Context.setNext(Next);
+}
+
 void ASanInstrumentation::instrumentStart(Cfg *Func) {
   Constant *ShadowMemInit =
       Ctx->getConstantExternSym(Ctx->getGlobalString("__asan_init"));
   constexpr SizeT NumArgs = 0;
   constexpr Variable *Void = nullptr;
   constexpr bool NoTailCall = false;
   auto *Call = InstCall::create(Func, NumArgs, Void, ShadowMemInit, NoTailCall);
   Func->getEntryNode()->getInsts().push_front(Call);
 }
 
+// TODO(tlively): make this more efficient with swap idiom
+void ASanInstrumentation::finishFunc(Cfg *Func) {
+  (void)Func;
+  ICE_TLS_GET_FIELD(LocalDtors)->clear();
+}
+
 } // end of namespace Ice