Skip to content

Commit 6217b66

Browse files
bgergely0bgergely0
committed
[BOLT] Support for OpNegateRAState - second half
 - create InsertNegateRAStatePass - this pass explores the CFG, and finds if a BB has signed or unsigned return address state - after exploration, it inserts OpNegateRAState at the end of BBs, where the following BB has a different RA State (or where RA State boundaries are)
1 parent 0353abb commit 6217b66

File tree

5 files changed

+300
-0
lines changed

5 files changed

+300
-0
lines changed

‎bolt/include/bolt/Core/BinaryBasicBlock.h

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,11 @@ class JumpTable;
3737

3838
class BinaryBasicBlock {
3939
public:
40+
enum class RAStateEnum : char {
41+
Unknown, /// Not discovered yet
42+
Signed,
43+
Unsigned,
44+
};
4045
/// Profile execution information for a given edge in CFG.
4146
///
4247
/// If MispredictedCount equals COUNT_INFERRED, then we have a profile
@@ -350,6 +355,17 @@ class BinaryBasicBlock {
350355
BranchInfo.end());
351356
}
352357

358+
RAStateEnum RAState{RAStateEnum::Unknown};
359+
void setRASigned() { RAState = RAStateEnum::Signed; }
360+
bool isRAStateUnknown() { return RAState == RAStateEnum::Unknown; }
361+
bool isRAStateSigned() { return RAState == RAStateEnum::Signed; }
362+
/// Unsigned should only overwrite Unknown state, and not Signed
363+
void setRAUnsigned() {
364+
if (RAState == RAStateEnum::Unknown) {
365+
RAState = RAStateEnum::Unsigned;
366+
}
367+
}
368+
353369
/// Get instruction at given index.
354370
MCInst &getInstructionAtIndex(unsigned Index) { return Instructions[Index]; }
355371

‎bolt/include/bolt/Passes/InsertNegateRAStatePass.h

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
#ifndef BOLT_PASSES_INSERT_NEGATE_RA_STATE_PASS
2+
#define BOLT_PASSES_INSERT_NEGATE_RA_STATE_PASS
3+
4+
#include "bolt/Passes/BinaryPasses.h"
5+
#include <stack>
6+
7+
namespace llvm {
8+
namespace bolt {
9+
10+
class InsertNegateRAState : public BinaryFunctionPass {
11+
public:
12+
explicit InsertNegateRAState() : BinaryFunctionPass(false) {}
13+
14+
const char *getName() const override { return "insert-negate-ra-state-pass"; }
15+
16+
/// Pass entry point
17+
Error runOnFunctions(BinaryContext &BC) override;
18+
void runOnFunction(BinaryFunction &BF);
19+
bool addNegateRAStateAfterPacOrAuth(BinaryFunction &BF);
20+
bool BBhasAUTH(BinaryContext &BC, BinaryBasicBlock *BB);
21+
bool BBhasSIGN(BinaryContext &BC, BinaryBasicBlock *BB);
22+
void explore_call_graph(BinaryContext &BC, BinaryBasicBlock *BB);
23+
void process_signed_BB(BinaryContext &BC, BinaryBasicBlock *BB,
24+
std::stack<BinaryBasicBlock *> *SignedStack,
25+
std::stack<BinaryBasicBlock *> *UnsignedStack);
26+
void process_unsigned_BB(BinaryContext &BC, BinaryBasicBlock *BB,
27+
std::stack<BinaryBasicBlock *> *SignedStack,
28+
std::stack<BinaryBasicBlock *> *UnsignedStack);
29+
};
30+
31+
} // namespace bolt
32+
} // namespace llvm
33+
#endif

‎bolt/lib/Passes/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ add_llvm_library(LLVMBOLTPasses
1717
IdenticalCodeFolding.cpp
1818
IndirectCallPromotion.cpp
1919
Inliner.cpp
20+
InsertNegateRAStatePass.cpp
2021
Instrumentation.cpp
2122
JTFootprintReduction.cpp
2223
LongJmp.cpp

‎bolt/lib/Passes/InsertNegateRAStatePass.cpp

Lines changed: 247 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,247 @@
1+
#include "bolt/Passes/InsertNegateRAStatePass.h"
2+
#include "bolt/Core/BinaryFunction.h"
3+
#include "bolt/Core/ParallelUtilities.h"
4+
#include "bolt/Utils/CommandLineOpts.h"
5+
#include <cstdlib>
6+
#include <fstream>
7+
#include <iterator>
8+
9+
using namespace llvm;
10+
11+
namespace llvm {
12+
namespace bolt {
13+
14+
void InsertNegateRAState::runOnFunction(BinaryFunction &BF) {
15+
BinaryContext &BC = BF.getBinaryContext();
16+
17+
if (BF.getState() == BinaryFunction::State::Empty) {
18+
return;
19+
}
20+
21+
if (BF.getState() != BinaryFunction::State::CFG &&
22+
BF.getState() != BinaryFunction::State::CFG_Finalized) {
23+
BC.errs() << "BOLT-WARNING: No CFG for " << BF.getPrintName()
24+
<< " in InsertNegateRAStatePass\n";
25+
return;
26+
}
27+
28+
if (BF.getState() == BinaryFunction::State::CFG_Finalized) {
29+
BC.errs() << "BOLT-WARNING: CFG finalized for " << BF.getPrintName()
30+
<< " in InsertNegateRAStatePass\n";
31+
return;
32+
}
33+
34+
if (BF.isIgnored())
35+
return;
36+
37+
if (!addNegateRAStateAfterPacOrAuth(BF)) {
38+
// none inserted, function doesn't need more work
39+
return;
40+
}
41+
42+
auto FirstBB = BF.begin();
43+
explore_call_graph(BC, &(*FirstBB));
44+
45+
// We have to do the walk again, starting from any undiscovered autiasp
46+
// instructions, because some autiasp might not be reachable because of
47+
// indirect branches but we know that autiasp block should have a Signed
48+
// state, so we can work out other Unkown states starting from these nodes.
49+
for (BinaryBasicBlock &BB : BF) {
50+
if (BBhasAUTH(BC, &BB) && BB.isRAStateUnknown()) {
51+
BB.setRASigned();
52+
explore_call_graph(BC, &BB);
53+
}
54+
}
55+
56+
// insert negateRAState-s where there is a State boundary:
57+
// that is: two consecutive BBs have different RA State
58+
BinaryFunction::iterator PrevBB;
59+
bool FirstIter = true;
60+
for (auto BB = BF.begin(); BB != BF.end(); ++BB) {
61+
if (!FirstIter) {
62+
if ((PrevBB->RAState == BinaryBasicBlock::RAStateEnum::Signed &&
63+
(*BB).RAState == BinaryBasicBlock::RAStateEnum::Unsigned &&
64+
!BBhasAUTH(BC, &(*PrevBB))) ||
65+
(PrevBB->RAState == BinaryBasicBlock::RAStateEnum::Signed &&
66+
(*BB).RAState == BinaryBasicBlock::RAStateEnum::Signed &&
67+
BBhasAUTH(BC, &(*PrevBB)))) {
68+
auto InstRevIter = PrevBB->getLastNonPseudo();
69+
MCInst LastNonPseudo = *InstRevIter;
70+
auto InstIter = InstRevIter.base();
71+
BF.addCFIInstruction(&(*PrevBB), InstIter,
72+
MCCFIInstruction::createNegateRAState(nullptr));
73+
}
74+
} else {
75+
FirstIter = false;
76+
}
77+
PrevBB = BB;
78+
}
79+
}
80+
81+
void InsertNegateRAState::explore_call_graph(BinaryContext &BC,
82+
BinaryBasicBlock *BB) {
83+
std::stack<BinaryBasicBlock *> SignedStack;
84+
std::stack<BinaryBasicBlock *> UnsignedStack;
85+
86+
// start according to the first BB
87+
if (BBhasSIGN(BC, BB)) {
88+
SignedStack.push(BB);
89+
process_signed_BB(BC, BB, &SignedStack, &UnsignedStack);
90+
} else {
91+
UnsignedStack.push(BB);
92+
process_unsigned_BB(BC, BB, &SignedStack, &UnsignedStack);
93+
}
94+
95+
while (!(SignedStack.empty() && UnsignedStack.empty())) {
96+
if (!SignedStack.empty()) {
97+
BB = SignedStack.top();
98+
SignedStack.pop();
99+
process_signed_BB(BC, BB, &SignedStack, &UnsignedStack);
100+
} else if (!UnsignedStack.empty()) {
101+
BB = UnsignedStack.top();
102+
UnsignedStack.pop();
103+
process_unsigned_BB(BC, BB, &SignedStack, &UnsignedStack);
104+
}
105+
}
106+
}
107+
void InsertNegateRAState::process_signed_BB(
108+
BinaryContext &BC, BinaryBasicBlock *BB,
109+
std::stack<BinaryBasicBlock *> *SignedStack,
110+
std::stack<BinaryBasicBlock *> *UnsignedStack) {
111+
112+
BB->setRASigned();
113+
114+
if (BBhasAUTH(BC, BB)) {
115+
// successors of block with autiasp are stored in the Unsigned Stack
116+
for (BinaryBasicBlock *Succ : BB->successors()) {
117+
if (Succ->getFunction() == BB->getFunction() &&
118+
Succ->isRAStateUnknown()) {
119+
UnsignedStack->push(Succ);
120+
}
121+
}
122+
} else {
123+
for (BinaryBasicBlock *Succ : BB->successors()) {
124+
if (Succ->getFunction() == BB->getFunction() &&
125+
!Succ->isRAStateSigned()) {
126+
SignedStack->push(Succ);
127+
}
128+
}
129+
}
130+
// process predecessors
131+
if (BBhasSIGN(BC, BB)) {
132+
for (BinaryBasicBlock *Pred : BB->predecessors()) {
133+
if (Pred->getFunction() == BB->getFunction() &&
134+
Pred->isRAStateUnknown()) {
135+
UnsignedStack->push(Pred);
136+
}
137+
}
138+
} else {
139+
for (BinaryBasicBlock *Pred : BB->predecessors()) {
140+
if (Pred->getFunction() == BB->getFunction() &&
141+
!Pred->isRAStateSigned()) {
142+
SignedStack->push(Pred);
143+
}
144+
}
145+
}
146+
}
147+
148+
void InsertNegateRAState::process_unsigned_BB(
149+
BinaryContext &BC, BinaryBasicBlock *BB,
150+
std::stack<BinaryBasicBlock *> *SignedStack,
151+
std::stack<BinaryBasicBlock *> *UnsignedStack) {
152+
153+
BB->setRAUnsigned();
154+
155+
if (BBhasSIGN(BC, BB)) {
156+
BB->setRASigned();
157+
// successors of block with paciasp are stored in the Signed Stack
158+
for (BinaryBasicBlock *Succ : BB->successors()) {
159+
if (Succ->getFunction() == BB->getFunction() &&
160+
!Succ->isRAStateSigned()) {
161+
SignedStack->push(Succ);
162+
}
163+
}
164+
} else {
165+
for (BinaryBasicBlock *Succ : BB->successors()) {
166+
if (Succ->getFunction() == BB->getFunction() &&
167+
Succ->isRAStateUnknown()) {
168+
UnsignedStack->push(Succ);
169+
}
170+
}
171+
}
172+
173+
// process predecessors
174+
if (BBhasAUTH(BC, BB)) {
175+
BB->setRASigned();
176+
for (BinaryBasicBlock *Pred : BB->predecessors()) {
177+
if (Pred->getFunction() == BB->getFunction() &&
178+
!Pred->isRAStateSigned()) {
179+
SignedStack->push(Pred);
180+
}
181+
}
182+
} else {
183+
for (BinaryBasicBlock *Pred : BB->predecessors()) {
184+
if (Pred->getFunction() == BB->getFunction() &&
185+
Pred->isRAStateUnknown()) {
186+
UnsignedStack->push(Pred);
187+
}
188+
}
189+
}
190+
}
191+
192+
bool InsertNegateRAState::BBhasAUTH(BinaryContext &BC, BinaryBasicBlock *BB) {
193+
for (auto Iter = BB->begin(); Iter != BB->end(); ++Iter) {
194+
MCInst Inst = *Iter;
195+
if (BC.MIB->isPAuth(Inst)) {
196+
return true;
197+
}
198+
}
199+
return false;
200+
}
201+
202+
bool InsertNegateRAState::BBhasSIGN(BinaryContext &BC, BinaryBasicBlock *BB) {
203+
for (auto Iter = BB->begin(); Iter != BB->end(); ++Iter) {
204+
MCInst Inst = *Iter;
205+
if (BC.MIB->isPSign(Inst)) {
206+
return true;
207+
}
208+
}
209+
return false;
210+
}
211+
212+
bool InsertNegateRAState::addNegateRAStateAfterPacOrAuth(BinaryFunction &BF) {
213+
BinaryContext &BC = BF.getBinaryContext();
214+
bool FoundAny = false;
215+
for (BinaryBasicBlock &BB : BF) {
216+
for (auto Iter = BB.begin(); Iter != BB.end(); ++Iter) {
217+
MCInst Inst = *Iter;
218+
if (BC.MIB->isPSign(Inst)) {
219+
Iter = BF.addCFIInstruction(
220+
&BB, Iter + 1, MCCFIInstruction::createNegateRAState(nullptr));
221+
FoundAny = true;
222+
}
223+
224+
if (BC.MIB->isPAuth(Inst)) {
225+
Iter = BF.addCFIInstruction(
226+
&BB, Iter + 1, MCCFIInstruction::createNegateRAState(nullptr));
227+
FoundAny = true;
228+
}
229+
}
230+
}
231+
return FoundAny;
232+
}
233+
234+
Error InsertNegateRAState::runOnFunctions(BinaryContext &BC) {
235+
ParallelUtilities::WorkFuncTy WorkFun = [&](BinaryFunction &BF) {
236+
runOnFunction(BF);
237+
};
238+
239+
ParallelUtilities::runOnEachFunction(
240+
BC, ParallelUtilities::SchedulingPolicy::SP_TRIVIAL, WorkFun, nullptr,
241+
"InsertNegateRAStatePass");
242+
243+
return Error::success();
244+
}
245+
246+
} // end namespace bolt
247+
} // end namespace llvm

‎bolt/lib/Rewrite/BinaryPassManager.cpp

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
#include "bolt/Passes/IdenticalCodeFolding.h"
2121
#include "bolt/Passes/IndirectCallPromotion.h"
2222
#include "bolt/Passes/Inliner.h"
23+
#include "bolt/Passes/InsertNegateRAStatePass.h"
2324
#include "bolt/Passes/Instrumentation.h"
2425
#include "bolt/Passes/JTFootprintReduction.h"
2526
#include "bolt/Passes/LongJmp.h"
@@ -499,6 +500,8 @@ Error BinaryFunctionPassManager::runAllPasses(BinaryContext &BC) {
499500
// targets. No extra instructions after this pass, otherwise we may have
500501
// relocations out of range and crash during linking.
501502
Manager.registerPass(std::make_unique<LongJmpPass>(PrintLongJmp));
503+
504+
Manager.registerPass(std::make_unique<InsertNegateRAState>());
502505
}
503506

504507
// This pass should always run last.*

0 commit comments

Comments
 (0)