1//===- ShadowStackGCLowering.cpp - Custom lowering for shadow-stack gc ----===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This file contains the custom lowering code required by the shadow-stack GC
10// strategy.
11//
12// This pass implements the code transformation described in this paper:
13// "Accurate Garbage Collection in an Uncooperative Environment"
14// Fergus Henderson, ISMM, 2002
15//
16//===----------------------------------------------------------------------===//
17
18#include "llvm/CodeGen/ShadowStackGCLowering.h"
19#include "llvm/ADT/SmallVector.h"
20#include "llvm/ADT/StringExtras.h"
21#include "llvm/Analysis/DomTreeUpdater.h"
22#include "llvm/CodeGen/GCMetadata.h"
23#include "llvm/CodeGen/Passes.h"
24#include "llvm/IR/BasicBlock.h"
25#include "llvm/IR/Constant.h"
26#include "llvm/IR/Constants.h"
27#include "llvm/IR/DerivedTypes.h"
28#include "llvm/IR/Dominators.h"
29#include "llvm/IR/Function.h"
30#include "llvm/IR/GlobalValue.h"
31#include "llvm/IR/GlobalVariable.h"
32#include "llvm/IR/IRBuilder.h"
33#include "llvm/IR/Instructions.h"
34#include "llvm/IR/IntrinsicInst.h"
35#include "llvm/IR/Intrinsics.h"
36#include "llvm/IR/Module.h"
37#include "llvm/IR/Type.h"
38#include "llvm/IR/Value.h"
39#include "llvm/InitializePasses.h"
40#include "llvm/Pass.h"
41#include "llvm/Support/Casting.h"
42#include "llvm/Transforms/Utils/EscapeEnumerator.h"
43#include <cassert>
44#include <optional>
45#include <utility>
46#include <vector>
47
48using namespace llvm;
49
50#define DEBUG_TYPE "shadow-stack-gc-lowering"
51
52namespace {
53
54class ShadowStackGCLoweringImpl {
55 /// RootChain - This is the global linked-list that contains the chain of GC
56 /// roots.
57 GlobalVariable *Head = nullptr;
58
59 /// StackEntryTy - Abstract type of a link in the shadow stack.
60 StructType *StackEntryTy = nullptr;
61 StructType *FrameMapTy = nullptr;
62
63 /// Roots - GC roots in the current function. Each is a pair of the
64 /// intrinsic call and its corresponding alloca.
65 std::vector<std::pair<CallInst *, AllocaInst *>> Roots;
66
67public:
68 ShadowStackGCLoweringImpl() = default;
69
70 bool doInitialization(Module &M);
71 bool runOnFunction(Function &F, DomTreeUpdater *DTU);
72
73private:
74 bool IsNullValue(Value *V);
75 Constant *GetFrameMap(Function &F);
76 Type *GetConcreteStackEntryType(Function &F);
77 void CollectRoots(Function &F);
78
79 static GetElementPtrInst *CreateGEP(LLVMContext &Context, IRBuilder<> &B,
80 Type *Ty, Value *BasePtr, int Idx1,
81 const char *Name);
82 static GetElementPtrInst *CreateGEP(LLVMContext &Context, IRBuilder<> &B,
83 Type *Ty, Value *BasePtr, int Idx1, int Idx2,
84 const char *Name);
85};
86
87class ShadowStackGCLowering : public FunctionPass {
88 ShadowStackGCLoweringImpl Impl;
89
90public:
91 static char ID;
92
93 ShadowStackGCLowering();
94
95 bool doInitialization(Module &M) override { return Impl.doInitialization(M); }
96 void getAnalysisUsage(AnalysisUsage &AU) const override {
97 AU.addPreserved<DominatorTreeWrapperPass>();
98 }
99 bool runOnFunction(Function &F) override {
100 std::optional<DomTreeUpdater> DTU;
101 if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
102 DTU.emplace(args&: DTWP->getDomTree(), args: DomTreeUpdater::UpdateStrategy::Lazy);
103 return Impl.runOnFunction(F, DTU: DTU ? &*DTU : nullptr);
104 }
105};
106
107} // end anonymous namespace
108
109PreservedAnalyses ShadowStackGCLoweringPass::run(Module &M,
110 ModuleAnalysisManager &MAM) {
111 auto &Map = MAM.getResult<CollectorMetadataAnalysis>(IR&: M);
112 if (!Map.contains(GCName: "shadow-stack"))
113 return PreservedAnalyses::all();
114
115 ShadowStackGCLoweringImpl Impl;
116 bool Changed = Impl.doInitialization(M);
117 for (auto &F : M) {
118 auto &FAM =
119 MAM.getResult<FunctionAnalysisManagerModuleProxy>(IR&: M).getManager();
120 auto *DT = FAM.getCachedResult<DominatorTreeAnalysis>(IR&: F);
121 DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
122 Changed |= Impl.runOnFunction(F, DTU: DT ? &DTU : nullptr);
123 }
124
125 if (!Changed)
126 return PreservedAnalyses::all();
127 PreservedAnalyses PA;
128 PA.preserve<DominatorTreeAnalysis>();
129 return PA;
130}
131
132char ShadowStackGCLowering::ID = 0;
133char &llvm::ShadowStackGCLoweringID = ShadowStackGCLowering::ID;
134
135INITIALIZE_PASS_BEGIN(ShadowStackGCLowering, DEBUG_TYPE,
136 "Shadow Stack GC Lowering", false, false)
137INITIALIZE_PASS_DEPENDENCY(GCModuleInfo)
138INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
139INITIALIZE_PASS_END(ShadowStackGCLowering, DEBUG_TYPE,
140 "Shadow Stack GC Lowering", false, false)
141
142FunctionPass *llvm::createShadowStackGCLoweringPass() { return new ShadowStackGCLowering(); }
143
144ShadowStackGCLowering::ShadowStackGCLowering() : FunctionPass(ID) {}
145
146Constant *ShadowStackGCLoweringImpl::GetFrameMap(Function &F) {
147 // doInitialization creates the abstract type of this value.
148 Type *VoidPtr = PointerType::getUnqual(C&: F.getContext());
149
150 // Truncate the ShadowStackDescriptor if some metadata is null.
151 unsigned NumMeta = 0;
152 SmallVector<Constant *, 16> Metadata;
153 for (unsigned I = 0; I != Roots.size(); ++I) {
154 Constant *C = cast<Constant>(Val: Roots[I].first->getArgOperand(i: 1));
155 if (!C->isNullValue())
156 NumMeta = I + 1;
157 Metadata.push_back(Elt: C);
158 }
159 Metadata.resize(N: NumMeta);
160
161 Type *Int32Ty = Type::getInt32Ty(C&: F.getContext());
162
163 Constant *BaseElts[] = {
164 ConstantInt::get(Ty: Int32Ty, V: Roots.size(), IsSigned: false),
165 ConstantInt::get(Ty: Int32Ty, V: NumMeta, IsSigned: false),
166 };
167
168 Constant *DescriptorElts[] = {
169 ConstantStruct::get(T: FrameMapTy, V: BaseElts),
170 ConstantArray::get(T: ArrayType::get(ElementType: VoidPtr, NumElements: NumMeta), V: Metadata)};
171
172 Type *EltTys[] = {DescriptorElts[0]->getType(), DescriptorElts[1]->getType()};
173 StructType *STy = StructType::create(Elements: EltTys, Name: "gc_map." + utostr(X: NumMeta));
174
175 Constant *FrameMap = ConstantStruct::get(T: STy, V: DescriptorElts);
176
177 // FIXME: Is this actually dangerous as WritingAnLLVMPass.html claims? Seems
178 // that, short of multithreaded LLVM, it should be safe; all that is
179 // necessary is that a simple Module::iterator loop not be invalidated.
180 // Appending to the GlobalVariable list is safe in that sense.
181 //
182 // All of the output passes emit globals last. The ExecutionEngine
183 // explicitly supports adding globals to the module after
184 // initialization.
185 //
186 // Still, if it isn't deemed acceptable, then this transformation needs
187 // to be a ModulePass (which means it cannot be in the 'llc' pipeline
188 // (which uses a FunctionPassManager (which segfaults (not asserts) if
189 // provided a ModulePass))).
190 Constant *GV = new GlobalVariable(*F.getParent(), FrameMap->getType(), true,
191 GlobalVariable::InternalLinkage, FrameMap,
192 "__gc_" + F.getName());
193
194 Constant *GEPIndices[2] = {
195 ConstantInt::get(Ty: Type::getInt32Ty(C&: F.getContext()), V: 0),
196 ConstantInt::get(Ty: Type::getInt32Ty(C&: F.getContext()), V: 0)};
197 return ConstantExpr::getGetElementPtr(Ty: FrameMap->getType(), C: GV, IdxList: GEPIndices);
198}
199
200Type *ShadowStackGCLoweringImpl::GetConcreteStackEntryType(Function &F) {
201 // doInitialization creates the generic version of this type.
202 std::vector<Type *> EltTys;
203 EltTys.push_back(x: StackEntryTy);
204 for (const std::pair<CallInst *, AllocaInst *> &Root : Roots)
205 EltTys.push_back(x: Root.second->getAllocatedType());
206
207 return StructType::create(Elements: EltTys, Name: ("gc_stackentry." + F.getName()).str());
208}
209
210/// doInitialization - If this module uses the GC intrinsics, find them now. If
211/// not, exit fast.
212bool ShadowStackGCLoweringImpl::doInitialization(Module &M) {
213 bool Active = false;
214 for (Function &F : M) {
215 if (F.hasGC() && F.getGC() == "shadow-stack") {
216 Active = true;
217 break;
218 }
219 }
220 if (!Active)
221 return false;
222
223 // struct FrameMap {
224 // int32_t NumRoots; // Number of roots in stack frame.
225 // int32_t NumMeta; // Number of metadata descriptors. May be < NumRoots.
226 // void *Meta[]; // May be absent for roots without metadata.
227 // };
228 std::vector<Type *> EltTys;
229 // 32 bits is ok up to a 32GB stack frame. :)
230 EltTys.push_back(x: Type::getInt32Ty(C&: M.getContext()));
231 // Specifies length of variable length array.
232 EltTys.push_back(x: Type::getInt32Ty(C&: M.getContext()));
233 FrameMapTy = StructType::create(Elements: EltTys, Name: "gc_map");
234 PointerType *FrameMapPtrTy = PointerType::getUnqual(C&: M.getContext());
235
236 // struct StackEntry {
237 // ShadowStackEntry *Next; // Caller's stack entry.
238 // FrameMap *Map; // Pointer to constant FrameMap.
239 // void *Roots[]; // Stack roots (in-place array, so we pretend).
240 // };
241
242 PointerType *StackEntryPtrTy = PointerType::getUnqual(C&: M.getContext());
243
244 EltTys.clear();
245 EltTys.push_back(x: StackEntryPtrTy);
246 EltTys.push_back(x: FrameMapPtrTy);
247 StackEntryTy = StructType::create(Elements: EltTys, Name: "gc_stackentry");
248
249 // Get the root chain if it already exists.
250 Head = M.getGlobalVariable(Name: "llvm_gc_root_chain");
251 if (!Head) {
252 // If the root chain does not exist, insert a new one with linkonce
253 // linkage!
254 Head = new GlobalVariable(
255 M, StackEntryPtrTy, false, GlobalValue::LinkOnceAnyLinkage,
256 Constant::getNullValue(Ty: StackEntryPtrTy), "llvm_gc_root_chain");
257 } else if (Head->hasExternalLinkage() && Head->isDeclaration()) {
258 Head->setInitializer(Constant::getNullValue(Ty: StackEntryPtrTy));
259 Head->setLinkage(GlobalValue::LinkOnceAnyLinkage);
260 }
261
262 return true;
263}
264
265bool ShadowStackGCLoweringImpl::IsNullValue(Value *V) {
266 if (Constant *C = dyn_cast<Constant>(Val: V))
267 return C->isNullValue();
268 return false;
269}
270
271void ShadowStackGCLoweringImpl::CollectRoots(Function &F) {
272 // FIXME: Account for original alignment. Could fragment the root array.
273 // Approach 1: Null initialize empty slots at runtime. Yuck.
274 // Approach 2: Emit a map of the array instead of just a count.
275
276 assert(Roots.empty() && "Not cleaned up?");
277
278 SmallVector<std::pair<CallInst *, AllocaInst *>, 16> MetaRoots;
279
280 for (BasicBlock &BB : F)
281 for (Instruction &I : BB)
282 if (IntrinsicInst *CI = dyn_cast<IntrinsicInst>(Val: &I))
283 if (Function *F = CI->getCalledFunction())
284 if (F->getIntrinsicID() == Intrinsic::gcroot) {
285 std::pair<CallInst *, AllocaInst *> Pair = std::make_pair(
286 x&: CI,
287 y: cast<AllocaInst>(Val: CI->getArgOperand(i: 0)->stripPointerCasts()));
288 if (IsNullValue(V: CI->getArgOperand(i: 1)))
289 Roots.push_back(x: Pair);
290 else
291 MetaRoots.push_back(Elt: Pair);
292 }
293
294 // Number roots with metadata (usually empty) at the beginning, so that the
295 // FrameMap::Meta array can be elided.
296 Roots.insert(position: Roots.begin(), first: MetaRoots.begin(), last: MetaRoots.end());
297}
298
299GetElementPtrInst *
300ShadowStackGCLoweringImpl::CreateGEP(LLVMContext &Context, IRBuilder<> &B,
301 Type *Ty, Value *BasePtr, int Idx,
302 int Idx2, const char *Name) {
303 Value *Indices[] = {ConstantInt::get(Ty: Type::getInt32Ty(C&: Context), V: 0),
304 ConstantInt::get(Ty: Type::getInt32Ty(C&: Context), V: Idx),
305 ConstantInt::get(Ty: Type::getInt32Ty(C&: Context), V: Idx2)};
306 Value *Val = B.CreateGEP(Ty, Ptr: BasePtr, IdxList: Indices, Name);
307
308 assert(isa<GetElementPtrInst>(Val) && "Unexpected folded constant");
309
310 return dyn_cast<GetElementPtrInst>(Val);
311}
312
313GetElementPtrInst *ShadowStackGCLoweringImpl::CreateGEP(LLVMContext &Context,
314 IRBuilder<> &B,
315 Type *Ty,
316 Value *BasePtr, int Idx,
317 const char *Name) {
318 Value *Indices[] = {ConstantInt::get(Ty: Type::getInt32Ty(C&: Context), V: 0),
319 ConstantInt::get(Ty: Type::getInt32Ty(C&: Context), V: Idx)};
320 Value *Val = B.CreateGEP(Ty, Ptr: BasePtr, IdxList: Indices, Name);
321
322 assert(isa<GetElementPtrInst>(Val) && "Unexpected folded constant");
323
324 return dyn_cast<GetElementPtrInst>(Val);
325}
326
327/// runOnFunction - Insert code to maintain the shadow stack.
328bool ShadowStackGCLoweringImpl::runOnFunction(Function &F,
329 DomTreeUpdater *DTU) {
330 // Quick exit for functions that do not use the shadow stack GC.
331 if (!F.hasGC() || F.getGC() != "shadow-stack")
332 return false;
333
334 LLVMContext &Context = F.getContext();
335
336 // Find calls to llvm.gcroot.
337 CollectRoots(F);
338
339 // If there are no roots in this function, then there is no need to add a
340 // stack map entry for it.
341 if (Roots.empty())
342 return false;
343
344 // Build the constant map and figure the type of the shadow stack entry.
345 Value *FrameMap = GetFrameMap(F);
346 Type *ConcreteStackEntryTy = GetConcreteStackEntryType(F);
347
348 // Build the shadow stack entry at the very start of the function.
349 BasicBlock::iterator IP = F.getEntryBlock().begin();
350 IRBuilder<> AtEntry(IP->getParent(), IP);
351
352 Instruction *StackEntry =
353 AtEntry.CreateAlloca(Ty: ConcreteStackEntryTy, ArraySize: nullptr, Name: "gc_frame");
354
355 AtEntry.SetInsertPointPastAllocas(&F);
356 IP = AtEntry.GetInsertPoint();
357
358 // Initialize the map pointer and load the current head of the shadow stack.
359 Instruction *CurrentHead =
360 AtEntry.CreateLoad(Ty: AtEntry.getPtrTy(), Ptr: Head, Name: "gc_currhead");
361 Instruction *EntryMapPtr = CreateGEP(Context, B&: AtEntry, Ty: ConcreteStackEntryTy,
362 BasePtr: StackEntry, Idx: 0, Idx2: 1, Name: "gc_frame.map");
363 AtEntry.CreateStore(Val: FrameMap, Ptr: EntryMapPtr);
364
365 // After all the allocas...
366 for (unsigned I = 0, E = Roots.size(); I != E; ++I) {
367 // For each root, find the corresponding slot in the aggregate...
368 Value *SlotPtr = CreateGEP(Context, B&: AtEntry, Ty: ConcreteStackEntryTy,
369 BasePtr: StackEntry, Idx: 1 + I, Name: "gc_root");
370
371 // And use it in lieu of the alloca.
372 AllocaInst *OriginalAlloca = Roots[I].second;
373 SlotPtr->takeName(V: OriginalAlloca);
374 OriginalAlloca->replaceAllUsesWith(V: SlotPtr);
375 }
376
377 // Move past the original stores inserted by GCStrategy::InitRoots. This isn't
378 // really necessary (the collector would never see the intermediate state at
379 // runtime), but it's nicer not to push the half-initialized entry onto the
380 // shadow stack.
381 while (isa<StoreInst>(Val: IP))
382 ++IP;
383 AtEntry.SetInsertPoint(TheBB: IP->getParent(), IP);
384
385 // Push the entry onto the shadow stack.
386 Instruction *EntryNextPtr = CreateGEP(Context, B&: AtEntry, Ty: ConcreteStackEntryTy,
387 BasePtr: StackEntry, Idx: 0, Idx2: 0, Name: "gc_frame.next");
388 Instruction *NewHeadVal = CreateGEP(Context, B&: AtEntry, Ty: ConcreteStackEntryTy,
389 BasePtr: StackEntry, Idx: 0, Name: "gc_newhead");
390 AtEntry.CreateStore(Val: CurrentHead, Ptr: EntryNextPtr);
391 AtEntry.CreateStore(Val: NewHeadVal, Ptr: Head);
392
393 // For each instruction that escapes...
394 EscapeEnumerator EE(F, "gc_cleanup", /*HandleExceptions=*/true, DTU);
395 while (IRBuilder<> *AtExit = EE.Next()) {
396 // Pop the entry from the shadow stack. Don't reuse CurrentHead from
397 // AtEntry, since that would make the value live for the entire function.
398 Instruction *EntryNextPtr2 =
399 CreateGEP(Context, B&: *AtExit, Ty: ConcreteStackEntryTy, BasePtr: StackEntry, Idx: 0, Idx2: 0,
400 Name: "gc_frame.next");
401 Value *SavedHead =
402 AtExit->CreateLoad(Ty: AtExit->getPtrTy(), Ptr: EntryNextPtr2, Name: "gc_savedhead");
403 AtExit->CreateStore(Val: SavedHead, Ptr: Head);
404 }
405
406 // Delete the original allocas (which are no longer used) and the intrinsic
407 // calls (which are no longer valid). Doing this last avoids invalidating
408 // iterators.
409 for (std::pair<CallInst *, AllocaInst *> &Root : Roots) {
410 Root.first->eraseFromParent();
411 Root.second->eraseFromParent();
412 }
413
414 Roots.clear();
415 return true;
416}
417