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 initializeShadowStackGCLoweringPass(Registry&: *PassRegistry::getPassRegistry());
146}
147
148Constant *ShadowStackGCLoweringImpl::GetFrameMap(Function &F) {
149 // doInitialization creates the abstract type of this value.
150 Type *VoidPtr = PointerType::getUnqual(C&: F.getContext());
151
152 // Truncate the ShadowStackDescriptor if some metadata is null.
153 unsigned NumMeta = 0;
154 SmallVector<Constant *, 16> Metadata;
155 for (unsigned I = 0; I != Roots.size(); ++I) {
156 Constant *C = cast<Constant>(Val: Roots[I].first->getArgOperand(i: 1));
157 if (!C->isNullValue())
158 NumMeta = I + 1;
159 Metadata.push_back(Elt: C);
160 }
161 Metadata.resize(N: NumMeta);
162
163 Type *Int32Ty = Type::getInt32Ty(C&: F.getContext());
164
165 Constant *BaseElts[] = {
166 ConstantInt::get(Ty: Int32Ty, V: Roots.size(), IsSigned: false),
167 ConstantInt::get(Ty: Int32Ty, V: NumMeta, IsSigned: false),
168 };
169
170 Constant *DescriptorElts[] = {
171 ConstantStruct::get(T: FrameMapTy, V: BaseElts),
172 ConstantArray::get(T: ArrayType::get(ElementType: VoidPtr, NumElements: NumMeta), V: Metadata)};
173
174 Type *EltTys[] = {DescriptorElts[0]->getType(), DescriptorElts[1]->getType()};
175 StructType *STy = StructType::create(Elements: EltTys, Name: "gc_map." + utostr(X: NumMeta));
176
177 Constant *FrameMap = ConstantStruct::get(T: STy, V: DescriptorElts);
178
179 // FIXME: Is this actually dangerous as WritingAnLLVMPass.html claims? Seems
180 // that, short of multithreaded LLVM, it should be safe; all that is
181 // necessary is that a simple Module::iterator loop not be invalidated.
182 // Appending to the GlobalVariable list is safe in that sense.
183 //
184 // All of the output passes emit globals last. The ExecutionEngine
185 // explicitly supports adding globals to the module after
186 // initialization.
187 //
188 // Still, if it isn't deemed acceptable, then this transformation needs
189 // to be a ModulePass (which means it cannot be in the 'llc' pipeline
190 // (which uses a FunctionPassManager (which segfaults (not asserts) if
191 // provided a ModulePass))).
192 Constant *GV = new GlobalVariable(*F.getParent(), FrameMap->getType(), true,
193 GlobalVariable::InternalLinkage, FrameMap,
194 "__gc_" + F.getName());
195
196 Constant *GEPIndices[2] = {
197 ConstantInt::get(Ty: Type::getInt32Ty(C&: F.getContext()), V: 0),
198 ConstantInt::get(Ty: Type::getInt32Ty(C&: F.getContext()), V: 0)};
199 return ConstantExpr::getGetElementPtr(Ty: FrameMap->getType(), C: GV, IdxList: GEPIndices);
200}
201
202Type *ShadowStackGCLoweringImpl::GetConcreteStackEntryType(Function &F) {
203 // doInitialization creates the generic version of this type.
204 std::vector<Type *> EltTys;
205 EltTys.push_back(x: StackEntryTy);
206 for (const std::pair<CallInst *, AllocaInst *> &Root : Roots)
207 EltTys.push_back(x: Root.second->getAllocatedType());
208
209 return StructType::create(Elements: EltTys, Name: ("gc_stackentry." + F.getName()).str());
210}
211
212/// doInitialization - If this module uses the GC intrinsics, find them now. If
213/// not, exit fast.
214bool ShadowStackGCLoweringImpl::doInitialization(Module &M) {
215 bool Active = false;
216 for (Function &F : M) {
217 if (F.hasGC() && F.getGC() == "shadow-stack") {
218 Active = true;
219 break;
220 }
221 }
222 if (!Active)
223 return false;
224
225 // struct FrameMap {
226 // int32_t NumRoots; // Number of roots in stack frame.
227 // int32_t NumMeta; // Number of metadata descriptors. May be < NumRoots.
228 // void *Meta[]; // May be absent for roots without metadata.
229 // };
230 std::vector<Type *> EltTys;
231 // 32 bits is ok up to a 32GB stack frame. :)
232 EltTys.push_back(x: Type::getInt32Ty(C&: M.getContext()));
233 // Specifies length of variable length array.
234 EltTys.push_back(x: Type::getInt32Ty(C&: M.getContext()));
235 FrameMapTy = StructType::create(Elements: EltTys, Name: "gc_map");
236 PointerType *FrameMapPtrTy = PointerType::getUnqual(C&: M.getContext());
237
238 // struct StackEntry {
239 // ShadowStackEntry *Next; // Caller's stack entry.
240 // FrameMap *Map; // Pointer to constant FrameMap.
241 // void *Roots[]; // Stack roots (in-place array, so we pretend).
242 // };
243
244 PointerType *StackEntryPtrTy = PointerType::getUnqual(C&: M.getContext());
245
246 EltTys.clear();
247 EltTys.push_back(x: StackEntryPtrTy);
248 EltTys.push_back(x: FrameMapPtrTy);
249 StackEntryTy = StructType::create(Elements: EltTys, Name: "gc_stackentry");
250
251 // Get the root chain if it already exists.
252 Head = M.getGlobalVariable(Name: "llvm_gc_root_chain");
253 if (!Head) {
254 // If the root chain does not exist, insert a new one with linkonce
255 // linkage!
256 Head = new GlobalVariable(
257 M, StackEntryPtrTy, false, GlobalValue::LinkOnceAnyLinkage,
258 Constant::getNullValue(Ty: StackEntryPtrTy), "llvm_gc_root_chain");
259 } else if (Head->hasExternalLinkage() && Head->isDeclaration()) {
260 Head->setInitializer(Constant::getNullValue(Ty: StackEntryPtrTy));
261 Head->setLinkage(GlobalValue::LinkOnceAnyLinkage);
262 }
263
264 return true;
265}
266
267bool ShadowStackGCLoweringImpl::IsNullValue(Value *V) {
268 if (Constant *C = dyn_cast<Constant>(Val: V))
269 return C->isNullValue();
270 return false;
271}
272
273void ShadowStackGCLoweringImpl::CollectRoots(Function &F) {
274 // FIXME: Account for original alignment. Could fragment the root array.
275 // Approach 1: Null initialize empty slots at runtime. Yuck.
276 // Approach 2: Emit a map of the array instead of just a count.
277
278 assert(Roots.empty() && "Not cleaned up?");
279
280 SmallVector<std::pair<CallInst *, AllocaInst *>, 16> MetaRoots;
281
282 for (BasicBlock &BB : F)
283 for (Instruction &I : BB)
284 if (IntrinsicInst *CI = dyn_cast<IntrinsicInst>(Val: &I))
285 if (Function *F = CI->getCalledFunction())
286 if (F->getIntrinsicID() == Intrinsic::gcroot) {
287 std::pair<CallInst *, AllocaInst *> Pair = std::make_pair(
288 x&: CI,
289 y: cast<AllocaInst>(Val: CI->getArgOperand(i: 0)->stripPointerCasts()));
290 if (IsNullValue(V: CI->getArgOperand(i: 1)))
291 Roots.push_back(x: Pair);
292 else
293 MetaRoots.push_back(Elt: Pair);
294 }
295
296 // Number roots with metadata (usually empty) at the beginning, so that the
297 // FrameMap::Meta array can be elided.
298 Roots.insert(position: Roots.begin(), first: MetaRoots.begin(), last: MetaRoots.end());
299}
300
301GetElementPtrInst *
302ShadowStackGCLoweringImpl::CreateGEP(LLVMContext &Context, IRBuilder<> &B,
303 Type *Ty, Value *BasePtr, int Idx,
304 int Idx2, const char *Name) {
305 Value *Indices[] = {ConstantInt::get(Ty: Type::getInt32Ty(C&: Context), V: 0),
306 ConstantInt::get(Ty: Type::getInt32Ty(C&: Context), V: Idx),
307 ConstantInt::get(Ty: Type::getInt32Ty(C&: Context), V: Idx2)};
308 Value *Val = B.CreateGEP(Ty, Ptr: BasePtr, IdxList: Indices, Name);
309
310 assert(isa<GetElementPtrInst>(Val) && "Unexpected folded constant");
311
312 return dyn_cast<GetElementPtrInst>(Val);
313}
314
315GetElementPtrInst *ShadowStackGCLoweringImpl::CreateGEP(LLVMContext &Context,
316 IRBuilder<> &B,
317 Type *Ty,
318 Value *BasePtr, int Idx,
319 const char *Name) {
320 Value *Indices[] = {ConstantInt::get(Ty: Type::getInt32Ty(C&: Context), V: 0),
321 ConstantInt::get(Ty: Type::getInt32Ty(C&: Context), V: Idx)};
322 Value *Val = B.CreateGEP(Ty, Ptr: BasePtr, IdxList: Indices, Name);
323
324 assert(isa<GetElementPtrInst>(Val) && "Unexpected folded constant");
325
326 return dyn_cast<GetElementPtrInst>(Val);
327}
328
329/// runOnFunction - Insert code to maintain the shadow stack.
330bool ShadowStackGCLoweringImpl::runOnFunction(Function &F,
331 DomTreeUpdater *DTU) {
332 // Quick exit for functions that do not use the shadow stack GC.
333 if (!F.hasGC() || F.getGC() != "shadow-stack")
334 return false;
335
336 LLVMContext &Context = F.getContext();
337
338 // Find calls to llvm.gcroot.
339 CollectRoots(F);
340
341 // If there are no roots in this function, then there is no need to add a
342 // stack map entry for it.
343 if (Roots.empty())
344 return false;
345
346 // Build the constant map and figure the type of the shadow stack entry.
347 Value *FrameMap = GetFrameMap(F);
348 Type *ConcreteStackEntryTy = GetConcreteStackEntryType(F);
349
350 // Build the shadow stack entry at the very start of the function.
351 BasicBlock::iterator IP = F.getEntryBlock().begin();
352 IRBuilder<> AtEntry(IP->getParent(), IP);
353
354 Instruction *StackEntry =
355 AtEntry.CreateAlloca(Ty: ConcreteStackEntryTy, ArraySize: nullptr, Name: "gc_frame");
356
357 AtEntry.SetInsertPointPastAllocas(&F);
358 IP = AtEntry.GetInsertPoint();
359
360 // Initialize the map pointer and load the current head of the shadow stack.
361 Instruction *CurrentHead =
362 AtEntry.CreateLoad(Ty: AtEntry.getPtrTy(), Ptr: Head, Name: "gc_currhead");
363 Instruction *EntryMapPtr = CreateGEP(Context, B&: AtEntry, Ty: ConcreteStackEntryTy,
364 BasePtr: StackEntry, Idx: 0, Idx2: 1, Name: "gc_frame.map");
365 AtEntry.CreateStore(Val: FrameMap, Ptr: EntryMapPtr);
366
367 // After all the allocas...
368 for (unsigned I = 0, E = Roots.size(); I != E; ++I) {
369 // For each root, find the corresponding slot in the aggregate...
370 Value *SlotPtr = CreateGEP(Context, B&: AtEntry, Ty: ConcreteStackEntryTy,
371 BasePtr: StackEntry, Idx: 1 + I, Name: "gc_root");
372
373 // And use it in lieu of the alloca.
374 AllocaInst *OriginalAlloca = Roots[I].second;
375 SlotPtr->takeName(V: OriginalAlloca);
376 OriginalAlloca->replaceAllUsesWith(V: SlotPtr);
377 }
378
379 // Move past the original stores inserted by GCStrategy::InitRoots. This isn't
380 // really necessary (the collector would never see the intermediate state at
381 // runtime), but it's nicer not to push the half-initialized entry onto the
382 // shadow stack.
383 while (isa<StoreInst>(Val: IP))
384 ++IP;
385 AtEntry.SetInsertPoint(TheBB: IP->getParent(), IP);
386
387 // Push the entry onto the shadow stack.
388 Instruction *EntryNextPtr = CreateGEP(Context, B&: AtEntry, Ty: ConcreteStackEntryTy,
389 BasePtr: StackEntry, Idx: 0, Idx2: 0, Name: "gc_frame.next");
390 Instruction *NewHeadVal = CreateGEP(Context, B&: AtEntry, Ty: ConcreteStackEntryTy,
391 BasePtr: StackEntry, Idx: 0, Name: "gc_newhead");
392 AtEntry.CreateStore(Val: CurrentHead, Ptr: EntryNextPtr);
393 AtEntry.CreateStore(Val: NewHeadVal, Ptr: Head);
394
395 // For each instruction that escapes...
396 EscapeEnumerator EE(F, "gc_cleanup", /*HandleExceptions=*/true, DTU);
397 while (IRBuilder<> *AtExit = EE.Next()) {
398 // Pop the entry from the shadow stack. Don't reuse CurrentHead from
399 // AtEntry, since that would make the value live for the entire function.
400 Instruction *EntryNextPtr2 =
401 CreateGEP(Context, B&: *AtExit, Ty: ConcreteStackEntryTy, BasePtr: StackEntry, Idx: 0, Idx2: 0,
402 Name: "gc_frame.next");
403 Value *SavedHead =
404 AtExit->CreateLoad(Ty: AtExit->getPtrTy(), Ptr: EntryNextPtr2, Name: "gc_savedhead");
405 AtExit->CreateStore(Val: SavedHead, Ptr: Head);
406 }
407
408 // Delete the original allocas (which are no longer used) and the intrinsic
409 // calls (which are no longer valid). Doing this last avoids invalidating
410 // iterators.
411 for (std::pair<CallInst *, AllocaInst *> &Root : Roots) {
412 Root.first->eraseFromParent();
413 Root.second->eraseFromParent();
414 }
415
416 Roots.clear();
417 return true;
418}
419