diff --git a/llvm/include/llvm/Transforms/Yk/ModuleClone.h b/llvm/include/llvm/Transforms/Yk/ModuleClone.h index 94ad9358395cae5..6f3bb41eaeb6d9f 100644 --- a/llvm/include/llvm/Transforms/Yk/ModuleClone.h +++ b/llvm/include/llvm/Transforms/Yk/ModuleClone.h @@ -3,7 +3,8 @@ #include "llvm/Pass.h" -#define YK_CLONE_PREFIX "__yk_unopt_" +#define YK_UNOPT_PREFIX "__yk_unopt_" +#define YK_UNOPT_MAIN "__yk_unopt_main" #define YK_CLONE_MODULE_CP_COUNT 2 namespace llvm { diff --git a/llvm/include/llvm/Transforms/Yk/ShadowStack.h b/llvm/include/llvm/Transforms/Yk/ShadowStack.h index b131df25fe657de..97aa097f55e8b51 100644 --- a/llvm/include/llvm/Transforms/Yk/ShadowStack.h +++ b/llvm/include/llvm/Transforms/Yk/ShadowStack.h @@ -4,7 +4,7 @@ #include "llvm/Pass.h" namespace llvm { -ModulePass *createYkShadowStackPass(); +ModulePass *createYkShadowStackPass(uint64_t controlPointCount); } // namespace llvm #endif diff --git a/llvm/lib/CodeGen/TargetPassConfig.cpp b/llvm/lib/CodeGen/TargetPassConfig.cpp index 1f90426f97aa4fa..a8189402ea40a91 100644 --- a/llvm/lib/CodeGen/TargetPassConfig.cpp +++ b/llvm/lib/CodeGen/TargetPassConfig.cpp @@ -1162,7 +1162,7 @@ bool TargetPassConfig::addISelPasses() { } if (YkShadowStack) { - addPass(createYkShadowStackPass()); + addPass(createYkShadowStackPass(numberOfControlPoints)); } // We insert the yk control point pass as late as possible. It has to run // before instruction selection (or the machine IR won't reflect our diff --git a/llvm/lib/Transforms/Yk/ModuleClone.cpp b/llvm/lib/Transforms/Yk/ModuleClone.cpp index f9a847f7940067b..f747e2f83fd7e70 100644 --- a/llvm/lib/Transforms/Yk/ModuleClone.cpp +++ b/llvm/lib/Transforms/Yk/ModuleClone.cpp @@ -79,7 +79,7 @@ struct YkModuleClone : public ModulePass { continue; } // Skip already cloned functions or functions with address taken. - if (F.hasAddressTaken() || F.getName().startswith(YK_CLONE_PREFIX)) { + if (F.hasAddressTaken() || F.getName().startswith(YK_UNOPT_PREFIX)) { continue; } ValueToValueMapTy VMap; @@ -96,7 +96,7 @@ struct YkModuleClone : public ModulePass { } // Rename function auto originalName = F.getName().str(); - auto cloneName = Twine(YK_CLONE_PREFIX) + originalName; + auto cloneName = Twine(YK_UNOPT_PREFIX) + originalName; ClonedFunc->setName(cloneName); ClonedFuncs[originalName] = ClonedFunc; } diff --git a/llvm/lib/Transforms/Yk/ShadowStack.cpp b/llvm/lib/Transforms/Yk/ShadowStack.cpp index 1f5b03660707610..60913693cda2035 100644 --- a/llvm/lib/Transforms/Yk/ShadowStack.cpp +++ b/llvm/lib/Transforms/Yk/ShadowStack.cpp @@ -118,9 +118,14 @@ class YkShadowStack : public ModulePass { Type *Int32Ty = nullptr; Type *PointerSizedIntTy = nullptr; +private: + uint64_t controlPointCount; + public: static char ID; - YkShadowStack() : ModulePass(ID) { + + YkShadowStack(uint64_t controlPointCount) + : ModulePass(ID), controlPointCount(controlPointCount) { initializeYkShadowStackPass(*PassRegistry::getPassRegistry()); } @@ -271,6 +276,64 @@ class YkShadowStack : public ModulePass { } } + // Handle main() separatley, since it works slightly differently to other + // functions: it allocates the shadow stack + // + // For main clone: `__yk_unopt_main` we need to update the allocas + // while reusing the same shadow stack pointer as the original main. + // + // Note that since we assuming main() doesn't call main(), we can consider + // the shadow stack disused at the point main() returns. For this reason, + // there's no need to emit a shadow epilogue for main(). + // + // YKFIXME: Investigate languages that don't have/use main as the first + // entry point. + void updateMainFunctions(DataLayout &DL, Module &M, GlobalVariable *SSGlobal, + LLVMContext &Context) { + Function *Main = M.getFunction(MAIN); + if (Main == nullptr || Main->isDeclaration()) { + Context.emitError("Unable to add shadow stack: could not find definition " + "of \"main\" function!"); + return; + } + + // Update the original main function. + std::map MainAllocas; + std::vector MainRets; + size_t SFrameSize = analyseFunction(*Main, DL, MainAllocas, MainRets); + CallInst *Malloc = insertMainPrologue(Main, SSGlobal, SFrameSize); + rewriteAllocas(DL, MainAllocas, Malloc); + + // If we have two control points, we need to update the cloned main + // function as well. + if (controlPointCount == 2) { + auto clonedMainFunctionName = (Twine(YK_UNOPT_PREFIX) + MAIN).str(); + Function *MainClone = M.getFunction(clonedMainFunctionName); + if (MainClone == nullptr || MainClone->isDeclaration()) { + Context.emitError( + "Unable to add shadow stack: could not find definition " + "of \"" + + clonedMainFunctionName + "\" function!"); + return; + } + + std::map MainCloneAllocas; + std::vector MainCloneRets; + size_t SFrameSizeClone = + analyseFunction(*MainClone, DL, MainCloneAllocas, MainCloneRets); + + // Insert a load of SSGlobal at the beginning of MainClone. + BasicBlock &EntryBB = MainClone->getEntryBlock(); + Instruction *FirstNonPHI = EntryBB.getFirstNonPHI(); + IRBuilder<> BuilderClone(FirstNonPHI); + + // Load the current shadow stack pointer from the global variable. + LoadInst *LoadedSSPtr = BuilderClone.CreateLoad(Int8PtrTy, SSGlobal); + // Replace all allocas in MainClone with the loaded shadow stack pointer. + rewriteAllocas(DL, MainCloneAllocas, LoadedSSPtr); + } + } + bool runOnModule(Module &M) override { LLVMContext &Context = M.getContext(); @@ -291,25 +354,7 @@ class YkShadowStack : public ModulePass { SSGlobal->setInitializer( ConstantPointerNull::get(cast(Int8PtrTy))); - // Handle main() separatley, since it works slightly differently to other - // functions: it allocates the shadow stack. - // - // Note that since we assuming main() doesn't call main(), we can consider - // the shadow stack disused at the point main() returns. For this reason, - // there's no need to emit a shadow epilogue for main(). - // - // YKFIXME: Investigate languages that don't have/use main as the first - // entry point. - Function *Main = M.getFunction(MAIN); - if (Main == nullptr || Main->isDeclaration()) { - Context.emitError("Unable to add shadow stack: could not find definition " - "of \"main\" function!"); - } - std::map MainAllocas; - std::vector MainRets; - size_t SFrameSize = analyseFunction(*Main, DL, MainAllocas, MainRets); - CallInst *Malloc = insertMainPrologue(Main, SSGlobal, SFrameSize); - rewriteAllocas(DL, MainAllocas, Malloc); + updateMainFunctions(DL, M, SSGlobal, Context); // Instrument each remaining function with shadow stack code. for (Function &F : M) { @@ -317,9 +362,9 @@ class YkShadowStack : public ModulePass { // skip declarations. continue; } - - if (F.getName() == MAIN || F.getName().startswith(YK_CLONE_PREFIX)) { - // We've handled main already. + // skip already handled main and unopt functions + if (F.getName() == MAIN || F.getName() == YK_UNOPT_MAIN || + !F.getName().startswith(YK_UNOPT_PREFIX)) { continue; } @@ -349,4 +394,6 @@ class YkShadowStack : public ModulePass { char YkShadowStack::ID = 0; INITIALIZE_PASS(YkShadowStack, DEBUG_TYPE, "yk shadowstack", false, false) -ModulePass *llvm::createYkShadowStackPass() { return new YkShadowStack(); } +ModulePass *llvm::createYkShadowStackPass(uint64_t controlPointCount) { + return new YkShadowStack(controlPointCount); +} diff --git a/llvm/lib/Transforms/Yk/StackMaps.cpp b/llvm/lib/Transforms/Yk/StackMaps.cpp index aa11c386224544e..ba1fbabb66783a6 100644 --- a/llvm/lib/Transforms/Yk/StackMaps.cpp +++ b/llvm/lib/Transforms/Yk/StackMaps.cpp @@ -56,7 +56,7 @@ class YkStackmaps : public ModulePass { for (Function &F : M) { if (F.empty()) // skip declarations. continue; - if (F.getName().startswith(YK_CLONE_PREFIX)) // skip cloned functions + if (F.getName().startswith(YK_UNOPT_PREFIX)) // skip cloned functions continue; LivenessAnalysis LA(&F); diff --git a/llvm/lib/YkIR/YkIRWriter.cpp b/llvm/lib/YkIR/YkIRWriter.cpp index e257b2c907e8e20..aecf54722d2964c 100644 --- a/llvm/lib/YkIR/YkIRWriter.cpp +++ b/llvm/lib/YkIR/YkIRWriter.cpp @@ -1795,7 +1795,7 @@ class YkIRWriter { int functionCount = 0; for (llvm::Function &F : M) { // Skip cloned functions - if (!StringRef(F.getName()).startswith(YK_CLONE_PREFIX)) { + if (!StringRef(F.getName()).startswith(YK_UNOPT_PREFIX)) { FunctionIndexMap[&F] = functionCount; functionCount++; } @@ -1805,7 +1805,7 @@ class YkIRWriter { // funcs: for (llvm::Function &F : M) { // Skip cloned functions - if (StringRef(F.getName()).startswith(YK_CLONE_PREFIX)) { + if (StringRef(F.getName()).startswith(YK_UNOPT_PREFIX)) { continue; } serialiseFunc(F);