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