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/DataLayout.h"
28#include "llvm/IR/DerivedTypes.h"
29#include "llvm/IR/Dominators.h"
30#include "llvm/IR/Function.h"
31#include "llvm/IR/GlobalValue.h"
32#include "llvm/IR/GlobalVariable.h"
33#include "llvm/IR/IRBuilder.h"
34#include "llvm/IR/Instructions.h"
35#include "llvm/IR/IntrinsicInst.h"
36#include "llvm/IR/Intrinsics.h"
37#include "llvm/IR/Module.h"
38#include "llvm/IR/Type.h"
39#include "llvm/IR/Value.h"
40#include "llvm/InitializePasses.h"
41#include "llvm/Pass.h"
42#include "llvm/Support/Alignment.h"
43#include "llvm/Support/Casting.h"
44#include "llvm/Transforms/Utils/EscapeEnumerator.h"
45#include <cassert>
46#include <optional>
47#include <utility>
48#include <vector>
49
50using namespace llvm;
51
52#define DEBUG_TYPE "shadow-stack-gc-lowering"
53
54namespace {
55
56class ShadowStackGCLoweringImpl {
57 /// RootChain - This is the global linked-list that contains the chain of GC
58 /// roots.
59 GlobalVariable *Head = nullptr;
60
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
67 /// RootOffsets - Byte offsets and sizes of each root within the frame.
68 /// Each element is a pair of (offset, size).
69 std::vector<std::pair<uint64_t, uint64_t>> RootOffsets;
70
71public:
72 ShadowStackGCLoweringImpl() = default;
73
74 bool doInitialization(Module &M);
75 bool runOnFunction(Function &F, DomTreeUpdater *DTU);
76
77private:
78 bool IsNullValue(Value *V);
79 Constant *GetFrameMap(Function &F, uint64_t FrameSizeInPtrs);
80 std::pair<uint64_t, Align> ComputeFrameLayout(Function &F);
81 void CollectRoots(Function &F);
82};
83
84class ShadowStackGCLowering : public FunctionPass {
85 ShadowStackGCLoweringImpl Impl;
86
87public:
88 static char ID;
89
90 ShadowStackGCLowering();
91
92 bool doInitialization(Module &M) override { return Impl.doInitialization(M); }
93 void getAnalysisUsage(AnalysisUsage &AU) const override {
94 AU.addPreserved<DominatorTreeWrapperPass>();
95 }
96 bool runOnFunction(Function &F) override {
97 std::optional<DomTreeUpdater> DTU;
98 if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
99 DTU.emplace(args&: DTWP->getDomTree(), args: DomTreeUpdater::UpdateStrategy::Lazy);
100 return Impl.runOnFunction(F, DTU: DTU ? &*DTU : nullptr);
101 }
102};
103
104} // end anonymous namespace
105
106PreservedAnalyses ShadowStackGCLoweringPass::run(Module &M,
107 ModuleAnalysisManager &MAM) {
108 auto &Map = MAM.getResult<CollectorMetadataAnalysis>(IR&: M);
109 if (!Map.contains(GCName: "shadow-stack"))
110 return PreservedAnalyses::all();
111
112 ShadowStackGCLoweringImpl Impl;
113 bool Changed = Impl.doInitialization(M);
114 for (auto &F : M) {
115 auto &FAM =
116 MAM.getResult<FunctionAnalysisManagerModuleProxy>(IR&: M).getManager();
117 auto *DT = FAM.getCachedResult<DominatorTreeAnalysis>(IR&: F);
118 DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
119 Changed |= Impl.runOnFunction(F, DTU: DT ? &DTU : nullptr);
120 }
121
122 if (!Changed)
123 return PreservedAnalyses::all();
124 PreservedAnalyses PA;
125 PA.preserve<DominatorTreeAnalysis>();
126 return PA;
127}
128
129char ShadowStackGCLowering::ID = 0;
130char &llvm::ShadowStackGCLoweringID = ShadowStackGCLowering::ID;
131
132INITIALIZE_PASS_BEGIN(ShadowStackGCLowering, DEBUG_TYPE,
133 "Shadow Stack GC Lowering", false, false)
134INITIALIZE_PASS_DEPENDENCY(GCModuleInfo)
135INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
136INITIALIZE_PASS_END(ShadowStackGCLowering, DEBUG_TYPE,
137 "Shadow Stack GC Lowering", false, false)
138
139FunctionPass *llvm::createShadowStackGCLoweringPass() { return new ShadowStackGCLowering(); }
140
141ShadowStackGCLowering::ShadowStackGCLowering() : FunctionPass(ID) {}
142
143Constant *ShadowStackGCLoweringImpl::GetFrameMap(Function &F,
144 uint64_t FrameSizeInPtrs) {
145 // doInitialization creates the abstract type of this value.
146 Type *VoidPtr = PointerType::getUnqual(C&: F.getContext());
147
148 // Truncate the ShadowStackDescriptor if some metadata is null.
149 unsigned NumMeta = 0;
150 SmallVector<Constant *, 16> Metadata;
151 for (unsigned I = 0; I != Roots.size(); ++I) {
152 Constant *C = cast<Constant>(Val: Roots[I].first->getArgOperand(i: 1));
153 if (!C->isNullValue())
154 NumMeta = I + 1;
155 Metadata.push_back(Elt: C);
156 }
157 Metadata.resize(N: NumMeta);
158
159 Type *Int32Ty = Type::getInt32Ty(C&: F.getContext());
160
161 Constant *BaseElts[] = {
162 ConstantInt::get(Ty: Int32Ty, V: FrameSizeInPtrs, IsSigned: false),
163 ConstantInt::get(Ty: Int32Ty, V: NumMeta, IsSigned: false),
164 };
165
166 Constant *DescriptorElts[] = {
167 ConstantStruct::get(T: FrameMapTy, V: BaseElts),
168 ConstantArray::get(T: ArrayType::get(ElementType: VoidPtr, NumElements: NumMeta), V: Metadata)};
169
170 Type *EltTys[] = {DescriptorElts[0]->getType(), DescriptorElts[1]->getType()};
171 StructType *STy = StructType::create(Elements: EltTys, Name: "gc_map." + utostr(X: NumMeta));
172
173 Constant *FrameMap = ConstantStruct::get(T: STy, V: DescriptorElts);
174
175 // FIXME: Is this actually dangerous as WritingAnLLVMPass.html claims? Seems
176 // that, short of multithreaded LLVM, it should be safe; all that is
177 // necessary is that a simple Module::iterator loop not be invalidated.
178 // Appending to the GlobalVariable list is safe in that sense.
179 //
180 // All of the output passes emit globals last. The ExecutionEngine
181 // explicitly supports adding globals to the module after
182 // initialization.
183 //
184 // Still, if it isn't deemed acceptable, then this transformation needs
185 // to be a ModulePass (which means it cannot be in the 'llc' pipeline
186 // (which uses a FunctionPassManager (which segfaults (not asserts) if
187 // provided a ModulePass))).
188 return new GlobalVariable(*F.getParent(), FrameMap->getType(), true,
189 GlobalVariable::InternalLinkage, FrameMap,
190 "__gc_" + F.getName());
191}
192
193std::pair<uint64_t, Align>
194ShadowStackGCLoweringImpl::ComputeFrameLayout(Function &F) {
195 // Compute the layout of the shadow stack frame using byte offsets.
196 // Layout: [Next ptr | Map ptr | Root 0 | Root 1 | ... | Root N]
197
198 const DataLayout &DL = F.getParent()->getDataLayout();
199 uint64_t PtrSize = DL.getPointerSize(AS: 0);
200 Align PtrAlign = DL.getPointerABIAlignment(AS: 0);
201
202 RootOffsets.clear();
203 Align MaxAlign = PtrAlign;
204
205 // Offset 0: Next pointer
206 // Offset PtrSize: Map pointer
207 uint64_t Offset = 2 * PtrSize;
208
209 // Compute offsets and sizes for each root
210 for (const std::pair<CallInst *, AllocaInst *> &Root : Roots) {
211 AllocaInst *AI = Root.second;
212 std::optional<TypeSize> RootSize = AI->getAllocationSize(DL);
213 if (!RootSize || !RootSize->isFixed())
214 reportFatalUsageError(
215 reason: "Intrinsic::gcroot requires a fixed size stack object");
216 uint64_t Size = RootSize->getFixedValue();
217 Align RootAlign = AI->getAlign();
218 MaxAlign = std::max(a: MaxAlign, b: RootAlign);
219
220 // Align the offset for this root
221 uint64_t AlignedOffset = alignTo(Size: Offset, A: RootAlign);
222
223 // Store both offset and size as a pair
224 RootOffsets.push_back(x: {AlignedOffset, Size});
225 Offset = AlignedOffset + Size;
226 }
227
228 // Final frame size, aligned to maximum alignment
229 uint64_t FrameSize = alignTo(Size: Offset, A: MaxAlign);
230 return {FrameSize, MaxAlign};
231}
232
233/// doInitialization - If this module uses the GC intrinsics, find them now. If
234/// not, exit fast.
235bool ShadowStackGCLoweringImpl::doInitialization(Module &M) {
236 bool Active = false;
237 for (Function &F : M) {
238 if (F.hasGC() && F.getGC() == "shadow-stack") {
239 Active = true;
240 break;
241 }
242 }
243 if (!Active)
244 return false;
245
246 // struct FrameMap {
247 // int32_t NumRoots; // Number of roots in stack frame.
248 // int32_t NumMeta; // Number of metadata descriptors. May be < NumRoots.
249 // void *Meta[]; // May be absent for roots without metadata.
250 // };
251 std::vector<Type *> EltTys;
252 // 32 bits is ok up to a 32GB stack frame. :)
253 EltTys.push_back(x: Type::getInt32Ty(C&: M.getContext()));
254 // Specifies length of variable length array.
255 EltTys.push_back(x: Type::getInt32Ty(C&: M.getContext()));
256 FrameMapTy = StructType::create(Elements: EltTys, Name: "gc_map");
257
258 // The shadow stack linked list uses opaque pointers.
259 // Each frame is a byte array with: [Next ptr | Map ptr | Roots...]
260 PointerType *StackEntryPtrTy = PointerType::getUnqual(C&: M.getContext());
261
262 // Get the root chain if it already exists.
263 Head = M.getGlobalVariable(Name: "llvm_gc_root_chain");
264 if (!Head) {
265 // If the root chain does not exist, insert a new one with linkonce
266 // linkage!
267 Head = new GlobalVariable(
268 M, StackEntryPtrTy, false, GlobalValue::LinkOnceAnyLinkage,
269 Constant::getNullValue(Ty: StackEntryPtrTy), "llvm_gc_root_chain");
270 } else if (Head->hasExternalLinkage() && Head->isDeclaration()) {
271 Head->setInitializer(Constant::getNullValue(Ty: StackEntryPtrTy));
272 Head->setLinkage(GlobalValue::LinkOnceAnyLinkage);
273 }
274
275 return true;
276}
277
278bool ShadowStackGCLoweringImpl::IsNullValue(Value *V) {
279 if (Constant *C = dyn_cast<Constant>(Val: V))
280 return C->isNullValue();
281 return false;
282}
283
284void ShadowStackGCLoweringImpl::CollectRoots(Function &F) {
285 assert(Roots.empty() && "Not cleaned up?");
286
287 SmallVector<std::pair<CallInst *, AllocaInst *>, 16> MetaRoots;
288
289 for (BasicBlock &BB : F)
290 for (Instruction &I : BB)
291 if (IntrinsicInst *CI = dyn_cast<IntrinsicInst>(Val: &I))
292 if (Function *F = CI->getCalledFunction())
293 if (F->getIntrinsicID() == Intrinsic::gcroot) {
294 std::pair<CallInst *, AllocaInst *> Pair = std::make_pair(
295 x&: CI,
296 y: cast<AllocaInst>(Val: CI->getArgOperand(i: 0)->stripPointerCasts()));
297 if (IsNullValue(V: CI->getArgOperand(i: 1)))
298 Roots.push_back(x: Pair);
299 else
300 MetaRoots.push_back(Elt: Pair);
301 }
302
303 // Number roots with metadata (usually empty) at the beginning, so that the
304 // FrameMap::Meta array can be elided.
305 Roots.insert(position: Roots.begin(), first: MetaRoots.begin(), last: MetaRoots.end());
306}
307
308/// runOnFunction - Insert code to maintain the shadow stack.
309bool ShadowStackGCLoweringImpl::runOnFunction(Function &F,
310 DomTreeUpdater *DTU) {
311 // Quick exit for functions that do not use the shadow stack GC.
312 if (!F.hasGC() || F.getGC() != "shadow-stack")
313 return false;
314
315 LLVMContext &Context = F.getContext();
316 const DataLayout &DL = F.getParent()->getDataLayout();
317
318 // Find calls to llvm.gcroot.
319 CollectRoots(F);
320
321 // If there are no roots in this function, then there is no need to add a
322 // stack map entry for it.
323 if (Roots.empty())
324 return false;
325
326 // Compute frame layout using byte offsets first.
327 auto [FrameSize, FrameAlign] = ComputeFrameLayout(F);
328
329 // Build the constant map with frame size in pointer-sized units.
330 uint64_t PtrSize = DL.getPointerSize();
331 Value *FrameMap = GetFrameMap(F, FrameSizeInPtrs: FrameSize / PtrSize - 2);
332
333 // Build the shadow stack entry at the very start of the function.
334 BasicBlock::iterator IP = F.getEntryBlock().begin();
335 IRBuilder<> AtEntry(IP->getParent(), IP);
336 Type *Int8Ty = Type::getInt8Ty(C&: Context);
337 AllocaInst *StackEntry = AtEntry.CreateAlloca(
338 Ty: ArrayType::get(ElementType: Int8Ty, NumElements: FrameSize), ArraySize: nullptr, Name: "gc_frame");
339 StackEntry->setAlignment(FrameAlign);
340
341 AtEntry.SetInsertPointPastAllocas(&F);
342 IP = AtEntry.GetInsertPoint();
343
344 // Initialize the map pointer and load the current head of the shadow stack.
345 Instruction *CurrentHead =
346 AtEntry.CreateLoad(Ty: AtEntry.getPtrTy(), Ptr: Head, Name: "gc_currhead");
347
348 // Map pointer is at offset PtrSize (after the Next pointer)
349 Value *EntryMapPtr = AtEntry.CreatePtrAdd(
350 Ptr: StackEntry, Offset: AtEntry.getInt64(C: PtrSize), Name: "gc_frame.map");
351 AtEntry.CreateStore(Val: FrameMap, Ptr: EntryMapPtr);
352
353 // Zero out any padding between roots to ensure deterministic frame contents.
354 // This includes the region after the map pointer up to the first root.
355 uint64_t LastEnd = 2 * PtrSize; // End of Map pointer field
356 assert(RootOffsets.size() == Roots.size());
357 for (unsigned I = 0, E = Roots.size(); I != E; ++I) {
358 auto [RootOffset, RootSize] = RootOffsets[I];
359
360 // Zero any padding before this root
361 if (RootOffset > LastEnd) {
362 Value *PaddingPtr =
363 AtEntry.CreatePtrAdd(Ptr: StackEntry, Offset: AtEntry.getInt64(C: LastEnd));
364 AtEntry.CreateMemSet(Ptr: PaddingPtr, Val: AtEntry.getInt8(C: 0), Size: RootOffset - LastEnd,
365 Align: Align(1));
366 }
367
368 // For each root, compute pointer using precomputed offset
369 Value *SlotPtr = AtEntry.CreatePtrAdd(
370 Ptr: StackEntry, Offset: AtEntry.getInt64(C: RootOffset), Name: "gc_root");
371
372 // And use it in lieu of the alloca.
373 AllocaInst *OriginalAlloca = Roots[I].second;
374 SlotPtr->takeName(V: OriginalAlloca);
375 OriginalAlloca->replaceAllUsesWith(V: SlotPtr);
376
377 LastEnd = RootOffset + RootSize;
378 }
379
380 // Zero any padding at the end of the frame
381 if (FrameSize > LastEnd) {
382 Value *PaddingPtr =
383 AtEntry.CreatePtrAdd(Ptr: StackEntry, Offset: AtEntry.getInt64(C: LastEnd));
384 AtEntry.CreateMemSet(Ptr: PaddingPtr, Val: AtEntry.getInt8(C: 0), Size: FrameSize - LastEnd,
385 Align: Align(1));
386 }
387
388 // Move past the original stores inserted by GCStrategy::InitRoots. This isn't
389 // really necessary (the collector would never see the intermediate state at
390 // runtime), but it's nicer not to push the half-initialized entry onto the
391 // shadow stack.
392 while (isa<StoreInst>(Val: IP))
393 ++IP;
394 AtEntry.SetInsertPoint(TheBB: IP->getParent(), IP);
395
396 // Push the entry onto the shadow stack.
397 // Next pointer is at offset 0, so it's just the frame pointer
398 AtEntry.CreateStore(Val: CurrentHead, Ptr: StackEntry);
399 // The new head value is also the frame pointer (the linked list links to
400 // frame base)
401 AtEntry.CreateStore(Val: StackEntry, Ptr: Head);
402
403 // For each instruction that escapes...
404 EscapeEnumerator EE(F, "gc_cleanup", /*HandleExceptions=*/true, DTU);
405 while (IRBuilder<> *AtExit = EE.Next()) {
406 // Pop the entry from the shadow stack. Don't reuse CurrentHead from
407 // AtEntry, since that would make the value live for the entire function.
408 // Next pointer is at offset 0, so load from the frame base
409 Value *SavedHead =
410 AtExit->CreateLoad(Ty: AtExit->getPtrTy(), Ptr: StackEntry, Name: "gc_savedhead");
411 AtExit->CreateStore(Val: SavedHead, Ptr: Head);
412 }
413
414 // Delete the original allocas (which are no longer used) and the intrinsic
415 // calls (which are no longer valid). Doing this last avoids invalidating
416 // iterators.
417 for (std::pair<CallInst *, AllocaInst *> &Root : Roots) {
418 Root.first->eraseFromParent();
419 Root.second->eraseFromParent();
420 }
421
422 Roots.clear();
423 RootOffsets.clear();
424 return true;
425}
426