[llvm] 53c1579 - [StackProtector] Fix phi handling in HasAddressTaken() (#129248)
via llvm-commits
llvm-commits at lists.llvm.org
Wed Mar 5 03:45:16 PST 2025
Author: Nikita Popov
Date: 2025-03-05T12:45:13+01:00
New Revision: 53c157939e5ac9acc8e1f8853325a021bc925501
URL: https://github.com/llvm/llvm-project/commit/53c157939e5ac9acc8e1f8853325a021bc925501
DIFF: https://github.com/llvm/llvm-project/commit/53c157939e5ac9acc8e1f8853325a021bc925501.diff
LOG: [StackProtector] Fix phi handling in HasAddressTaken() (#129248)
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.
Added:
Modified:
llvm/lib/CodeGen/StackProtector.cpp
llvm/test/CodeGen/X86/stack-protector-phi.ll
Removed:
################################################################################
diff --git a/llvm/lib/CodeGen/StackProtector.cpp b/llvm/lib/CodeGen/StackProtector.cpp
index 232e84fb672e0..8aa3f9d4f465a 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 observed 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
More information about the llvm-commits
mailing list