Reuse shadowstack in cloned functions.

ykjit · Jan 20, 2025 · 564318a · 564318a
1 parent 7386572
commit 564318a
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 23 deletions.
diff --git 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
@@ -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/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<AllocaInst *, size_t> MainAllocas;
+    std::vector<ReturnInst *> MainRets;
+    size_t SFrameSize = analyseFunction(*Main, DL, MainAllocas, MainRets);
+    CallInst *Malloc = insertMainPrologue(Main, SSGlobal, SFrameSize);
+    rewriteAllocas(DL, MainAllocas, Malloc);
+
+    // Update the cloned main function.
+    if (controlPointCount == 2) {
+      auto clonedMainFunctionName = (Twine(YK_CLONE_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<AllocaInst *, size_t> MainCloneAllocas;
+      std::vector<ReturnInst *> 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,
+                                                      "loaded_shadowstack_ptr");
+      // 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<PointerType>(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<AllocaInst *, size_t> MainAllocas;
-    std::vector<ReturnInst *> 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) {
@@ -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);
+}