[llvm] MachineVerifier: Check stack protector is top-most in frame (PR #121481)

[Guy David via llvm-commits]

https://github.com/guy-david updated https://github.com/llvm/llvm-project/pull/121481

>From 2b6d2be1446eb86ff8883e37957fabd0e3330855 Mon Sep 17 00:00:00 2001
From: Guy David <guyda96 at gmail.com>
Date: Thu, 2 Jan 2025 16:28:13 +0200
Subject: [PATCH] MachineVerifier: Check stack protector is top-most in frame

Mitigate against potential bugs that might place it elsewhere and render
the mechanism useless.
---
 llvm/lib/CodeGen/MachineVerifier.cpp          | 47 ++++++++++++++++++-
 .../stack-protector-offset.mir                | 17 +++++++
 2 files changed, 63 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/MachineVerifier/stack-protector-offset.mir

diff --git a/llvm/lib/CodeGen/MachineVerifier.cpp b/llvm/lib/CodeGen/MachineVerifier.cpp
index bec36b728ae328..968586be099119 100644
--- a/llvm/lib/CodeGen/MachineVerifier.cpp
+++ b/llvm/lib/CodeGen/MachineVerifier.cpp
@@ -353,6 +353,8 @@ struct MachineVerifier {
                        LaneBitmask LaneMask = LaneBitmask::getNone());
 
   void verifyStackFrame();
+  // Check that the stack protector is the top-most object in the stack.
+  void verifyStackProtector();
 
   void verifySlotIndexes() const;
   void verifyProperties(const MachineFunction &MF);
@@ -709,8 +711,10 @@ void MachineVerifier::visitMachineFunctionBefore() {
   // Check that the register use lists are sane.
   MRI->verifyUseLists();
 
-  if (!MF->empty())
+  if (!MF->empty()) {
     verifyStackFrame();
+    verifyStackProtector();
+  }
 }
 
 void
@@ -4038,3 +4042,44 @@ void MachineVerifier::verifyStackFrame() {
     }
   }
 }
+
+void MachineVerifier::verifyStackProtector() {
+  const MachineFrameInfo &MFI = MF->getFrameInfo();
+  if (!MFI.hasStackProtectorIndex())
+    return;
+  // Only applicable when the offsets of frame objects have been determined,
+  // which is indicated by a non-zero stack size.
+  if (!MFI.getStackSize())
+    return;
+  const TargetFrameLowering &TFI = *MF->getSubtarget().getFrameLowering();
+  bool StackGrowsDown =
+      TFI.getStackGrowthDirection() == TargetFrameLowering::StackGrowsDown;
+  assert(StackGrowsDown && "Only support stack growth down");
+  // Collect the frame indices of the callee-saved registers which are spilled
+  // to the stack. These are the registers that are stored above the stack
+  // protector.
+  SmallSet<unsigned, 4> CalleeSavedFrameIndices;
+  if (MFI.isCalleeSavedInfoValid()) {
+    for (const CalleeSavedInfo &Info : MFI.getCalleeSavedInfo()) {
+      if (!Info.isSpilledToReg())
+        CalleeSavedFrameIndices.insert(Info.getFrameIdx());
+    }
+  }
+  unsigned FI = MFI.getStackProtectorIndex();
+  int64_t SPOffset = MFI.getObjectOffset(FI);
+  for (unsigned I = 0, E = MFI.getObjectIndexEnd(); I != E; ++I) {
+    if (I == FI)
+      continue;
+    // Variable-sized objects do not have a fixed offset.
+    if (MFI.isVariableSizedObjectIndex(I))
+      continue;
+    if (CalleeSavedFrameIndices.contains(I))
+      continue;
+    if (SPOffset < MFI.getObjectOffset(I)) {
+      report("Stack protector is not the top-most object on the stack", MF);
+      OS << "Stack protector is not the top-most object on the stack in "
+         << MF->getName() << '\n';
+      break;
+    }
+  }
+}
diff --git a/llvm/test/MachineVerifier/stack-protector-offset.mir b/llvm/test/MachineVerifier/stack-protector-offset.mir
new file mode 100644
index 00000000000000..3fbeab2eb3b46b
--- /dev/null
+++ b/llvm/test/MachineVerifier/stack-protector-offset.mir
@@ -0,0 +1,17 @@
+# RUN: not --crash llc -mtriple=aarch64 -run-pass=machineverifier -o - %s 2>&1 | FileCheck %s
+# REQUIRES: aarch64-registered-target
+
+# CHECK: *** Bad machine code: Stack protector is not the top-most object on the stack ***
+
+---
+name:            main
+frameInfo:
+  stackSize:       16
+  stackProtector:  '%stack.1'
+  isCalleeSavedInfoValid: true
+stack:
+  - { id: 0, offset: -16, size: 8, alignment: 8, stack-id: default }
+  - { id: 1, offset: -32, size: 8, alignment: 8, stack-id: default }
+body:             |
+  bb.0:
+...