Subzero: Implemented codegen for poisoning and unpoisoning stack redzones

src/IceASanInstrumentation.cpp

Index: src/IceASanInstrumentation.cpp
diff --git a/src/IceASanInstrumentation.cpp b/src/IceASanInstrumentation.cpp
index e3203b36225ed5b4fa6cdc7799f079dd9ee5c2b6..2c84cc05ea327896a34b9bdc09ed116e9fc98eab 100644
--- a/src/IceASanInstrumentation.cpp
+++ b/src/IceASanInstrumentation.cpp
@@ -31,8 +31,12 @@ namespace Ice {
 
 namespace {
 
-constexpr const char *ASanPrefix = "__asan";
 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 const char *RzPrefix = "__$rz";
 constexpr const char *RzArrayName = "__$rz_array";
 constexpr const char *RzSizesName = "__$rz_sizes";
@@ -63,7 +67,7 @@ llvm::NaClBitcodeRecord::RecordVector sizeToByteVec(SizeT Size) {
 } // 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) {
@@ -161,25 +165,49 @@ std::string ASanInstrumentation::nextRzName() {
 // 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<Entry> PoisonVals;
+  Variable *FirstShadowLocVar;
+  InstArithmetic *ShadowIndexCalc;
+  InstArithmetic *ShadowLocCalc;
   InstAlloca *Cur;
   ConstantInteger32 *VarSizeOp;
   while (
-      (Cur = llvm::dyn_cast<InstAlloca>(iteratorToInst(C.getCur()))) &&
+      (Cur = llvm::dyn_cast<InstAlloca>(iteratorToInst(Context.getCur()))) &&
       (VarSizeOp = llvm::dyn_cast<ConstantInteger32>(Cur->getSizeInBytes()))) {

[Karl, 2016/08/01 15:40:21]

Doesn't this code assume that all Allocas are at the beginning? Is this a good
assumption?

Also, if the size of an alloc (in bytes) is zero, won't this stop the loop?

[tlively, 2016/08/01 17:08:50]

I believe that allocas of known size corresponding to stack variables are always
at the beginning. I don't currently do anything with dynamic allocas (though I
should).

The loops would only stop if a zero byte alloca is represented with a nullptr as
its size operand. I'm not sure whether that's the case. Also I'm not sure that a
stack variable would ever be zero bytes. Can you confirm?

[Karl, 2016/08/04 14:26:34]

I don't know the answer to this. I was asking because I didn't know.

[Jim Stichnoth, 2016/08/04 14:42:54]

Good point Karl.  All the fixed-size alloca instructions (at least those we care
about) will appear in the entry block, but they could be interleaved with other
instructions.

See Cfg::processAllocas() which iterates through the entire entry block without
an early exit.

[tlively, 2016/08/05 18:24:52]

Done.

-    HasLocals = true;
+    Cur->setDeleted();
+
+    if (PoisonVals.empty()) {
+      // insert leftmost redzone
+      Variable *LastRzVar = Func->makeVariable(IceType_i32);

[Jim Stichnoth, 2016/08/01 20:38:00]

auto * ?

(you used it below for ShadowIndexVar)

[tlively, 2016/08/05 18:24:52]

Done.

+      LastRzVar->setName(Func, nextRzName());
+      auto *ByteCount = ConstantInteger32::create(Ctx, IceType_i32, RzSize);
+      constexpr SizeT Alignment = 8;
+      Context.insert(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();
@@ -190,67 +218,58 @@ void ASanInstrumentation::instrumentFuncStart(LoweringContext &Context) {
         ConstantInteger32::create(Ctx, IceType_i32, VarSize + RzPadding);
     constexpr SizeT Alignment = 8;
     auto *NewVar = InstAlloca::create(Func, Dest, ByteCount, Alignment);

[Jim Stichnoth, 2016/08/01 20:38:00]

Could you name this something like "NewAllocaVar" or similar?

When I see "auto *NewVar = ...", I assume its type is Variable * .

[tlively, 2016/08/05 18:24:52]

Done.

-
-    // 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();
+    Context.insert(NewVar);
+
+    SizeT Zeros = VarSize >> ShadowScaleLog2;

[Jim Stichnoth, 2016/08/01 20:38:01]

Declare scalar locals as const when possible/practical/reasonable.

[tlively, 2016/08/05 18:24:52]

Done.

+    SizeT Offset = VarSize % ShadowScale;
+    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());
-
-  // 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);
-  }
+  if (PoisonVals.empty())
+    return;
 
-  // insert initializers for the redzones
-  for (Inst *Init : Initializations) {
-    C.insert(Init);
+  Context.setNext(Context.getCur());
+  Context.insert(ShadowIndexCalc);
+  Context.insert(ShadowLocCalc);
+
+  // Poison redzones
+  std::vector<Entry>::iterator Iter = PoisonVals.begin();
+  for (SizeT Offset = 0; Iter != PoisonVals.end(); Offset += 4) {

[Jim Stichnoth, 2016/08/01 20:38:01]

Instead of magic constant "4", can you make it a static constexpr with a
meaningful name?  e.g. BytesPerWord or RedZoneWords

[tlively, 2016/08/05 18:24:52]

Done.

+    int32_t CurVals[4] = {0};
+    for (int I = 0; I < 4; ++I) {

[Jim Stichnoth, 2016/08/01 20:38:01]

Can you use int32_t or uint32_t instead of int here, for consistency across the
code base?

Also, for silly index variables like this, feel free to name them lower-case
"i", "j", etc.  Even the LLVM code base does this sometimes.

[tlively, 2016/08/05 18:24:53]

Done.

+      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));
   }
+  Context.advanceCur();
+  Context.advanceNext();
 }
 
 void ASanInstrumentation::instrumentCall(LoweringContext &Context,
@@ -327,22 +346,13 @@ bool ASanInstrumentation::isOkGlobalAccess(Operand *Op, SizeT 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) {