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@

7 files changed, 69 insertions, 14 deletions

diff --git a/gnu/llvm/include/llvm/CodeGen/MachineFrameInfo.h b/gnu/llvm/include/llvm/CodeGen/MachineFrameInfo.h
index 9f590f77dc0..412c1ef7e5a 100644
--- a/gnu/llvm/include/llvm/CodeGen/MachineFrameInfo.h
+++ b/gnu/llvm/include/llvm/CodeGen/MachineFrameInfo.h
@@ -262,6 +262,8 @@ private:
     unsigned Register = 0;
     /// Set to true if this function needs return protectors
     bool Needed = false;
+    /// Does the return protector cookie need to be stored in frame
+    bool NeedsStore = true;
   } RPI;
 
   /// The frame index for the function context. Used for SjLj exceptions.
@@ -348,6 +350,9 @@ public:
   /// Get / Set if this frame needs a return protector
   void setReturnProtectorNeeded(bool I) { RPI.Needed = I; }
   bool getReturnProtectorNeeded() const { return RPI.Needed; }
+  /// Get / Set if the return protector cookie needs to be stored in frame
+  void setReturnProtectorNeedsStore(bool I) { RPI.NeedsStore = I; }
+  bool getReturnProtectorNeedsStore() const { return RPI.NeedsStore; }
 
   /// Return the index for the function context object.
   /// This object is used for SjLj exceptions.
diff --git a/gnu/llvm/include/llvm/CodeGen/ReturnProtectorLowering.h b/gnu/llvm/include/llvm/CodeGen/ReturnProtectorLowering.h
index fb9ef0c908f..5e0b4c818cd 100644
--- a/gnu/llvm/include/llvm/CodeGen/ReturnProtectorLowering.h
+++ b/gnu/llvm/include/llvm/CodeGen/ReturnProtectorLowering.h
@@ -28,6 +28,7 @@ class MachineInstr;
 
 class ReturnProtectorLowering {
 public:
+  virtual ~ReturnProtectorLowering() {}
   /// Subclass interface - subclasses need to implement these functions.
 
   /// insertReturnProtectorPrologue/Epilogue - insert return protector
@@ -57,7 +58,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;
 
   /// determineReturnProtectorTempRegister - Find a register that can be used
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;
     }