|
| 1 | +//===- ShadowStack.cpp - Pass to add shadow stacks to the AOT module --===// |
| 2 | +// |
| 3 | +// Add shadow stacks to store variables that may have their references taken. |
| 4 | +// Storing such variables on a shadow stack allows AOT to share them with |
| 5 | +// compiled traces, and back (i.e. references created inside a trace will still |
| 6 | +// be valid when we return from the trace via deoptimisation). |
| 7 | +// YKFIXME: This can be optimised by only putting variables on the shadow stack |
| 8 | +// that actually have their reference taken. |
| 9 | + |
| 10 | +#include "llvm/Transforms/Yk/ShadowStack.h" |
| 11 | +#include "llvm/IR/BasicBlock.h" |
| 12 | +#include "llvm/IR/DataLayout.h" |
| 13 | +#include "llvm/IR/Function.h" |
| 14 | +#include "llvm/IR/IRBuilder.h" |
| 15 | +#include "llvm/IR/Instructions.h" |
| 16 | +#include "llvm/IR/Module.h" |
| 17 | +#include "llvm/IR/Value.h" |
| 18 | +#include "llvm/IR/Verifier.h" |
| 19 | +#include "llvm/InitializePasses.h" |
| 20 | +#include "llvm/Pass.h" |
| 21 | +#include "llvm/Transforms/Yk/LivenessAnalysis.h" |
| 22 | + |
| 23 | +#include <map> |
| 24 | + |
| 25 | +#define DEBUG_TYPE "yk-shadowstack" |
| 26 | +#define YK_MT_NEW "yk_mt_new" |
| 27 | +#define G_SHADOW_STACK "shadowstack_0" |
| 28 | +// The size of the shadow stack. Defaults to 1MB. |
| 29 | +// YKFIXME: Make this adjustable by a compiler flag. |
| 30 | +#define SHADOW_STACK_SIZE 1000000 |
| 31 | + |
| 32 | +using namespace llvm; |
| 33 | + |
| 34 | +namespace llvm { |
| 35 | +void initializeYkShadowStackPass(PassRegistry &); |
| 36 | +} // namespace llvm |
| 37 | + |
| 38 | +namespace { |
| 39 | +class YkShadowStack : public ModulePass { |
| 40 | +public: |
| 41 | + static char ID; |
| 42 | + YkShadowStack() : ModulePass(ID) { |
| 43 | + initializeYkShadowStackPass(*PassRegistry::getPassRegistry()); |
| 44 | + } |
| 45 | + |
| 46 | + // Checks whether the given instruction is the alloca of the call to |
| 47 | + // `yk_mt_new`. |
| 48 | + bool isYkMTNewAlloca(Instruction *I) { |
| 49 | + for (User *U : I->users()) { |
| 50 | + if (U && isa<StoreInst>(U)) { |
| 51 | + Value *V = cast<StoreInst>(U)->getValueOperand(); |
| 52 | + if (isa<CallInst>(V)) { |
| 53 | + CallInst *CI = cast<CallInst>(V); |
| 54 | + if (CI->isInlineAsm()) |
| 55 | + return false; |
| 56 | + if (!CI->getCalledFunction()) |
| 57 | + return false; |
| 58 | + return (CI->getCalledFunction()->getName() == YK_MT_NEW); |
| 59 | + } |
| 60 | + } |
| 61 | + } |
| 62 | + return false; |
| 63 | + } |
| 64 | + |
| 65 | + bool runOnModule(Module &M) override { |
| 66 | + LLVMContext &Context = M.getContext(); |
| 67 | + |
| 68 | + DataLayout DL(&M); |
| 69 | + Type *Int8Ty = Type::getInt8Ty(Context); |
| 70 | + Type *Int32Ty = Type::getInt32Ty(Context); |
| 71 | + Type *PointerSizedIntTy = DL.getIntPtrType(Context); |
| 72 | + Type *Int8PtrTy = Type::getInt8PtrTy(Context); |
| 73 | + |
| 74 | + // Create a global variable which will store the pointer to the heap memory |
| 75 | + // allocated for the shadow stack. |
| 76 | + Constant *GShadowStackPtr = M.getOrInsertGlobal(G_SHADOW_STACK, Int8PtrTy); |
| 77 | + GlobalVariable *GVar = M.getNamedGlobal(G_SHADOW_STACK); |
| 78 | + GVar->setInitializer( |
| 79 | + ConstantPointerNull::get(cast<PointerType>(Int8PtrTy))); |
| 80 | + |
| 81 | + // We only need to create one shadow stack per module so we'll do this |
| 82 | + // inside the module's entry point. |
| 83 | + // YKFIXME: Investigate languages that don't have/use main as the first |
| 84 | + // entry point. |
| 85 | + Function *Main = M.getFunction("main"); |
| 86 | + if (Main == nullptr) { |
| 87 | + Context.emitError( |
| 88 | + "Unable to add shadow stack: could not find \"main\" function!"); |
| 89 | + return false; |
| 90 | + } |
| 91 | + Instruction *First = Main->getEntryBlock().getFirstNonPHI(); |
| 92 | + IRBuilder<> Builder(First); |
| 93 | + |
| 94 | + // Now create some memory on the heap for the shadow stack. |
| 95 | + FunctionCallee MF = |
| 96 | + M.getOrInsertFunction("malloc", Int8PtrTy, PointerSizedIntTy); |
| 97 | + CallInst *Malloc = Builder.CreateCall( |
| 98 | + MF, {ConstantInt::get(PointerSizedIntTy, SHADOW_STACK_SIZE)}, ""); |
| 99 | + Builder.CreateStore(Malloc, GShadowStackPtr); |
| 100 | + |
| 101 | + Value *SSPtr; |
| 102 | + for (Function &F : M) { |
| 103 | + if (F.empty()) // skip declarations. |
| 104 | + continue; |
| 105 | + |
| 106 | + if (&F != Main) { |
| 107 | + // At the top of each function in the module, load the heap pointer |
| 108 | + // from the global shadow stack variable. |
| 109 | + Builder.SetInsertPoint(F.getEntryBlock().getFirstNonPHI()); |
| 110 | + SSPtr = Builder.CreateLoad(Int8PtrTy, GShadowStackPtr); |
| 111 | + } else { |
| 112 | + SSPtr = cast<Value>(Malloc); |
| 113 | + } |
| 114 | + |
| 115 | + size_t Offset = 0; |
| 116 | + // Remember which allocas were replaced, so we can remove them later in |
| 117 | + // one swoop. Removing them here messes up the loop. |
| 118 | + std::vector<Instruction *> RemoveAllocas; |
| 119 | + for (BasicBlock &BB : F) { |
| 120 | + for (Instruction &I : BB) { |
| 121 | + if (isa<AllocaInst>(I)) { |
| 122 | + // Replace allocas with pointers into the shadow stack. |
| 123 | + AllocaInst &AI = cast<AllocaInst>(I); |
| 124 | + if (isYkMTNewAlloca(&AI)) { |
| 125 | + // The variable created by `yk_mt_new` will never be traced, so |
| 126 | + // there's no need to store it on the shadow stack. |
| 127 | + continue; |
| 128 | + } |
| 129 | + if (isa<StructType>(AI.getAllocatedType())) { |
| 130 | + StructType *ST = cast<StructType>(AI.getAllocatedType()); |
| 131 | + // Some yk specific variables that will never be traced and thus |
| 132 | + // can live happily on the normal stack. |
| 133 | + // YKFIXME: This is somewhat fragile since `struct.YkLocation` is |
| 134 | + // a name given by LLVM which could theoretically change. Luckily, |
| 135 | + // this should all go away once we only move variables to the |
| 136 | + // shadowstack that have their reference taken. |
| 137 | + if (ST->getName() == "YkCtrlPointVars" || |
| 138 | + ST->getName() == "struct.YkLocation") { |
| 139 | + continue; |
| 140 | + } |
| 141 | + } |
| 142 | + Builder.SetInsertPoint(&I); |
| 143 | + auto AllocaSizeInBits = AI.getAllocationSizeInBits(DL); |
| 144 | + if (!AllocaSizeInBits) { |
| 145 | + // YKFIXME: Deal with functions where the stack size isn't know at |
| 146 | + // compile time, e.g. when `alloca` is used. |
| 147 | + Context.emitError("Unable to add shadow stack: function has " |
| 148 | + "dynamically sized stack!"); |
| 149 | + return false; |
| 150 | + } |
| 151 | + // Calculate this `AllocaInst`s size and create a replacement |
| 152 | + // pointer into the shadow stack. |
| 153 | + size_t AllocaSize = *AllocaSizeInBits / 8; |
| 154 | + GetElementPtrInst *GEP = GetElementPtrInst::Create( |
| 155 | + Int8Ty, SSPtr, {ConstantInt::get(Int32Ty, Offset)}, "", |
| 156 | + cast<Instruction>(&AI)); |
| 157 | + Builder.SetInsertPoint(GEP); |
| 158 | + Builder.CreateBitCast(GEP, AI.getAllocatedType()->getPointerTo()); |
| 159 | + cast<Value>(I).replaceAllUsesWith(GEP); |
| 160 | + RemoveAllocas.push_back(cast<Instruction>(&AI)); |
| 161 | + Offset += AllocaSize; |
| 162 | + } else if (isa<CallInst>(I)) { |
| 163 | + // When we see a call, we need make space for a new stack frame. We |
| 164 | + // do this by simply adjusting the pointer stored in the global |
| 165 | + // shadow stack. When the function returns the global is reset. This |
| 166 | + // is similar to how the RSP is adjusted inside the |
| 167 | + // prologue/epilogue of a function, but here the prologue/epilogue |
| 168 | + // are handled by the caller. |
| 169 | + CallInst &CI = cast<CallInst>(I); |
| 170 | + if (&CI == Malloc) { |
| 171 | + // Don't do this for the `malloc` that created the shadow stack. |
| 172 | + continue; |
| 173 | + } |
| 174 | + // Inline asm can't be traced. |
| 175 | + if (CI.isInlineAsm()) { |
| 176 | + continue; |
| 177 | + } |
| 178 | + |
| 179 | + // YKFIXME: Skip functions that are marked with `yk_outline` |
| 180 | + // (as those won't be traced and thus don't require a shadow |
| 181 | + // stack). |
| 182 | + // YKFIXME: Skip functions (direct or indirect) that we don't have |
| 183 | + // IR for. |
| 184 | + |
| 185 | + if (CI.getCalledFunction()) { |
| 186 | + if (CI.getCalledFunction()->getName() == "pthread_create") { |
| 187 | + // The global shadow stack needs to be thread local. In each new |
| 188 | + // thread created we need to malloc a new shadow stack and |
| 189 | + // assign it to that threads shadow stack global. Note: it's not |
| 190 | + // enough to look for `pthread_create` in here, as this call |
| 191 | + // could be hidden inside an external library. |
| 192 | + Context.emitError( |
| 193 | + "Unable to add shadow stack: No support for threads yet!"); |
| 194 | + return false; |
| 195 | + } |
| 196 | + // Skip some known intrinsics. YKFIXME: Is there a more general |
| 197 | + // solution, e.g. skip all intrinsics? |
| 198 | + else if (CI.getCalledFunction()->getName() == |
| 199 | + "llvm.experimental.stackmap") { |
| 200 | + continue; |
| 201 | + } else if (CI.getCalledFunction()->getName() == |
| 202 | + "llvm.dbg.declare") { |
| 203 | + continue; |
| 204 | + } |
| 205 | + } |
| 206 | + |
| 207 | + // Adjust shadow stack pointer before a call, and reset it back to |
| 208 | + // its previous value upon returning. |
| 209 | + GetElementPtrInst *GEP = GetElementPtrInst::Create( |
| 210 | + Int8Ty, SSPtr, {ConstantInt::get(Int32Ty, Offset)}, "", &I); |
| 211 | + Builder.SetInsertPoint(&I); |
| 212 | + Builder.CreateStore(GEP, GShadowStackPtr); |
| 213 | + Builder.SetInsertPoint(I.getNextNonDebugInstruction()); |
| 214 | + Builder.CreateStore(SSPtr, GShadowStackPtr); |
| 215 | + } |
| 216 | + } |
| 217 | + } |
| 218 | + for (Instruction *I : RemoveAllocas) { |
| 219 | + I->removeFromParent(); |
| 220 | + } |
| 221 | + RemoveAllocas.clear(); |
| 222 | + } |
| 223 | + |
| 224 | +#ifndef NDEBUG |
| 225 | + // Our pass runs after LLVM normally does its verify pass. In debug builds |
| 226 | + // we run it again to check that our pass is generating valid IR. |
| 227 | + if (verifyModule(M, &errs())) { |
| 228 | + Context.emitError("ShadowStack insertion pass generated invalid IR!"); |
| 229 | + return false; |
| 230 | + } |
| 231 | +#endif |
| 232 | + return true; |
| 233 | + } |
| 234 | +}; |
| 235 | +} // namespace |
| 236 | + |
| 237 | +char YkShadowStack::ID = 0; |
| 238 | +INITIALIZE_PASS(YkShadowStack, DEBUG_TYPE, "yk shadowstack", false, false) |
| 239 | + |
| 240 | +ModulePass *llvm::createYkShadowStackPass() { return new YkShadowStack(); } |
0 commit comments