Subzero: Implemented codegen for poisoning and unpoisoning stack redzones

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
@@ skipping 13 matching lines @@
 
 #include <sstream>
 #include <unordered_map>
 #include <unordered_set>
 #include <vector>
 
 namespace Ice {
 
 namespace {
 
+constexpr SizeT BytesPerWord = sizeof(uint32_t);
+constexpr SizeT RzSize = 32;
+constexpr SizeT ShadowScaleLog2 = 3;
+constexpr SizeT ShadowScale = 1 << ShadowScaleLog2;
+constexpr SizeT ShadowLength32 = 1 << (32 - ShadowScaleLog2);
+constexpr int32_t StackPoisonVal = -1;
 constexpr const char *ASanPrefix = "__asan";
-constexpr SizeT RzSize = 32;
 constexpr const char *RzPrefix = "__$rz";
 constexpr const char *RzArrayName = "__$rz_array";
 constexpr const char *RzSizesName = "__$rz_sizes";
-constexpr char RzStackPoison = -1;
 const llvm::NaClBitcodeRecord::RecordVector RzContents =
     llvm::NaClBitcodeRecord::RecordVector(RzSize, 'R');
 
 // In order to instrument the code correctly, the .pexe must not have had its
 // symbols stripped.
 using StringMap = std::unordered_map<std::string, std::string>;
 using StringSet = std::unordered_set<std::string>;
 // TODO(tlively): Handle all allocation functions
 const StringMap FuncSubstitutions = {{"malloc", "__asan_malloc"},
                                      {"free", "__asan_free"},
                                      {"calloc", "__asan_calloc"},
                                      {"__asan_dummy_calloc", "__asan_calloc"},
                                      {"realloc", "__asan_realloc"}};
 const StringSet FuncBlackList = {"_Balloc"};
 
 llvm::NaClBitcodeRecord::RecordVector sizeToByteVec(SizeT Size) {
   llvm::NaClBitcodeRecord::RecordVector SizeContents;
   for (unsigned i = 0; i < sizeof(Size); ++i) {
     SizeContents.emplace_back(Size % (1 << CHAR_BIT));
     Size >>= CHAR_BIT;
   }
   return SizeContents;
 }
 
 } // end of anonymous namespace
 
 ICE_TLS_DEFINE_FIELD(VarSizeMap *, ASanInstrumentation, LocalVars);
-ICE_TLS_DEFINE_FIELD(std::vector<InstCall *> *, ASanInstrumentation,
+ICE_TLS_DEFINE_FIELD(std::vector<InstStore *> *, ASanInstrumentation,
                      LocalDtors);
 
 bool ASanInstrumentation::isInstrumentable(Cfg *Func) {
   std::string FuncName = Func->getFunctionName().toStringOrEmpty();
   return FuncName == "" ||
          (FuncBlackList.count(FuncName) == 0 && FuncName.find(ASanPrefix) != 0);
 }
 
 // Create redzones around all global variables, ensuring that the initializer
 // types of the redzones and their associated globals match so that they are
@@ skipping 77 matching lines @@
 std::string ASanInstrumentation::nextRzName() {
   std::stringstream Name;
   Name << RzPrefix << RzNum++;
   return Name.str();
 }
 
 // Check for an alloca signaling the presence of local variables and add a
 // redzone if it is found
 void ASanInstrumentation::instrumentFuncStart(LoweringContext &Context) {
   if (ICE_TLS_GET_FIELD(LocalDtors) == nullptr) {
-    ICE_TLS_SET_FIELD(LocalDtors, new std::vector<InstCall *>());
+    ICE_TLS_SET_FIELD(LocalDtors, new std::vector<InstStore *>());
     ICE_TLS_SET_FIELD(LocalVars, new VarSizeMap());
   }
   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"));
-
+  using Entry = std::pair<SizeT, int32_t>;
+  std::vector<InstAlloca *> NewAllocas;
+  std::vector<Entry> PoisonVals;
+  Variable *FirstShadowLocVar;
+  InstArithmetic *ShadowIndexCalc;
+  InstArithmetic *ShadowLocCalc;
   InstAlloca *Cur;
   ConstantInteger32 *VarSizeOp;
-  while (
-      (Cur = llvm::dyn_cast<InstAlloca>(iteratorToInst(C.getCur()))) &&
-      (VarSizeOp = llvm::dyn_cast<ConstantInteger32>(Cur->getSizeInBytes()))) {
-    HasLocals = true;
+  while (!Context.atEnd()) {
+    Cur = llvm::dyn_cast<InstAlloca>(iteratorToInst(Context.getCur()));
+    VarSizeOp = (Cur == nullptr)
+                    ? nullptr
+                    : llvm::dyn_cast<ConstantInteger32>(Cur->getSizeInBytes());
+    if (Cur == nullptr || VarSizeOp == nullptr) {
+      Context.advanceCur();
+      Context.advanceNext();
+      continue;
+    }
+
+    Cur->setDeleted();
+
+    if (PoisonVals.empty()) {
+      // insert leftmost redzone
+      auto *LastRzVar = Func->makeVariable(IceType_i32);
+      LastRzVar->setName(Func, nextRzName());
+      auto *ByteCount = ConstantInteger32::create(Ctx, IceType_i32, RzSize);
+      constexpr SizeT Alignment = 8;
+      NewAllocas.emplace_back(
+          InstAlloca::create(Func, LastRzVar, ByteCount, Alignment));
+      PoisonVals.emplace_back(Entry{RzSize >> ShadowScaleLog2, StackPoisonVal});
+
+      // Calculate starting address for poisoning
+      FirstShadowLocVar = Func->makeVariable(IceType_i32);
+      FirstShadowLocVar->setName(Func, "firstShadowLoc");
+      auto *ShadowIndexVar = Func->makeVariable(IceType_i32);
+      ShadowIndexVar->setName(Func, "shadowIndex");
+
+      auto *ShadowScaleLog2Const =
+          ConstantInteger32::create(Ctx, IceType_i32, ShadowScaleLog2);
+      auto *ShadowMemLocConst =
+          ConstantInteger32::create(Ctx, IceType_i32, ShadowLength32);
+
+      ShadowIndexCalc =
+          InstArithmetic::create(Func, InstArithmetic::Lshr, ShadowIndexVar,
+                                 LastRzVar, ShadowScaleLog2Const);
+      ShadowLocCalc =
+          InstArithmetic::create(Func, InstArithmetic::Add, FirstShadowLocVar,
+                                 ShadowIndexVar, ShadowMemLocConst);
+    }
 
     // create the new alloca that includes a redzone
     SizeT VarSize = VarSizeOp->getValue();
     Variable *Dest = Cur->getDest();
     ICE_TLS_GET_FIELD(LocalVars)->insert({Dest, VarSize});
     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);
+    NewAllocas.emplace_back(
+        InstAlloca::create(Func, Dest, ByteCount, Alignment));
 
-    // 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 *RzSizeConst = ConstantInteger32::create(Ctx, IceType_i32, RzPadding);
-    auto *RzPoisonConst =
-        ConstantInteger32::create(Ctx, IceType_i32, RzStackPoison);
-    auto *Init = InstCall::create(Func, NumArgs, Void, InitFunc, NoTailcall);
-    Init->addArg(RzLocVar);
-    Init->addArg(RzSizeConst);
-    Init->addArg(RzPoisonConst);
-    auto *Destroy =
-        InstCall::create(Func, NumArgs, Void, DestroyFunc, NoTailcall);
-    Destroy->addArg(RzLocVar);
-    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();
+    const SizeT Zeros = VarSize >> ShadowScaleLog2;
+    const SizeT Offset = VarSize % ShadowScale;
+    const SizeT PoisonBytes =
+        ((VarSize + RzPadding) >> ShadowScaleLog2) - Zeros - 1;
+    if (Zeros > 0)
+      PoisonVals.emplace_back(Entry{Zeros, 0});
+    PoisonVals.emplace_back(Entry{1, (Offset == 0) ? StackPoisonVal : Offset});
+    PoisonVals.emplace_back(Entry{PoisonBytes, StackPoisonVal});
+    Context.advanceCur();
+    Context.advanceNext();
   }
 
-  C.setInsertPoint(C.getCur());
+  Context.rewind();
+  if (PoisonVals.empty()) {
+    Context.advanceNext();
+    return;
+  }
+  for (InstAlloca *RzAlloca : NewAllocas) {
+    Context.insert(RzAlloca);
+  }
+  Context.insert(ShadowIndexCalc);
+  Context.insert(ShadowLocCalc);
 
-  // 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 *RzPoisonConst =
-        ConstantInteger32::create(Ctx, IceType_i32, RzStackPoison);
-    auto *Init = InstCall::create(Func, NumArgs, Void, InitFunc, NoTailcall);
-    Init->addArg(LastRz);
-    Init->addArg(RzAlloca->getSizeInBytes());
-    Init->addArg(RzPoisonConst);
-    auto *Destroy =
-        InstCall::create(Func, NumArgs, Void, DestroyFunc, NoTailcall);
-    Destroy->addArg(LastRz);
-    Destroy->addArg(RzAlloca->getSizeInBytes());
-    ICE_TLS_GET_FIELD(LocalDtors)->emplace_back(Destroy);
-    C.insert(RzAlloca);
-    C.insert(Init);
+  // Poison redzones
+  std::vector<Entry>::iterator Iter = PoisonVals.begin();
+  for (SizeT Offset = 0; Iter != PoisonVals.end(); Offset += BytesPerWord) {
+    int32_t CurVals[BytesPerWord] = {0};
+    for (uint32_t i = 0; i < BytesPerWord; ++i) {
+      if (Iter == PoisonVals.end())
+        break;
+      Entry Val = *Iter;
+      CurVals[i] = Val.second;
+      --Val.first;
+      if (Val.first > 0)
+        *Iter = Val;
+      else
+        ++Iter;
+    }
+    int32_t Poison = ((CurVals[3] & 0xff) << 24) | ((CurVals[2] & 0xff) << 16) |
+                     ((CurVals[1] & 0xff) << 8) | (CurVals[0] & 0xff);
+    if (Poison == 0)
+      continue;
+    auto *PoisonConst = ConstantInteger32::create(Ctx, IceType_i32, Poison);
+    auto *ZeroConst = ConstantInteger32::create(Ctx, IceType_i32, 0);
+    auto *OffsetConst = ConstantInteger32::create(Ctx, IceType_i32, Offset);
+    auto *PoisonAddrVar = Func->makeVariable(IceType_i32);
+    Context.insert(InstArithmetic::create(Func, InstArithmetic::Add,
+                                          PoisonAddrVar, FirstShadowLocVar,
+                                          OffsetConst));
+    Context.insert(InstStore::create(Func, PoisonConst, PoisonAddrVar));
+    ICE_TLS_GET_FIELD(LocalDtors)
+        ->emplace_back(InstStore::create(Func, ZeroConst, PoisonAddrVar));
   }
-
-  // insert initializers for the redzones
-  for (Inst *Init : Initializations) {
-    C.insert(Init);
-  }
+  Context.advanceNext();
 }
 
 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 56 matching lines @@
   auto *Reloc = llvm::dyn_cast<ConstantRelocatable>(Op);
   if (Reloc == nullptr)
     return false;
   RelocOffsetT Offset = Reloc->getOffset();
   GlobalSizeMap::iterator GlobalSize = GlobalSizes.find(Reloc->getName());
   return GlobalSize != GlobalSizes.end() && GlobalSize->second - Offset >= Size;
 }
 
 void ASanInstrumentation::instrumentRet(LoweringContext &Context, InstRet *) {
   Cfg *Func = Context.getNode()->getCfg();
-  InstList::iterator Next = Context.getNext();
   Context.setInsertPoint(Context.getCur());
-  for (InstCall *RzUnpoison : *ICE_TLS_GET_FIELD(LocalDtors)) {
-    SizeT NumArgs = RzUnpoison->getNumArgs();
-    Variable *Dest = RzUnpoison->getDest();
-    Operand *CallTarget = RzUnpoison->getCallTarget();
-    bool HasTailCall = RzUnpoison->isTailcall();
-    bool IsTargetHelperCall = RzUnpoison->isTargetHelperCall();
-    auto *RzUnpoisonCpy = InstCall::create(Func, NumArgs, Dest, CallTarget,
-                                           HasTailCall, IsTargetHelperCall);
-    for (int I = 0, Args = RzUnpoison->getNumArgs(); I < Args; ++I) {
-      RzUnpoisonCpy->addArg(RzUnpoison->getArg(I));
-    }
-    Context.insert(RzUnpoisonCpy);
+  for (InstStore *RzUnpoison : *ICE_TLS_GET_FIELD(LocalDtors)) {
+    Context.insert(
+        InstStore::create(Func, RzUnpoison->getData(), RzUnpoison->getAddr()));
   }
-  Context.setNext(Next);
+  Context.advanceCur();
+  Context.advanceNext();
 }
 
 void ASanInstrumentation::instrumentStart(Cfg *Func) {
   Constant *ShadowMemInit =
       Ctx->getConstantExternSym(Ctx->getGlobalString("__asan_init"));
   constexpr SizeT NumArgs = 3;
   constexpr Variable *Void = nullptr;
   constexpr bool NoTailCall = false;
   auto *Call = InstCall::create(Func, NumArgs, Void, ShadowMemInit, NoTailCall);
   Func->getEntryNode()->getInsts().push_front(Call);
 
   instrumentGlobals(*getGlobals());
 
   Call->addArg(ConstantInteger32::create(Ctx, IceType_i32, RzGlobalsNum));
   Call->addArg(Ctx->getConstantSym(0, Ctx->getGlobalString(RzArrayName)));
   Call->addArg(Ctx->getConstantSym(0, Ctx->getGlobalString(RzSizesName)));
 }
 
 // TODO(tlively): make this more efficient with swap idiom
 void ASanInstrumentation::finishFunc(Cfg *) {
   ICE_TLS_GET_FIELD(LocalVars)->clear();
   ICE_TLS_GET_FIELD(LocalDtors)->clear();
 }
 
 } // end of namespace Ice