[llvm] [StackProtector] Fix phi handling in HasAddressTaken() (PR #129248)
via llvm-commits
llvm-commits at lists.llvm.org
Fri Feb 28 06:11:40 PST 2025
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-backend-x86
Author: Nikita Popov (nikic)
<details>
<summary>Changes</summary>
Despite the name, the HasAddressTaken() heuristic identifies not only allocas that have their address taken, but also those that have accesses that cannot be proven to be in-bounds.
However, the current handling for phi nodes is incorrect. Phi nodes are only visited once, and will perform the analysis using whichever (remaining) allocation size is passed the first time the phi node is visited. If it is later visited with a smaller remaining size, which may lead to out of bounds accesses, it will not be detected.
Fix this by keeping track of the smallest seen remaining allocation size and redo the analysis if it is decreased. To avoid degenerate cases (including via loops), limit the number of allowed decreases to a small number.
---
Full diff: https://github.com/llvm/llvm-project/pull/129248.diff
2 Files Affected:
- (modified) llvm/lib/CodeGen/StackProtector.cpp (+26-4)
- (modified) llvm/test/CodeGen/X86/stack-protector-phi.ll (+26-4)
``````````diff
diff --git a/llvm/lib/CodeGen/StackProtector.cpp b/llvm/lib/CodeGen/StackProtector.cpp
index 232e84fb672e0..83f90b4ffd3db 100644
--- a/llvm/lib/CodeGen/StackProtector.cpp
+++ b/llvm/lib/CodeGen/StackProtector.cpp
@@ -251,10 +251,21 @@ static bool ContainsProtectableArray(Type *Ty, Module *M, unsigned SSPBufferSize
return NeedsProtector;
}
+/// Maximum remaining allocation size obsorved for a phi node, and how often
+/// the allocation size has already been decreased. We only allow a limited
+/// number of decreases.
+struct PhiInfo {
+ TypeSize AllocSize;
+ unsigned NumDecreased = 0;
+ static constexpr unsigned MaxNumDecreased = 3;
+ PhiInfo(TypeSize AllocSize) : AllocSize(AllocSize) {}
+};
+using PhiMap = SmallDenseMap<const PHINode *, PhiInfo, 16>;
+
/// Check whether a stack allocation has its address taken.
static bool HasAddressTaken(const Instruction *AI, TypeSize AllocSize,
Module *M,
- SmallPtrSet<const PHINode *, 16> &VisitedPHIs) {
+ PhiMap &VisitedPHIs) {
const DataLayout &DL = M->getDataLayout();
for (const User *U : AI->users()) {
const auto *I = cast<Instruction>(U);
@@ -325,9 +336,20 @@ static bool HasAddressTaken(const Instruction *AI, TypeSize AllocSize,
// Keep track of what PHI nodes we have already visited to ensure
// they are only visited once.
const auto *PN = cast<PHINode>(I);
- if (VisitedPHIs.insert(PN).second)
- if (HasAddressTaken(PN, AllocSize, M, VisitedPHIs))
+ auto [It, Inserted] = VisitedPHIs.try_emplace(PN, AllocSize);
+ if (!Inserted) {
+ if (TypeSize::isKnownGE(AllocSize, It->second.AllocSize))
+ break;
+
+ // Check again with smaller size.
+ if (It->second.NumDecreased == PhiInfo::MaxNumDecreased)
return true;
+
+ It->second.AllocSize = AllocSize;
+ ++It->second.NumDecreased;
+ }
+ if (HasAddressTaken(PN, AllocSize, M, VisitedPHIs))
+ return true;
break;
}
case Instruction::Load:
@@ -377,7 +399,7 @@ bool SSPLayoutAnalysis::requiresStackProtector(Function *F,
// The set of PHI nodes visited when determining if a variable's reference has
// been taken. This set is maintained to ensure we don't visit the same PHI
// node multiple times.
- SmallPtrSet<const PHINode *, 16> VisitedPHIs;
+ PhiMap VisitedPHIs;
unsigned SSPBufferSize = F->getFnAttributeAsParsedInteger(
"stack-protector-buffer-size", SSPLayoutInfo::DefaultSSPBufferSize);
diff --git a/llvm/test/CodeGen/X86/stack-protector-phi.ll b/llvm/test/CodeGen/X86/stack-protector-phi.ll
index bf0442dbf47a1..8041869ad8783 100644
--- a/llvm/test/CodeGen/X86/stack-protector-phi.ll
+++ b/llvm/test/CodeGen/X86/stack-protector-phi.ll
@@ -4,16 +4,29 @@
define void @test_phi_diff_size(i1 %c) sspstrong {
; CHECK-LABEL: test_phi_diff_size:
; CHECK: # %bb.0: # %entry
+; CHECK-NEXT: subq $24, %rsp
+; CHECK-NEXT: .cfi_def_cfa_offset 32
+; CHECK-NEXT: movq %fs:40, %rax
+; CHECK-NEXT: movq %rax, {{[0-9]+}}(%rsp)
; CHECK-NEXT: testb $1, %dil
; CHECK-NEXT: je .LBB0_1
; CHECK-NEXT: # %bb.2: # %if
-; CHECK-NEXT: leaq -{{[0-9]+}}(%rsp), %rax
-; CHECK-NEXT: movq $0, (%rax)
-; CHECK-NEXT: retq
+; CHECK-NEXT: leaq {{[0-9]+}}(%rsp), %rax
+; CHECK-NEXT: jmp .LBB0_3
; CHECK-NEXT: .LBB0_1:
-; CHECK-NEXT: leaq -{{[0-9]+}}(%rsp), %rax
+; CHECK-NEXT: leaq {{[0-9]+}}(%rsp), %rax
+; CHECK-NEXT: .LBB0_3: # %join
; CHECK-NEXT: movq $0, (%rax)
+; CHECK-NEXT: movq %fs:40, %rax
+; CHECK-NEXT: cmpq {{[0-9]+}}(%rsp), %rax
+; CHECK-NEXT: jne .LBB0_5
+; CHECK-NEXT: # %bb.4: # %SP_return
+; CHECK-NEXT: addq $24, %rsp
+; CHECK-NEXT: .cfi_def_cfa_offset 8
; CHECK-NEXT: retq
+; CHECK-NEXT: .LBB0_5: # %CallStackCheckFailBlk
+; CHECK-NEXT: .cfi_def_cfa_offset 32
+; CHECK-NEXT: callq __stack_chk_fail at PLT
entry:
%a = alloca i64
br i1 %c, label %if, label %join
@@ -38,6 +51,8 @@ define void @test_phi_loop(i1 %c) sspstrong {
; CHECK-NEXT: .cfi_def_cfa_register %rbp
; CHECK-NEXT: andq $-131072, %rsp # imm = 0xFFFE0000
; CHECK-NEXT: subq $262144, %rsp # imm = 0x40000
+; CHECK-NEXT: movq %fs:40, %rax
+; CHECK-NEXT: movq %rax, {{[0-9]+}}(%rsp)
; CHECK-NEXT: movq %rsp, %rax
; CHECK-NEXT: .p2align 4
; CHECK-NEXT: .LBB1_1: # %loop
@@ -47,10 +62,17 @@ define void @test_phi_loop(i1 %c) sspstrong {
; CHECK-NEXT: testb $1, %dil
; CHECK-NEXT: jne .LBB1_1
; CHECK-NEXT: # %bb.2: # %exit
+; CHECK-NEXT: movq %fs:40, %rax
+; CHECK-NEXT: cmpq {{[0-9]+}}(%rsp), %rax
+; CHECK-NEXT: jne .LBB1_4
+; CHECK-NEXT: # %bb.3: # %SP_return
; CHECK-NEXT: movq %rbp, %rsp
; CHECK-NEXT: popq %rbp
; CHECK-NEXT: .cfi_def_cfa %rsp, 8
; CHECK-NEXT: retq
+; CHECK-NEXT: .LBB1_4: # %CallStackCheckFailBlk
+; CHECK-NEXT: .cfi_def_cfa %rbp, 16
+; CHECK-NEXT: callq __stack_chk_fail at PLT
entry:
%a = alloca <10000 x i64>
br label %loop
``````````
</details>
https://github.com/llvm/llvm-project/pull/129248
More information about the llvm-commits
mailing list