[llvm] [StackProtector] Fix phi handling in HasAddressTaken() (PR #129248)
Nikita Popov via llvm-commits
llvm-commits at lists.llvm.org
Fri Feb 28 06:16:39 PST 2025
https://github.com/nikic updated https://github.com/llvm/llvm-project/pull/129248
>From 9f7c0c2e15b2b08dc612759b23e75a8a41a9f81c Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Thu, 27 Feb 2025 18:23:30 +0100
Subject: [PATCH] [StackProtector] Fix phi handling in HasAddressTaken()
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
remainin 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.
---
llvm/lib/CodeGen/StackProtector.cpp | 30 +++++++++++++++++---
llvm/test/CodeGen/X86/stack-protector-phi.ll | 30 +++++++++++++++++---
2 files changed, 52 insertions(+), 8 deletions(-)
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
More information about the llvm-commits
mailing list