Skip to content

Commit cf02d92

Browse files
kbeylskbeyls
committed
[GS/StackClash] Also recognize non-power-of-2 constant values in registers.
1 parent 5c7dca7 commit cf02d92

File tree

3 files changed

+120
-52
lines changed

3 files changed

+120
-52
lines changed

‎bolt/lib/Passes/StackClashAnalysis.cpp

Lines changed: 85 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -57,26 +57,15 @@ namespace {
5757
// Step 1 - let's start by implementing detecting an issue on a bare minimal
5858
// example. Therefore, implement (2), (3), (4) and (a).
5959

60-
/// Returns true if the register R is present in the Map M.
61-
bool addToMaxMap(SmallDenseMap<MCPhysReg, uint64_t, 1> &M, MCPhysReg R,
62-
const uint64_t MaxValue) {
63-
auto MIt = M.find(R);
64-
if (MIt == M.end()) {
65-
MIt->second = MaxValue;
66-
return false;
67-
} else {
68-
MIt->second = std::max(MIt->second, MaxValue);
69-
return true;
70-
}
71-
}
72-
7360
template <typename T, auto MergeValLambda> class LatticeT {
7461
private:
7562
enum LValType { _Bottom, Value, _Top } LValType;
7663
T V;
7764
LatticeT(enum LValType ValType, T Val) : LValType(ValType), V(Val) {}
7865
static LatticeT _TopV; //(_Top, T());
7966
static LatticeT _BottomV; //(_Bottom, T());
67+
static T _DefaultVal;
68+
8069
public:
8170
LatticeT() : LatticeT(_Bottom, T()) {}
8271
LatticeT(T V) : LatticeT(Value, V) {}
@@ -143,22 +132,22 @@ template <typename T, auto MergeValLambda> class LatticeT {
143132
V = f(V, V2);
144133
return *this;
145134
}
135+
const T &getValOrDefault() const {
136+
if (hasVal())
137+
return getVal();
138+
else
139+
return _DefaultVal;
140+
}
146141
};
147142

148-
template <typename T, auto M> LatticeT<T, M> LatticeT<T, M>::_TopV(_Top, T());
143+
149144
template <typename T, auto M>
150-
LatticeT<T, M> LatticeT<T, M>::_BottomV(_Bottom, T());
145+
raw_ostream &operator<<(raw_ostream &OS, const LatticeT<T, M> &V);
151146

147+
template <typename T, auto M> LatticeT<T, M> LatticeT<T, M>::_TopV(_Top, T());
152148
template <typename T, auto M>
153-
raw_ostream &operator<<(raw_ostream &OS, const LatticeT<T, M> &V) {
154-
if (V == V.Top())
155-
OS << "(T)";
156-
else if (V == V.Bottom())
157-
OS << "(B)";
158-
else
159-
OS << V.getVal();
160-
return OS;
161-
}
149+
LatticeT<T, M> LatticeT<T, M>::_BottomV(_Bottom, T());
150+
template <typename T, auto M> T LatticeT<T, M>::_DefaultVal = T();
162151

163152
bool MaxOffsetMergeVal(int64_t &v1, const int64_t &v2) { return v1 == v2; }
164153
using MaxOffsetT = LatticeT<int64_t, MaxOffsetMergeVal>;
@@ -209,6 +198,55 @@ raw_ostream &operator<<(raw_ostream &OS, const Reg2MaxOffsetValT &M) {
209198

210199
using Reg2MaxOffsetT = LatticeT<Reg2MaxOffsetValT, Reg2MaxOffsetMergeVal>;
211200

201+
using Reg2MaxValT = SmallDenseMap<MCPhysReg, uint64_t, 1>;
202+
bool RegMaxValuesValMerge(Reg2MaxValT &v1, const Reg2MaxValT &v2) {
203+
SmallVector<MCPhysReg, 1> RegMaxValuesToRemove;
204+
for (auto Reg2MaxValue : v1) {
205+
const MCPhysReg R(Reg2MaxValue.first);
206+
auto v2Reg2MaxValue = v2.find(R);
207+
if (v2Reg2MaxValue == v2.end())
208+
RegMaxValuesToRemove.push_back(R);
209+
else
210+
Reg2MaxValue.second =
211+
std::max(Reg2MaxValue.second, v2Reg2MaxValue->second);
212+
// FIXME: this should be a "confluence" - similar
213+
// to MaxOffsetT? To avoid near infinite loops?
214+
}
215+
for (MCPhysReg R : RegMaxValuesToRemove)
216+
v1.erase(R);
217+
return true;
218+
}
219+
raw_ostream &operator<<(raw_ostream &OS, const Reg2MaxValT &M) {
220+
for (auto Reg2Value : M) {
221+
print_reg(OS, Reg2Value.first, nullptr);
222+
OS << ":" << Reg2Value.second << ",";
223+
}
224+
return OS;
225+
}
226+
using RegMaxValuesT = LatticeT<Reg2MaxValT, RegMaxValuesValMerge>;
227+
228+
void addToMaxMap(RegMaxValuesT &M, MCPhysReg R, const uint64_t Value) {
229+
if (!M.hasVal())
230+
return;
231+
auto MIt = M->find(R);
232+
if (MIt == M->end())
233+
MIt->second = Value;
234+
else
235+
MIt->second = std::max(MIt->second, Value);
236+
}
237+
238+
template <typename T, auto M>
239+
raw_ostream &operator<<(raw_ostream &OS, const LatticeT<T, M> &V) {
240+
if (V == V.Top())
241+
OS << "(T)";
242+
else if (V == V.Bottom())
243+
OS << "(B)";
244+
else
245+
OS << V.getVal();
246+
return OS;
247+
}
248+
249+
212250
struct State {
213251
// Store the maximum possible offset to which the stack extends
214252
// beyond the furthest probe seen.
@@ -220,7 +258,7 @@ struct State {
220258
/// range [0, MaxValue-1].
221259
// FIXME: also make this std::optional!!!
222260
// FIXME: same for RegConstValues.
223-
SmallDenseMap<MCPhysReg, uint64_t, 1> RegMaxValues;
261+
RegMaxValuesT RegMaxValues;
224262
/// Reg2MaxOffset contains the registers that contain the value
225263
/// of SP at some point during the running function, where it's
226264
/// guaranteed that at the time the SP value was stored in the register,
@@ -263,20 +301,7 @@ struct State {
263301
for (MCPhysReg R : RegConstValuesToRemove)
264302
RegConstValues.erase(R);
265303

266-
SmallVector<MCPhysReg, 1> RegMaxValuesToRemove;
267-
for (auto Reg2MaxValue : RegMaxValues) {
268-
const MCPhysReg R(Reg2MaxValue.first);
269-
auto SInReg2MaxValue = StateIn.RegMaxValues.find(R);
270-
if (SInReg2MaxValue == StateIn.RegMaxValues.end())
271-
RegMaxValuesToRemove.push_back(R);
272-
else
273-
Reg2MaxValue.second =
274-
std::max(Reg2MaxValue.second, SInReg2MaxValue->second);
275-
// FIXME: this should be a "confluence" - similar
276-
// to MaxOffsetT? To avoid near infinite loops?
277-
}
278-
for (MCPhysReg R : RegMaxValuesToRemove)
279-
RegMaxValues.erase(R);
304+
RegMaxValues &= StateIn.RegMaxValues;
280305

281306
if (!SPFixedOffsetFromOrig || !StateIn.SPFixedOffsetFromOrig)
282307
SPFixedOffsetFromOrig.reset();
@@ -318,13 +343,18 @@ raw_ostream &print_state(raw_ostream &OS, const State &S,
318343
OS << "), RegConstValues(";
319344
PrintRegMap(OS, S.RegConstValues, BC);
320345
OS << "), RegMaxValues(";
321-
PrintRegMap(OS, S.RegMaxValues, BC);
346+
if (S.RegMaxValues.hasVal()) {
347+
OS << "(";
348+
PrintRegMap(OS, *S.RegMaxValues, BC);
349+
OS << ")";
350+
} else
351+
OS << S.RegMaxValues;
322352
OS << "),";
323353
OS << "SPFixedOffsetFromOrig:" << S.SPFixedOffsetFromOrig << ",";
324354
OS << "Reg2MaxOffset:";
325355
if (S.Reg2MaxOffset.hasVal()) {
326356
OS << "(";
327-
PrintRegMap(OS, S.Reg2MaxOffset.getVal(), BC);
357+
PrintRegMap(OS, *S.Reg2MaxOffset, BC);
328358
OS << ")";
329359
} else
330360
OS << S.Reg2MaxOffset;
@@ -363,7 +393,7 @@ bool checkNonConstSPOffsetChange(const BinaryContext &BC, BinaryFunction &BF,
363393
if (Next)
364394
Next->LastStackGrowingInsts.insert(MCInstInBBReference::get(&Point, BF));
365395
if (auto OC = BC.MIB->getOffsetChange(Point, Cur.RegConstValues,
366-
Cur.RegMaxValues);
396+
Cur.RegMaxValues.getValOrDefault());
367397
OC && OC.ToReg == SP) {
368398
if (OC.FromReg == SP) {
369399
IsNonConstantSPOffsetChange = false;
@@ -464,8 +494,10 @@ class StackClashDFAnalysis
464494

465495
State getStartingStateAtBB(const BinaryBasicBlock &BB) {
466496
State Next;
467-
if (BB.isEntryPoint())
497+
if (BB.isEntryPoint()) {
468498
Next.Reg2MaxOffset = Reg2MaxOffsetValT();
499+
Next.RegMaxValues = Reg2MaxValT();
500+
}
469501
return Next;
470502
}
471503

@@ -527,12 +559,13 @@ class StackClashDFAnalysis
527559
<< "; MaxValueMask: " << MaxValueMask << "\n";
528560
});
529561
const uint64_t MaxValueInReg = MaxValueMask;
530-
auto MaxValueForRegI = Next.RegMaxValues.find(MaxValueReg);
531-
if (MaxValueForRegI == Next.RegMaxValues.end())
532-
Next.RegMaxValues[MaxValueReg] = MaxValueInReg;
533-
else {
534-
MaxValueForRegI->second =
535-
std::min(MaxValueForRegI->second, MaxValueInReg);
562+
if (Next.RegMaxValues.hasVal()) {
563+
if (auto MaxValueForRegI = Next.RegMaxValues->find(MaxValueReg);
564+
MaxValueForRegI == Next.RegMaxValues->end())
565+
(*Next.RegMaxValues)[MaxValueReg] = MaxValueInReg;
566+
else
567+
MaxValueForRegI->second =
568+
std::min(MaxValueForRegI->second, MaxValueInReg);
536569
}
537570
}
538571
// FIXME properly handle register aliases below. E.g. a call
@@ -545,8 +578,8 @@ class StackClashDFAnalysis
545578
#endif
546579
for (const MCOperand &Operand : BC.MIB->defOperands(Point)) {
547580
assert(Operand.isReg());
548-
if (Operand.getReg() != MaxValueReg)
549-
Next.RegMaxValues.erase(Operand.getReg());
581+
if (Next.RegMaxValues.hasVal() && Operand.getReg() != MaxValueReg)
582+
Next.RegMaxValues->erase(Operand.getReg());
550583
if (Operand.getReg() != ConstValueReg)
551584
Next.RegConstValues.erase(Operand.getReg());
552585
}
@@ -572,7 +605,7 @@ class StackClashDFAnalysis
572605

573606
MCPhysReg FixedOffsetRegJustSet = BC.MIB->getNoRegister();
574607
if (auto OC = BC.MIB->getOffsetChange(Point, Cur.RegConstValues,
575-
Cur.RegMaxValues))
608+
Cur.RegMaxValues.getValOrDefault()))
576609
if (Next.Reg2MaxOffset.hasVal() && OC.OffsetChange) {
577610
int64_t Offset = *OC.OffsetChange;
578611
if (OC.FromReg == SP) {

‎bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -347,7 +347,14 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
347347
}
348348
auto RegMaxValueI = RegMaxValues.find(OffsetReg);
349349
if (RegMaxValueI != RegMaxValues.end()) {
350+
uint64_t UValue = RegMaxValueI->second;
350351
int64_t Value = RegMaxValueI->second;
352+
// if Value is so large that it could result
353+
// in wrap-arounds, do not consider this an
354+
// offset change. Heuristically, we assume
355+
// any value larger than 2^30 can cause wrap-arounds
356+
if (UValue > (1L<<30))
357+
return Res;
351358
if (IsSub)
352359
Value = -Value;
353360
Res.FromReg = Inst.getOperand(1).getReg();
@@ -424,8 +431,10 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
424431
case AArch64::ANDXri:
425432
uint64_t MaskValue = AArch64_AM::decodeLogicalImmediate(
426433
Inst.getOperand(2).getImm(), Is32Bit ? 32 : 64);
434+
#if 0
427435
if (!isPowerOf2_64(MaskValue + 1) && MaskValue != 0)
428436
return false;
437+
#endif
429438
ToReg = Inst.getOperand(0).getReg();
430439
Mask = MaskValue;
431440
return true;

‎bolt/test/gadget-scanner/gs-stackclash.s

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,32 @@ f_variable_large_stack_protected:
105105
ldp x29, x30, [sp], 16
106106
ret
107107
.size f_variable_large_stack_protected, .-f_variable_large_stack_protected
108+
109+
// CHECK-NOT: GS-STACKCLASH:
110+
111+
.global f_var_stack_protected2
112+
.type f_var_stack_protected2 , %function
113+
f_var_stack_protected2:
114+
// From "__BOLT_FDE_FUNCat7f530" in /usr/lib64/libasound.so.2
115+
mov x1, sp
116+
and x0, x0, #0x70
117+
cmp sp, x1
118+
b.eq .Lf_var_stack_protected2_2
119+
.Lf_var_stack_protected2_1:
120+
sub sp, sp, #0x10, lsl #12
121+
str xzr, [sp, #0x400]
122+
cmp sp, x1
123+
b.ne .Lf_var_stack_protected2_1
124+
.Lf_var_stack_protected2_2:
125+
sub sp, sp, x0
126+
ret
127+
.size f_var_stack_protected2, .-f_var_stack_protected2
128+
129+
130+
131+
// Second example: register containing SP-offset value is
132+
// spill/filled at a location of x29+fixed offset.
133+
108134
// CHECK-NOT: GS-STACKCLASH:
109135

110136
.global f_verify_detect_fp_corruption

0 commit comments

Comments
 (0)