[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.

Author: nikic

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);

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@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@PLT
 entry:
   %a = alloca <10000 x i64>
   br label %loop