Skip to content

Commit f9a1d15

Browse files
paperchalicepaperchalice
authored
[CodeGen] Port StackProtector to new pass manager (#75334)
The original `StackProtector` is both transform and analysis pass, break it into two passes now. `getAnalysis<StackProtector>()` could be now replaced by `FAM.getResult<SSPLayoutAnalysis>(F)` in new pass system.
1 parent 4f68ee3 commit f9a1d15

File tree

6 files changed

+176
-80
lines changed

6 files changed

+176
-80
lines changed

‎llvm/include/llvm/CodeGen/CodeGenPassBuilder.h

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
#include "llvm/CodeGen/SelectOptimize.h"
4141
#include "llvm/CodeGen/ShadowStackGCLowering.h"
4242
#include "llvm/CodeGen/SjLjEHPrepare.h"
43+
#include "llvm/CodeGen/StackProtector.h"
4344
#include "llvm/CodeGen/UnreachableBlockElim.h"
4445
#include "llvm/CodeGen/WasmEHPrepare.h"
4546
#include "llvm/CodeGen/WinEHPrepare.h"
@@ -749,7 +750,7 @@ void CodeGenPassBuilder<Derived>::addISelPrepare(AddIRPass &addPass) const {
749750
// Add both the safe stack and the stack protection passes: each of them will
750751
// only protect functions that have corresponding attributes.
751752
addPass(SafeStackPass(&TM));
752-
addPass(StackProtectorPass());
753+
addPass(StackProtectorPass(&TM));
753754

754755
if (Opt.PrintISelInput)
755756
addPass(PrintFunctionPass(dbgs(),

‎llvm/include/llvm/CodeGen/MachinePassRegistry.def

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ MODULE_PASS("shadow-stack-gc-lowering", ShadowStackGCLoweringPass, ())
3434
#endif
3535
FUNCTION_ANALYSIS("gc-function", GCFunctionAnalysis, ())
3636
FUNCTION_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis, (PIC))
37+
FUNCTION_ANALYSIS("ssp-layout", SSPLayoutAnalysis, ())
3738
FUNCTION_ANALYSIS("targetir", TargetIRAnalysis,
3839
(std::move(TM.getTargetIRAnalysis())))
3940
#undef FUNCTION_ANALYSIS
@@ -65,6 +66,7 @@ FUNCTION_PASS("safe-stack", SafeStackPass, (TM))
6566
FUNCTION_PASS("scalarize-masked-mem-intrin", ScalarizeMaskedMemIntrinPass, ())
6667
FUNCTION_PASS("select-optimize", SelectOptimizePass, (TM))
6768
FUNCTION_PASS("sjlj-eh-prepare", SjLjEHPreparePass, (TM))
69+
FUNCTION_PASS("stack-protector", StackProtectorPass, (TM))
6870
FUNCTION_PASS("tlshoist", TLSVariableHoistPass, ())
6971
FUNCTION_PASS("unreachableblockelim", UnreachableBlockElimPass, ())
7072
FUNCTION_PASS("verify", VerifierPass, ())
@@ -134,7 +136,6 @@ MACHINE_FUNCTION_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis,
134136
#endif
135137
DUMMY_FUNCTION_PASS("atomic-expand", AtomicExpandPass, ())
136138
DUMMY_FUNCTION_PASS("codegenprepare", CodeGenPreparePass, ())
137-
DUMMY_FUNCTION_PASS("stack-protector", StackProtectorPass, ())
138139
#undef DUMMY_FUNCTION_PASS
139140

140141
#ifndef DUMMY_MACHINE_MODULE_PASS

‎llvm/include/llvm/CodeGen/StackProtector.h

Lines changed: 63 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
#include "llvm/Analysis/DomTreeUpdater.h"
2020
#include "llvm/CodeGen/MachineFrameInfo.h"
2121
#include "llvm/IR/Instructions.h"
22+
#include "llvm/IR/PassManager.h"
2223
#include "llvm/Pass.h"
2324
#include "llvm/TargetParser/Triple.h"
2425

@@ -30,25 +31,15 @@ class Module;
3031
class TargetLoweringBase;
3132
class TargetMachine;
3233

33-
class StackProtector : public FunctionPass {
34-
private:
34+
class SSPLayoutInfo {
35+
friend class StackProtectorPass;
36+
friend class SSPLayoutAnalysis;
37+
friend class StackProtector;
3538
static constexpr unsigned DefaultSSPBufferSize = 8;
3639

3740
/// A mapping of AllocaInsts to their required SSP layout.
38-
using SSPLayoutMap = DenseMap<const AllocaInst *,
39-
MachineFrameInfo::SSPLayoutKind>;
40-
41-
const TargetMachine *TM = nullptr;
42-
43-
/// TLI - Keep a pointer of a TargetLowering to consult for determining
44-
/// target type sizes.
45-
const TargetLoweringBase *TLI = nullptr;
46-
Triple Trip;
47-
48-
Function *F = nullptr;
49-
Module *M = nullptr;
50-
51-
std::optional<DomTreeUpdater> DTU;
41+
using SSPLayoutMap =
42+
DenseMap<const AllocaInst *, MachineFrameInfo::SSPLayoutKind>;
5243

5344
/// Layout - Mapping of allocations to the required SSPLayoutKind.
5445
/// StackProtector analysis will update this map when determining if an
@@ -59,23 +50,59 @@ class StackProtector : public FunctionPass {
5950
/// protection when -fstack-protection is used.
6051
unsigned SSPBufferSize = DefaultSSPBufferSize;
6152

53+
bool RequireStackProtector = false;
54+
6255
// A prologue is generated.
6356
bool HasPrologue = false;
6457

6558
// IR checking code is generated.
6659
bool HasIRCheck = false;
6760

68-
/// InsertStackProtectors - Insert code into the prologue and epilogue of
69-
/// the function.
70-
///
71-
/// - The prologue code loads and stores the stack guard onto the stack.
72-
/// - The epilogue checks the value stored in the prologue against the
73-
/// original value. It calls __stack_chk_fail if they differ.
74-
bool InsertStackProtectors();
61+
public:
62+
// Return true if StackProtector is supposed to be handled by SelectionDAG.
63+
bool shouldEmitSDCheck(const BasicBlock &BB) const;
64+
65+
void copyToMachineFrameInfo(MachineFrameInfo &MFI) const;
66+
};
67+
68+
class SSPLayoutAnalysis : public AnalysisInfoMixin<SSPLayoutAnalysis> {
69+
friend class AnalysisInfoMixin<SSPLayoutAnalysis>;
70+
using SSPLayoutMap = SSPLayoutInfo::SSPLayoutMap;
71+
72+
static AnalysisKey Key;
73+
74+
public:
75+
using Result = SSPLayoutInfo;
76+
77+
Result run(Function &F, FunctionAnalysisManager &FAM);
7578

76-
/// CreateFailBB - Create a basic block to jump to when the stack protector
77-
/// check fails.
78-
BasicBlock *CreateFailBB();
79+
/// Check whether or not \p F needs a stack protector based upon the stack
80+
/// protector level.
81+
static bool requiresStackProtector(Function *F,
82+
SSPLayoutMap *Layout = nullptr);
83+
};
84+
85+
class StackProtectorPass : public PassInfoMixin<StackProtectorPass> {
86+
const TargetMachine *TM;
87+
88+
public:
89+
explicit StackProtectorPass(const TargetMachine *TM) : TM(TM) {}
90+
PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM);
91+
};
92+
93+
class StackProtector : public FunctionPass {
94+
private:
95+
/// A mapping of AllocaInsts to their required SSP layout.
96+
using SSPLayoutMap = SSPLayoutInfo::SSPLayoutMap;
97+
98+
const TargetMachine *TM = nullptr;
99+
100+
Function *F = nullptr;
101+
Module *M = nullptr;
102+
103+
std::optional<DomTreeUpdater> DTU;
104+
105+
SSPLayoutInfo LayoutInfo;
79106

80107
public:
81108
static char ID; // Pass identification, replacement for typeid.
@@ -85,16 +112,22 @@ class StackProtector : public FunctionPass {
85112
void getAnalysisUsage(AnalysisUsage &AU) const override;
86113

87114
// Return true if StackProtector is supposed to be handled by SelectionDAG.
88-
bool shouldEmitSDCheck(const BasicBlock &BB) const;
115+
bool shouldEmitSDCheck(const BasicBlock &BB) const {
116+
return LayoutInfo.shouldEmitSDCheck(BB);
117+
}
89118

90119
bool runOnFunction(Function &Fn) override;
91120

92-
void copyToMachineFrameInfo(MachineFrameInfo &MFI) const;
121+
void copyToMachineFrameInfo(MachineFrameInfo &MFI) const {
122+
LayoutInfo.copyToMachineFrameInfo(MFI);
123+
}
93124

94125
/// Check whether or not \p F needs a stack protector based upon the stack
95126
/// protector level.
96-
static bool requiresStackProtector(Function *F, SSPLayoutMap *Layout = nullptr);
97-
127+
static bool requiresStackProtector(Function *F,
128+
SSPLayoutMap *Layout = nullptr) {
129+
return SSPLayoutAnalysis::requiresStackProtector(F, Layout);
130+
}
98131
};
99132

100133
} // end namespace llvm

‎llvm/lib/CodeGen/StackProtector.cpp

Lines changed: 106 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,90 @@ static cl::opt<bool> EnableSelectionDAGSP("enable-selectiondag-sp",
6464
static cl::opt<bool> DisableCheckNoReturn("disable-check-noreturn-call",
6565
cl::init(false), cl::Hidden);
6666

67+
/// InsertStackProtectors - Insert code into the prologue and epilogue of the
68+
/// function.
69+
///
70+
/// - The prologue code loads and stores the stack guard onto the stack.
71+
/// - The epilogue checks the value stored in the prologue against the original
72+
/// value. It calls __stack_chk_fail if they differ.
73+
static bool InsertStackProtectors(const TargetMachine *TM, Function *F,
74+
DomTreeUpdater *DTU, bool &HasPrologue,
75+
bool &HasIRCheck);
76+
77+
/// CreateFailBB - Create a basic block to jump to when the stack protector
78+
/// check fails.
79+
static BasicBlock *CreateFailBB(Function *F, const Triple &Trip);
80+
81+
bool SSPLayoutInfo::shouldEmitSDCheck(const BasicBlock &BB) const {
82+
return HasPrologue && !HasIRCheck && isa<ReturnInst>(BB.getTerminator());
83+
}
84+
85+
void SSPLayoutInfo::copyToMachineFrameInfo(MachineFrameInfo &MFI) const {
86+
if (Layout.empty())
87+
return;
88+
89+
for (int I = 0, E = MFI.getObjectIndexEnd(); I != E; ++I) {
90+
if (MFI.isDeadObjectIndex(I))
91+
continue;
92+
93+
const AllocaInst *AI = MFI.getObjectAllocation(I);
94+
if (!AI)
95+
continue;
96+
97+
SSPLayoutMap::const_iterator LI = Layout.find(AI);
98+
if (LI == Layout.end())
99+
continue;
100+
101+
MFI.setObjectSSPLayout(I, LI->second);
102+
}
103+
}
104+
105+
SSPLayoutInfo SSPLayoutAnalysis::run(Function &F,
106+
FunctionAnalysisManager &FAM) {
107+
108+
SSPLayoutInfo Info;
109+
Info.RequireStackProtector =
110+
SSPLayoutAnalysis::requiresStackProtector(&F, &Info.Layout);
111+
Info.SSPBufferSize = F.getFnAttributeAsParsedInteger(
112+
"stack-protector-buffer-size", SSPLayoutInfo::DefaultSSPBufferSize);
113+
return Info;
114+
}
115+
116+
AnalysisKey SSPLayoutAnalysis::Key;
117+
118+
PreservedAnalyses StackProtectorPass::run(Function &F,
119+
FunctionAnalysisManager &FAM) {
120+
auto &Info = FAM.getResult<SSPLayoutAnalysis>(F);
121+
auto *DT = FAM.getCachedResult<DominatorTreeAnalysis>(F);
122+
DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
123+
124+
if (!Info.RequireStackProtector)
125+
return PreservedAnalyses::all();
126+
127+
// TODO(etienneb): Functions with funclets are not correctly supported now.
128+
// Do nothing if this is funclet-based personality.
129+
if (F.hasPersonalityFn()) {
130+
EHPersonality Personality = classifyEHPersonality(F.getPersonalityFn());
131+
if (isFuncletEHPersonality(Personality))
132+
return PreservedAnalyses::all();
133+
}
134+
135+
++NumFunProtected;
136+
bool Changed = InsertStackProtectors(TM, &F, DT ? &DTU : nullptr,
137+
Info.HasPrologue, Info.HasIRCheck);
138+
#ifdef EXPENSIVE_CHECKS
139+
assert((!DT || DT->verify(DominatorTree::VerificationLevel::Full)) &&
140+
"Failed to maintain validity of domtree!");
141+
#endif
142+
143+
if (!Changed)
144+
return PreservedAnalyses::all();
145+
PreservedAnalyses PA;
146+
PA.preserve<SSPLayoutAnalysis>();
147+
PA.preserve<DominatorTreeAnalysis>();
148+
return PA;
149+
}
150+
67151
char StackProtector::ID = 0;
68152

69153
StackProtector::StackProtector() : FunctionPass(ID) {
@@ -90,14 +174,12 @@ bool StackProtector::runOnFunction(Function &Fn) {
90174
if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
91175
DTU.emplace(DTWP->getDomTree(), DomTreeUpdater::UpdateStrategy::Lazy);
92176
TM = &getAnalysis<TargetPassConfig>().getTM<TargetMachine>();
93-
Trip = TM->getTargetTriple();
94-
TLI = TM->getSubtargetImpl(Fn)->getTargetLowering();
95-
HasPrologue = false;
96-
HasIRCheck = false;
97-
98-
SSPBufferSize = Fn.getFnAttributeAsParsedInteger(
99-
"stack-protector-buffer-size", DefaultSSPBufferSize);
100-
if (!requiresStackProtector(F, &Layout))
177+
LayoutInfo.HasPrologue = false;
178+
LayoutInfo.HasIRCheck = false;
179+
180+
LayoutInfo.SSPBufferSize = Fn.getFnAttributeAsParsedInteger(
181+
"stack-protector-buffer-size", SSPLayoutInfo::DefaultSSPBufferSize);
182+
if (!requiresStackProtector(F, &LayoutInfo.Layout))
101183
return false;
102184

103185
// TODO(etienneb): Functions with funclets are not correctly supported now.
@@ -109,7 +191,9 @@ bool StackProtector::runOnFunction(Function &Fn) {
109191
}
110192

111193
++NumFunProtected;
112-
bool Changed = InsertStackProtectors();
194+
bool Changed =
195+
InsertStackProtectors(TM, F, DTU ? &*DTU : nullptr,
196+
LayoutInfo.HasPrologue, LayoutInfo.HasIRCheck);
113197
#ifdef EXPENSIVE_CHECKS
114198
assert((!DTU ||
115199
DTU->getDomTree().verify(DominatorTree::VerificationLevel::Full)) &&
@@ -284,7 +368,8 @@ static const CallInst *findStackProtectorIntrinsic(Function &F) {
284368
/// functions with aggregates that contain any buffer regardless of type and
285369
/// size, and functions that contain stack-based variables that have had their
286370
/// address taken.
287-
bool StackProtector::requiresStackProtector(Function *F, SSPLayoutMap *Layout) {
371+
bool SSPLayoutAnalysis::requiresStackProtector(Function *F,
372+
SSPLayoutMap *Layout) {
288373
Module *M = F->getParent();
289374
bool Strong = false;
290375
bool NeedsProtector = false;
@@ -295,7 +380,7 @@ bool StackProtector::requiresStackProtector(Function *F, SSPLayoutMap *Layout) {
295380
SmallPtrSet<const PHINode *, 16> VisitedPHIs;
296381

297382
unsigned SSPBufferSize = F->getFnAttributeAsParsedInteger(
298-
"stack-protector-buffer-size", DefaultSSPBufferSize);
383+
"stack-protector-buffer-size", SSPLayoutInfo::DefaultSSPBufferSize);
299384

300385
if (F->hasFnAttribute(Attribute::SafeStack))
301386
return false;
@@ -460,13 +545,12 @@ static bool CreatePrologue(Function *F, Module *M, Instruction *CheckLoc,
460545
return SupportsSelectionDAGSP;
461546
}
462547

463-
/// InsertStackProtectors - Insert code into the prologue and epilogue of the
464-
/// function.
465-
///
466-
/// - The prologue code loads and stores the stack guard onto the stack.
467-
/// - The epilogue checks the value stored in the prologue against the original
468-
/// value. It calls __stack_chk_fail if they differ.
469-
bool StackProtector::InsertStackProtectors() {
548+
bool InsertStackProtectors(const TargetMachine *TM, Function *F,
549+
DomTreeUpdater *DTU, bool &HasPrologue,
550+
bool &HasIRCheck) {
551+
auto *M = F->getParent();
552+
auto *TLI = TM->getSubtargetImpl(*F)->getTargetLowering();
553+
470554
// If the target wants to XOR the frame pointer into the guard value, it's
471555
// impossible to emit the check in IR, so the target *must* support stack
472556
// protection in SDAG.
@@ -574,7 +658,7 @@ bool StackProtector::InsertStackProtectors() {
574658
// merge pass will merge together all of the various BB into one including
575659
// fail BB generated by the stack protector pseudo instruction.
576660
if (!FailBB)
577-
FailBB = CreateFailBB();
661+
FailBB = CreateFailBB(F, TM->getTargetTriple());
578662

579663
IRBuilder<> B(CheckLoc);
580664
Value *Guard = getStackGuard(TLI, M, B);
@@ -589,8 +673,7 @@ bool StackProtector::InsertStackProtectors() {
589673
SuccessProb.getNumerator());
590674

591675
SplitBlockAndInsertIfThen(Cmp, CheckLoc,
592-
/*Unreachable=*/false, Weights,
593-
DTU ? &*DTU : nullptr,
676+
/*Unreachable=*/false, Weights, DTU,
594677
/*LI=*/nullptr, /*ThenBlock=*/FailBB);
595678

596679
auto *BI = cast<BranchInst>(Cmp->getParent()->getTerminator());
@@ -608,9 +691,8 @@ bool StackProtector::InsertStackProtectors() {
608691
return HasPrologue;
609692
}
610693

611-
/// CreateFailBB - Create a basic block to jump to when the stack protector
612-
/// check fails.
613-
BasicBlock *StackProtector::CreateFailBB() {
694+
BasicBlock *CreateFailBB(Function *F, const Triple &Trip) {
695+
auto *M = F->getParent();
614696
LLVMContext &Context = F->getContext();
615697
BasicBlock *FailBB = BasicBlock::Create(Context, "CallStackCheckFailBlk", F);
616698
IRBuilder<> B(FailBB);
@@ -633,27 +715,3 @@ BasicBlock *StackProtector::CreateFailBB() {
633715
B.CreateUnreachable();
634716
return FailBB;
635717
}
636-
637-
bool StackProtector::shouldEmitSDCheck(const BasicBlock &BB) const {
638-
return HasPrologue && !HasIRCheck && isa<ReturnInst>(BB.getTerminator());
639-
}
640-
641-
void StackProtector::copyToMachineFrameInfo(MachineFrameInfo &MFI) const {
642-
if (Layout.empty())
643-
return;
644-
645-
for (int I = 0, E = MFI.getObjectIndexEnd(); I != E; ++I) {
646-
if (MFI.isDeadObjectIndex(I))
647-
continue;
648-
649-
const AllocaInst *AI = MFI.getObjectAllocation(I);
650-
if (!AI)
651-
continue;
652-
653-
SSPLayoutMap::const_iterator LI = Layout.find(AI);
654-
if (LI == Layout.end())
655-
continue;
656-
657-
MFI.setObjectSSPLayout(I, LI->second);
658-
}
659-
}

0 commit comments

Comments
 (0)