Do not store the retguard cookie in frame in leaf functions if possible.

Makes things slightly faster and also improves security in these functions,
since the retguard cookie can't leak via the stack.

ok deraadt@

5 files changed, 62 insertions, 13 deletions

diff --git a/gnu/llvm/lib/CodeGen/ReturnProtectorLowering.cpp b/gnu/llvm/lib/CodeGen/ReturnProtectorLowering.cpp
index 60ad3ede793..99b743c4522 100644
--- a/gnu/llvm/lib/CodeGen/ReturnProtectorLowering.cpp
+++ b/gnu/llvm/lib/CodeGen/ReturnProtectorLowering.cpp
@@ -64,7 +64,7 @@ void ReturnProtectorLowering::setupReturnProtector(MachineFunction &MF) const {
 /// saveReturnProtectorRegister - Allows the target to save the
 /// ReturnProtectorRegister in the CalleeSavedInfo vector if needed.
 void ReturnProtectorLowering::saveReturnProtectorRegister(
-    const MachineFunction &MF, std::vector<CalleeSavedInfo> &CSI) const {
+    MachineFunction &MF, std::vector<CalleeSavedInfo> &CSI) const {
   const MachineFrameInfo &MFI = MF.getFrameInfo();
   if (!MFI.getReturnProtectorNeeded())
     return;
@@ -72,7 +72,15 @@ void ReturnProtectorLowering::saveReturnProtectorRegister(
   if (!MFI.hasReturnProtectorRegister())
     llvm_unreachable("Saving unset return protector register");
 
-  CSI.push_back(CalleeSavedInfo(MFI.getReturnProtectorRegister()));
+  unsigned Reg = MFI.getReturnProtectorRegister();
+  if (MFI.getReturnProtectorNeedsStore())
+    CSI.push_back(CalleeSavedInfo(Reg));
+  else {
+    for (auto &MBB : MF) {
+      if (!MBB.isLiveIn(Reg))
+        MBB.addLiveIn(Reg);
+    }
+  }
 }
 
 /// determineReturnProtectorTempRegister - Find a register that can be used
@@ -88,6 +96,32 @@ bool ReturnProtectorLowering::determineReturnProtectorRegister(
 
   const TargetRegisterInfo *TRI = MF.getSubtarget().getRegisterInfo();
 
+  std::vector<unsigned> TempRegs;
+  fillTempRegisters(MF, TempRegs);
+
+  // For leaf functions, try to find a free register that is available
+  // in every BB, so we do not need to store it in the frame at all.
+  if (!MFI.hasCalls() && !MFI.hasTailCall()) {
+    SmallSet<unsigned, 16> LeafUsed;
+    SmallSet<int, 24> LeafVisited;
+    markUsedRegsInSuccessors(MF.front(), LeafUsed, LeafVisited);
+    for (unsigned Reg : TempRegs) {
+      bool canUse = true;
+      for (MCRegAliasIterator AI(Reg, TRI, true); AI.isValid(); ++AI) {
+        if (LeafUsed.count(*AI)) {
+          canUse = false;
+          break;
+        }
+      }
+      if (canUse) {
+        MFI.setReturnProtectorRegister(Reg);
+        MFI.setReturnProtectorNeedsStore(false);
+        return true;
+      }
+    }
+  }
+
+  // For non-leaf functions, we only need to search save / restore blocks
   SmallSet<unsigned, 16> Used;
   SmallSet<int, 24> Visited;
 
@@ -115,9 +149,6 @@ bool ReturnProtectorLowering::determineReturnProtectorRegister(
 
   // Now we have gathered all the regs used outside the frame save / restore,
   // so we can see if we have a free reg to use for the retguard cookie.
-  std::vector<unsigned> TempRegs;
-  fillTempRegisters(MF, TempRegs);
-
   for (unsigned Reg : TempRegs) {
     bool canUse = true;
     for (MCRegAliasIterator AI(Reg, TRI, true); AI.isValid(); ++AI) {
@@ -157,12 +188,17 @@ void ReturnProtectorLowering::insertReturnProtectors(
   if (!cookie)
     llvm_unreachable("Function needs return protector but no cookie assigned");
 
+  unsigned Reg = MFI.getReturnProtectorRegister();
+
   std::vector<MachineInstr *> returns;
   for (auto &MBB : MF) {
     if (MBB.isReturnBlock()) {
       for (auto &MI : MBB.terminators()) {
-        if (opcodeIsReturn(MI.getOpcode()))
+        if (opcodeIsReturn(MI.getOpcode())) {
           returns.push_back(&MI);
+          if (!MBB.isLiveIn(Reg))
+            MBB.addLiveIn(Reg);
+        }
       }
     }
   }
@@ -174,4 +210,7 @@ void ReturnProtectorLowering::insertReturnProtectors(
     insertReturnProtectorEpilogue(MF, *MI, cookie);
 
   insertReturnProtectorPrologue(MF, MF.front(), cookie);
+
+  if (!MF.front().isLiveIn(Reg))
+    MF.front().addLiveIn(Reg);
 }
diff --git a/gnu/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp b/gnu/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
index 33cdc9d62ec..c0fccc3fd4d 100644
--- a/gnu/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
+++ b/gnu/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
@@ -1423,7 +1423,7 @@ void AArch64FrameLowering::determineCalleeSaves(MachineFunction &MF,
       SpillEstimate++;
   }
 
-  if (MFI.hasReturnProtectorRegister()) {
+  if (MFI.hasReturnProtectorRegister() && MFI.getReturnProtectorNeedsStore()) {
     SavedRegs.set(MFI.getReturnProtectorRegister());
     SpillEstimate++;
   }
diff --git a/gnu/llvm/lib/Target/AArch64/AArch64ReturnProtectorLowering.cpp b/gnu/llvm/lib/Target/AArch64/AArch64ReturnProtectorLowering.cpp
index 9d4bdd9591b..2ec92b17caf 100644
--- a/gnu/llvm/lib/Target/AArch64/AArch64ReturnProtectorLowering.cpp
+++ b/gnu/llvm/lib/Target/AArch64/AArch64ReturnProtectorLowering.cpp
@@ -39,7 +39,6 @@ void AArch64ReturnProtectorLowering::insertReturnProtectorPrologue(
   const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
   unsigned REG = MF.getFrameInfo().getReturnProtectorRegister();
 
-  MBB.addLiveIn(REG);
   BuildMI(MBB, MI, MBBDL, TII->get(AArch64::ADRP), REG)
       .addGlobalAddress(cookie, 0, AArch64II::MO_PAGE);
   BuildMI(MBB, MI, MBBDL, TII->get(AArch64::LDRXui), REG)
@@ -58,7 +57,6 @@ void AArch64ReturnProtectorLowering::insertReturnProtectorEpilogue(
   const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
   unsigned REG = MF.getFrameInfo().getReturnProtectorRegister();
 
-  MBB.addLiveIn(REG);
   MBB.addLiveIn(AArch64::X9);
   // REG holds the cookie we calculated in prologue. We use X9 as a
   // scratch reg to pull the random data. XOR REG with LR should yield
@@ -99,7 +97,7 @@ void AArch64ReturnProtectorLowering::fillTempRegisters(
 }
 
 void AArch64ReturnProtectorLowering::saveReturnProtectorRegister(
-    const MachineFunction &MF, std::vector<CalleeSavedInfo> &CSI) const {
+    MachineFunction &MF, std::vector<CalleeSavedInfo> &CSI) const {
 
   const MachineFrameInfo &MFI = MF.getFrameInfo();
   if (!MFI.getReturnProtectorNeeded())
@@ -108,6 +106,15 @@ void AArch64ReturnProtectorLowering::saveReturnProtectorRegister(
   if (!MFI.hasReturnProtectorRegister())
     llvm_unreachable("Saving unset return protector register");
 
+  unsigned Reg = MFI.getReturnProtectorRegister();
+  if (!MFI.getReturnProtectorNeedsStore()) {
+    for (auto &MBB : MF) {
+      if (!MBB.isLiveIn(Reg))
+        MBB.addLiveIn(Reg);
+    }
+    return;
+  }
+
   // Put the temp reg after FP and LR to avoid layout issues
   // with the D registers later.
   bool added = false;
diff --git a/gnu/llvm/lib/Target/AArch64/AArch64ReturnProtectorLowering.h b/gnu/llvm/lib/Target/AArch64/AArch64ReturnProtectorLowering.h
index 47feb002978..9537b4643f2 100644
--- a/gnu/llvm/lib/Target/AArch64/AArch64ReturnProtectorLowering.h
+++ b/gnu/llvm/lib/Target/AArch64/AArch64ReturnProtectorLowering.h
@@ -43,7 +43,7 @@ public:
   /// saveReturnProtectorRegister - Allows the target to save the
   /// CalculationRegister in the CalleeSavedInfo vector if needed.
   virtual void
-  saveReturnProtectorRegister(const MachineFunction &MF,
+  saveReturnProtectorRegister(MachineFunction &MF,
                               std::vector<CalleeSavedInfo> &CSI) const override;
 };
 
diff --git a/gnu/llvm/lib/Target/X86/X86ReturnProtectorLowering.cpp b/gnu/llvm/lib/Target/X86/X86ReturnProtectorLowering.cpp
index c881e2ab6f3..2eea9ac11fb 100644
--- a/gnu/llvm/lib/Target/X86/X86ReturnProtectorLowering.cpp
+++ b/gnu/llvm/lib/Target/X86/X86ReturnProtectorLowering.cpp
@@ -38,7 +38,6 @@ void X86ReturnProtectorLowering::insertReturnProtectorPrologue(
   const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
   unsigned REG = MF.getFrameInfo().getReturnProtectorRegister();
 
-  MBB.addLiveIn(REG);
   BuildMI(MBB, MI, MBBDL, TII->get(X86::MOV64rm), REG)
       .addReg(X86::RIP)
       .addImm(0)
@@ -57,7 +56,6 @@ void X86ReturnProtectorLowering::insertReturnProtectorEpilogue(
   const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
   unsigned REG = MF.getFrameInfo().getReturnProtectorRegister();
 
-  MBB.addLiveIn(REG);
   addDirectMem(BuildMI(MBB, MI, MBBDL, TII->get(X86::XOR64rm), REG).addReg(REG),
                X86::RSP);
   BuildMI(MBB, MI, MBBDL, TII->get(X86::CMP64rm))
@@ -102,15 +100,20 @@ void X86ReturnProtectorLowering::fillTempRegisters(
     switch (F.arg_size()) {
     case 0:
       TempRegs.push_back(X86::RDI);
+      LLVM_FALLTHROUGH;
     case 1:
       TempRegs.push_back(X86::RSI);
+      LLVM_FALLTHROUGH;
     case 2: // RDX is the 2nd return register
     case 3:
       TempRegs.push_back(X86::RCX);
+      LLVM_FALLTHROUGH;
     case 4:
       TempRegs.push_back(X86::R8);
+      LLVM_FALLTHROUGH;
     case 5:
       TempRegs.push_back(X86::R9);
+      LLVM_FALLTHROUGH;
     default:
       break;
     }