1 //===-- ShadowStackGCLowering.cpp - Custom lowering for shadow-stack gc ---===//
3 // The LLVM Compiler Infrastructure
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
8 //===----------------------------------------------------------------------===//
10 // This file contains the custom lowering code required by the shadow-stack GC
13 //===----------------------------------------------------------------------===//
15 #include "llvm/CodeGen/Passes.h"
16 #include "llvm/ADT/StringExtras.h"
17 #include "llvm/CodeGen/GCStrategy.h"
18 #include "llvm/IR/CallSite.h"
19 #include "llvm/IR/IRBuilder.h"
20 #include "llvm/IR/IntrinsicInst.h"
21 #include "llvm/IR/Module.h"
25 #define DEBUG_TYPE "shadowstackgclowering"
29 class ShadowStackGCLowering : public FunctionPass {
30 /// RootChain - This is the global linked-list that contains the chain of GC
34 /// StackEntryTy - Abstract type of a link in the shadow stack.
36 StructType *StackEntryTy;
37 StructType *FrameMapTy;
39 /// Roots - GC roots in the current function. Each is a pair of the
40 /// intrinsic call and its corresponding alloca.
41 std::vector<std::pair<CallInst *, AllocaInst *>> Roots;
45 ShadowStackGCLowering();
47 bool doInitialization(Module &M) override;
48 bool runOnFunction(Function &F) override;
51 bool IsNullValue(Value *V);
52 Constant *GetFrameMap(Function &F);
53 Type *GetConcreteStackEntryType(Function &F);
54 void CollectRoots(Function &F);
55 static GetElementPtrInst *CreateGEP(LLVMContext &Context, IRBuilder<> &B,
56 Value *BasePtr, int Idx1,
58 static GetElementPtrInst *CreateGEP(LLVMContext &Context, IRBuilder<> &B,
59 Value *BasePtr, int Idx1, int Idx2,
64 INITIALIZE_PASS_BEGIN(ShadowStackGCLowering, "shadow-stack-gc-lowering",
65 "Shadow Stack GC Lowering", false, false)
66 INITIALIZE_PASS_DEPENDENCY(GCModuleInfo)
67 INITIALIZE_PASS_END(ShadowStackGCLowering, "shadow-stack-gc-lowering",
68 "Shadow Stack GC Lowering", false, false)
70 FunctionPass *llvm::createShadowStackGCLoweringPass() { return new ShadowStackGCLowering(); }
72 char ShadowStackGCLowering::ID = 0;
74 ShadowStackGCLowering::ShadowStackGCLowering()
75 : FunctionPass(ID), Head(nullptr), StackEntryTy(nullptr),
77 initializeShadowStackGCLoweringPass(*PassRegistry::getPassRegistry());
81 /// EscapeEnumerator - This is a little algorithm to find all escape points
82 /// from a function so that "finally"-style code can be inserted. In addition
83 /// to finding the existing return and unwind instructions, it also (if
84 /// necessary) transforms any call instructions into invokes and sends them to
87 /// It's wrapped up in a state machine using the same transform C# uses for
88 /// 'yield return' enumerators, This transform allows it to be non-allocating.
89 class EscapeEnumerator {
91 const char *CleanupBBName;
95 Function::iterator StateBB, StateE;
99 EscapeEnumerator(Function &F, const char *N = "cleanup")
100 : F(F), CleanupBBName(N), State(0), Builder(F.getContext()) {}
102 IRBuilder<> *Next() {
113 // Find all 'return', 'resume', and 'unwind' instructions.
114 while (StateBB != StateE) {
115 BasicBlock *CurBB = StateBB++;
117 // Branches and invokes do not escape, only unwind, resume, and return
119 TerminatorInst *TI = CurBB->getTerminator();
120 if (!isa<ReturnInst>(TI) && !isa<ResumeInst>(TI))
123 Builder.SetInsertPoint(TI->getParent(), TI);
129 // Find all 'call' instructions.
130 SmallVector<Instruction *, 16> Calls;
131 for (Function::iterator BB = F.begin(), E = F.end(); BB != E; ++BB)
132 for (BasicBlock::iterator II = BB->begin(), EE = BB->end(); II != EE;
134 if (CallInst *CI = dyn_cast<CallInst>(II))
135 if (!CI->getCalledFunction() ||
136 !CI->getCalledFunction()->getIntrinsicID())
142 // Create a cleanup block.
143 LLVMContext &C = F.getContext();
144 BasicBlock *CleanupBB = BasicBlock::Create(C, CleanupBBName, &F);
146 StructType::get(Type::getInt8PtrTy(C), Type::getInt32Ty(C), nullptr);
147 Constant *PersFn = F.getParent()->getOrInsertFunction(
148 "__gcc_personality_v0", FunctionType::get(Type::getInt32Ty(C), true));
149 LandingPadInst *LPad =
150 LandingPadInst::Create(ExnTy, PersFn, 1, "cleanup.lpad", CleanupBB);
151 LPad->setCleanup(true);
152 ResumeInst *RI = ResumeInst::Create(LPad, CleanupBB);
154 // Transform the 'call' instructions into 'invoke's branching to the
155 // cleanup block. Go in reverse order to make prettier BB names.
156 SmallVector<Value *, 16> Args;
157 for (unsigned I = Calls.size(); I != 0;) {
158 CallInst *CI = cast<CallInst>(Calls[--I]);
160 // Split the basic block containing the function call.
161 BasicBlock *CallBB = CI->getParent();
163 CallBB->splitBasicBlock(CI, CallBB->getName() + ".cont");
165 // Remove the unconditional branch inserted at the end of CallBB.
166 CallBB->getInstList().pop_back();
167 NewBB->getInstList().remove(CI);
169 // Create a new invoke instruction.
172 Args.append(CS.arg_begin(), CS.arg_end());
175 InvokeInst::Create(CI->getCalledValue(), NewBB, CleanupBB, Args,
176 CI->getName(), CallBB);
177 II->setCallingConv(CI->getCallingConv());
178 II->setAttributes(CI->getAttributes());
179 CI->replaceAllUsesWith(II);
183 Builder.SetInsertPoint(RI->getParent(), RI);
191 Constant *ShadowStackGCLowering::GetFrameMap(Function &F) {
192 // doInitialization creates the abstract type of this value.
193 Type *VoidPtr = Type::getInt8PtrTy(F.getContext());
195 // Truncate the ShadowStackDescriptor if some metadata is null.
196 unsigned NumMeta = 0;
197 SmallVector<Constant *, 16> Metadata;
198 for (unsigned I = 0; I != Roots.size(); ++I) {
199 Constant *C = cast<Constant>(Roots[I].first->getArgOperand(1));
200 if (!C->isNullValue())
202 Metadata.push_back(ConstantExpr::getBitCast(C, VoidPtr));
204 Metadata.resize(NumMeta);
206 Type *Int32Ty = Type::getInt32Ty(F.getContext());
208 Constant *BaseElts[] = {
209 ConstantInt::get(Int32Ty, Roots.size(), false),
210 ConstantInt::get(Int32Ty, NumMeta, false),
213 Constant *DescriptorElts[] = {
214 ConstantStruct::get(FrameMapTy, BaseElts),
215 ConstantArray::get(ArrayType::get(VoidPtr, NumMeta), Metadata)};
217 Type *EltTys[] = {DescriptorElts[0]->getType(), DescriptorElts[1]->getType()};
218 StructType *STy = StructType::create(EltTys, "gc_map." + utostr(NumMeta));
220 Constant *FrameMap = ConstantStruct::get(STy, DescriptorElts);
222 // FIXME: Is this actually dangerous as WritingAnLLVMPass.html claims? Seems
223 // that, short of multithreaded LLVM, it should be safe; all that is
224 // necessary is that a simple Module::iterator loop not be invalidated.
225 // Appending to the GlobalVariable list is safe in that sense.
227 // All of the output passes emit globals last. The ExecutionEngine
228 // explicitly supports adding globals to the module after
231 // Still, if it isn't deemed acceptable, then this transformation needs
232 // to be a ModulePass (which means it cannot be in the 'llc' pipeline
233 // (which uses a FunctionPassManager (which segfaults (not asserts) if
234 // provided a ModulePass))).
235 Constant *GV = new GlobalVariable(*F.getParent(), FrameMap->getType(), true,
236 GlobalVariable::InternalLinkage, FrameMap,
237 "__gc_" + F.getName());
239 Constant *GEPIndices[2] = {
240 ConstantInt::get(Type::getInt32Ty(F.getContext()), 0),
241 ConstantInt::get(Type::getInt32Ty(F.getContext()), 0)};
242 return ConstantExpr::getGetElementPtr(GV, GEPIndices);
245 Type *ShadowStackGCLowering::GetConcreteStackEntryType(Function &F) {
246 // doInitialization creates the generic version of this type.
247 std::vector<Type *> EltTys;
248 EltTys.push_back(StackEntryTy);
249 for (size_t I = 0; I != Roots.size(); I++)
250 EltTys.push_back(Roots[I].second->getAllocatedType());
252 return StructType::create(EltTys, "gc_stackentry." + F.getName().str());
255 /// doInitialization - If this module uses the GC intrinsics, find them now. If
257 bool ShadowStackGCLowering::doInitialization(Module &M) {
259 for (Function &F : M) {
260 if (F.hasGC() && F.getGC() == std::string("shadow-stack")) {
269 // int32_t NumRoots; // Number of roots in stack frame.
270 // int32_t NumMeta; // Number of metadata descriptors. May be < NumRoots.
271 // void *Meta[]; // May be absent for roots without metadata.
273 std::vector<Type *> EltTys;
274 // 32 bits is ok up to a 32GB stack frame. :)
275 EltTys.push_back(Type::getInt32Ty(M.getContext()));
276 // Specifies length of variable length array.
277 EltTys.push_back(Type::getInt32Ty(M.getContext()));
278 FrameMapTy = StructType::create(EltTys, "gc_map");
279 PointerType *FrameMapPtrTy = PointerType::getUnqual(FrameMapTy);
281 // struct StackEntry {
282 // ShadowStackEntry *Next; // Caller's stack entry.
283 // FrameMap *Map; // Pointer to constant FrameMap.
284 // void *Roots[]; // Stack roots (in-place array, so we pretend).
287 StackEntryTy = StructType::create(M.getContext(), "gc_stackentry");
290 EltTys.push_back(PointerType::getUnqual(StackEntryTy));
291 EltTys.push_back(FrameMapPtrTy);
292 StackEntryTy->setBody(EltTys);
293 PointerType *StackEntryPtrTy = PointerType::getUnqual(StackEntryTy);
295 // Get the root chain if it already exists.
296 Head = M.getGlobalVariable("llvm_gc_root_chain");
298 // If the root chain does not exist, insert a new one with linkonce
300 Head = new GlobalVariable(
301 M, StackEntryPtrTy, false, GlobalValue::LinkOnceAnyLinkage,
302 Constant::getNullValue(StackEntryPtrTy), "llvm_gc_root_chain");
303 } else if (Head->hasExternalLinkage() && Head->isDeclaration()) {
304 Head->setInitializer(Constant::getNullValue(StackEntryPtrTy));
305 Head->setLinkage(GlobalValue::LinkOnceAnyLinkage);
311 bool ShadowStackGCLowering::IsNullValue(Value *V) {
312 if (Constant *C = dyn_cast<Constant>(V))
313 return C->isNullValue();
317 void ShadowStackGCLowering::CollectRoots(Function &F) {
318 // FIXME: Account for original alignment. Could fragment the root array.
319 // Approach 1: Null initialize empty slots at runtime. Yuck.
320 // Approach 2: Emit a map of the array instead of just a count.
322 assert(Roots.empty() && "Not cleaned up?");
324 SmallVector<std::pair<CallInst *, AllocaInst *>, 16> MetaRoots;
326 for (Function::iterator BB = F.begin(), E = F.end(); BB != E; ++BB)
327 for (BasicBlock::iterator II = BB->begin(), E = BB->end(); II != E;)
328 if (IntrinsicInst *CI = dyn_cast<IntrinsicInst>(II++))
329 if (Function *F = CI->getCalledFunction())
330 if (F->getIntrinsicID() == Intrinsic::gcroot) {
331 std::pair<CallInst *, AllocaInst *> Pair = std::make_pair(
333 cast<AllocaInst>(CI->getArgOperand(0)->stripPointerCasts()));
334 if (IsNullValue(CI->getArgOperand(1)))
335 Roots.push_back(Pair);
337 MetaRoots.push_back(Pair);
340 // Number roots with metadata (usually empty) at the beginning, so that the
341 // FrameMap::Meta array can be elided.
342 Roots.insert(Roots.begin(), MetaRoots.begin(), MetaRoots.end());
345 GetElementPtrInst *ShadowStackGCLowering::CreateGEP(LLVMContext &Context,
346 IRBuilder<> &B, Value *BasePtr,
349 Value *Indices[] = {ConstantInt::get(Type::getInt32Ty(Context), 0),
350 ConstantInt::get(Type::getInt32Ty(Context), Idx),
351 ConstantInt::get(Type::getInt32Ty(Context), Idx2)};
352 Value *Val = B.CreateGEP(BasePtr, Indices, Name);
354 assert(isa<GetElementPtrInst>(Val) && "Unexpected folded constant");
356 return dyn_cast<GetElementPtrInst>(Val);
359 GetElementPtrInst *ShadowStackGCLowering::CreateGEP(LLVMContext &Context,
360 IRBuilder<> &B, Value *BasePtr,
361 int Idx, const char *Name) {
362 Value *Indices[] = {ConstantInt::get(Type::getInt32Ty(Context), 0),
363 ConstantInt::get(Type::getInt32Ty(Context), Idx)};
364 Value *Val = B.CreateGEP(BasePtr, Indices, Name);
366 assert(isa<GetElementPtrInst>(Val) && "Unexpected folded constant");
368 return dyn_cast<GetElementPtrInst>(Val);
371 /// runOnFunction - Insert code to maintain the shadow stack.
372 bool ShadowStackGCLowering::runOnFunction(Function &F) {
373 // Quick exit for functions that do not use the shadow stack GC.
375 F.getGC() != std::string("shadow-stack"))
378 LLVMContext &Context = F.getContext();
380 // Find calls to llvm.gcroot.
383 // If there are no roots in this function, then there is no need to add a
384 // stack map entry for it.
388 // Build the constant map and figure the type of the shadow stack entry.
389 Value *FrameMap = GetFrameMap(F);
390 Type *ConcreteStackEntryTy = GetConcreteStackEntryType(F);
392 // Build the shadow stack entry at the very start of the function.
393 BasicBlock::iterator IP = F.getEntryBlock().begin();
394 IRBuilder<> AtEntry(IP->getParent(), IP);
396 Instruction *StackEntry =
397 AtEntry.CreateAlloca(ConcreteStackEntryTy, nullptr, "gc_frame");
399 while (isa<AllocaInst>(IP))
401 AtEntry.SetInsertPoint(IP->getParent(), IP);
403 // Initialize the map pointer and load the current head of the shadow stack.
404 Instruction *CurrentHead = AtEntry.CreateLoad(Head, "gc_currhead");
405 Instruction *EntryMapPtr =
406 CreateGEP(Context, AtEntry, StackEntry, 0, 1, "gc_frame.map");
407 AtEntry.CreateStore(FrameMap, EntryMapPtr);
409 // After all the allocas...
410 for (unsigned I = 0, E = Roots.size(); I != E; ++I) {
411 // For each root, find the corresponding slot in the aggregate...
412 Value *SlotPtr = CreateGEP(Context, AtEntry, StackEntry, 1 + I, "gc_root");
414 // And use it in lieu of the alloca.
415 AllocaInst *OriginalAlloca = Roots[I].second;
416 SlotPtr->takeName(OriginalAlloca);
417 OriginalAlloca->replaceAllUsesWith(SlotPtr);
420 // Move past the original stores inserted by GCStrategy::InitRoots. This isn't
421 // really necessary (the collector would never see the intermediate state at
422 // runtime), but it's nicer not to push the half-initialized entry onto the
424 while (isa<StoreInst>(IP))
426 AtEntry.SetInsertPoint(IP->getParent(), IP);
428 // Push the entry onto the shadow stack.
429 Instruction *EntryNextPtr =
430 CreateGEP(Context, AtEntry, StackEntry, 0, 0, "gc_frame.next");
431 Instruction *NewHeadVal =
432 CreateGEP(Context, AtEntry, StackEntry, 0, "gc_newhead");
433 AtEntry.CreateStore(CurrentHead, EntryNextPtr);
434 AtEntry.CreateStore(NewHeadVal, Head);
436 // For each instruction that escapes...
437 EscapeEnumerator EE(F, "gc_cleanup");
438 while (IRBuilder<> *AtExit = EE.Next()) {
439 // Pop the entry from the shadow stack. Don't reuse CurrentHead from
440 // AtEntry, since that would make the value live for the entire function.
441 Instruction *EntryNextPtr2 =
442 CreateGEP(Context, *AtExit, StackEntry, 0, 0, "gc_frame.next");
443 Value *SavedHead = AtExit->CreateLoad(EntryNextPtr2, "gc_savedhead");
444 AtExit->CreateStore(SavedHead, Head);
447 // Delete the original allocas (which are no longer used) and the intrinsic
448 // calls (which are no longer valid). Doing this last avoids invalidating
450 for (unsigned I = 0, E = Roots.size(); I != E; ++I) {
451 Roots[I].first->eraseFromParent();
452 Roots[I].second->eraseFromParent();