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