Cherrypick patches llvm lvi bug

llvm/lib/Target/X86/X86FrameLowering.cpp

@@ -148,60 +148,6 @@ static unsigned getLEArOpcode(bool IsLP64) {
   return IsLP64 ? X86::LEA64r : X86::LEA32r;
 }
 
-/// findDeadCallerSavedReg - Return a caller-saved register that isn't live
-/// when it reaches the "return" instruction. We can then pop a stack object
-/// to this register without worry about clobbering it.
-static unsigned findDeadCallerSavedReg(MachineBasicBlock &MBB,
-                                       MachineBasicBlock::iterator &MBBI,
-                                       const X86RegisterInfo *TRI,
-                                       bool Is64Bit) {
-  const MachineFunction *MF = MBB.getParent();
-  if (MF->callsEHReturn())
-    return 0;
-
-  const TargetRegisterClass &AvailableRegs = *TRI->getGPRsForTailCall(*MF);
-
-  if (MBBI == MBB.end())
-    return 0;
-
-  switch (MBBI->getOpcode()) {
-  default: return 0;
-  case TargetOpcode::PATCHABLE_RET:
-  case X86::RET:
-  case X86::RETL:
-  case X86::RETQ:
-  case X86::RETIL:
-  case X86::RETIQ:
-  case X86::TCRETURNdi:
-  case X86::TCRETURNri:
-  case X86::TCRETURNmi:
-  case X86::TCRETURNdi64:
-  case X86::TCRETURNri64:
-  case X86::TCRETURNmi64:
-  case X86::EH_RETURN:
-  case X86::EH_RETURN64: {
-    SmallSet<uint16_t, 8> Uses;
-    for (unsigned i = 0, e = MBBI->getNumOperands(); i != e; ++i) {
-      MachineOperand &MO = MBBI->getOperand(i);
-      if (!MO.isReg() || MO.isDef())
-        continue;
-      Register Reg = MO.getReg();
-      if (!Reg)
-        continue;
-      for (MCRegAliasIterator AI(Reg, TRI, true); AI.isValid(); ++AI)
-        Uses.insert(*AI);
-    }
-
-    for (auto CS : AvailableRegs)
-      if (!Uses.count(CS) && CS != X86::RIP && CS != X86::RSP &&
-          CS != X86::ESP)
-        return CS;
-  }
-  }
-
-  return 0;
-}
-
 static bool isEAXLiveIn(MachineBasicBlock &MBB) {
   for (MachineBasicBlock::RegisterMaskPair RegMask : MBB.liveins()) {
     unsigned Reg = RegMask.PhysReg;
@@ -288,7 +234,7 @@ void X86FrameLowering::emitSPUpdate(MachineBasicBlock &MBB,
     if (isSub && !isEAXLiveIn(MBB))
       Reg = Rax;
     else
-      Reg = findDeadCallerSavedReg(MBB, MBBI, TRI, Is64Bit);
+      Reg = TRI->findDeadCallerSavedReg(MBB, MBBI);
 
     unsigned MovRIOpc = Is64Bit ? X86::MOV64ri : X86::MOV32ri;
     unsigned AddSubRROpc =
@@ -345,7 +291,7 @@ void X86FrameLowering::emitSPUpdate(MachineBasicBlock &MBB,
       // need to find a dead register when using pop.
       unsigned Reg = isSub
         ? (unsigned)(Is64Bit ? X86::RAX : X86::EAX)
-        : findDeadCallerSavedReg(MBB, MBBI, TRI, Is64Bit);
+        : TRI->findDeadCallerSavedReg(MBB, MBBI);
       if (Reg) {
         unsigned Opc = isSub
           ? (Is64Bit ? X86::PUSH64r : X86::PUSH32r)

llvm/lib/Target/X86/X86LoadValueInjectionRetHardening.cpp

@@ -72,62 +72,39 @@ bool X86LoadValueInjectionRetHardeningPass::runOnMachineFunction(
   ++NumFunctionsConsidered;
   const X86RegisterInfo *TRI = Subtarget->getRegisterInfo();
   const X86InstrInfo *TII = Subtarget->getInstrInfo();
-  unsigned ClobberReg = X86::NoRegister;
-  std::bitset<X86::NUM_TARGET_REGS> UnclobberableGR64s;
-  UnclobberableGR64s.set(X86::RSP); // can't clobber stack pointer
-  UnclobberableGR64s.set(X86::RIP); // can't clobber instruction pointer
-  UnclobberableGR64s.set(X86::RAX); // used for function return
-  UnclobberableGR64s.set(X86::RDX); // used for function return
-
-  // We can clobber any register allowed by the function's calling convention.
-  for (const MCPhysReg *PR = TRI->getCalleeSavedRegs(&MF); auto Reg = *PR; ++PR)
-    UnclobberableGR64s.set(Reg);
-  for (auto &Reg : X86::GR64RegClass) {
-    if (!UnclobberableGR64s.test(Reg)) {
-      ClobberReg = Reg;
-      break;
-    }
-  }
-
-  if (ClobberReg != X86::NoRegister) {
-    LLVM_DEBUG(dbgs() << "Selected register "
-                      << Subtarget->getRegisterInfo()->getRegAsmName(ClobberReg)
-                      << " to clobber\n");
-  } else {
-    LLVM_DEBUG(dbgs() << "Could not find a register to clobber\n");
-  }
 
   bool Modified = false;
   for (auto &MBB : MF) {
-    if (MBB.empty())
-      continue;
-
-    MachineInstr &MI = MBB.back();
-    if (MI.getOpcode() != X86::RETQ)
-      continue;
-
-    if (ClobberReg != X86::NoRegister) {
-      MBB.erase_instr(&MI);
-      BuildMI(MBB, MBB.end(), DebugLoc(), TII->get(X86::POP64r))
-          .addReg(ClobberReg, RegState::Define)
-          .setMIFlag(MachineInstr::FrameDestroy);
-      BuildMI(MBB, MBB.end(), DebugLoc(), TII->get(X86::LFENCE));
-      BuildMI(MBB, MBB.end(), DebugLoc(), TII->get(X86::JMP64r))
-          .addReg(ClobberReg);
-    } else {
-      // In case there is no available scratch register, we can still read from
-      // RSP to assert that RSP points to a valid page. The write to RSP is
-      // also helpful because it verifies that the stack's write permissions
-      // are intact.
-      MachineInstr *Fence = BuildMI(MBB, MI, DebugLoc(), TII->get(X86::LFENCE));
-      addRegOffset(BuildMI(MBB, Fence, DebugLoc(), TII->get(X86::SHL64mi)),
-                   X86::RSP, false, 0)
-          .addImm(0)
-          ->addRegisterDead(X86::EFLAGS, TRI);
+    for (auto MBBI = MBB.begin(); MBBI != MBB.end(); ++MBBI) {
+      if (MBBI->getOpcode() != X86::RETQ)
+        continue;
+
+      unsigned ClobberReg = TRI->findDeadCallerSavedReg(MBB, MBBI);
+      if (ClobberReg != X86::NoRegister) {
+        BuildMI(MBB, MBBI, DebugLoc(), TII->get(X86::POP64r))
+            .addReg(ClobberReg, RegState::Define)
+            .setMIFlag(MachineInstr::FrameDestroy);
+        BuildMI(MBB, MBBI, DebugLoc(), TII->get(X86::LFENCE));
+        BuildMI(MBB, MBBI, DebugLoc(), TII->get(X86::JMP64r))
+            .addReg(ClobberReg);
+        MBB.erase(MBBI);
+      } else {
+        // In case there is no available scratch register, we can still read
+        // from RSP to assert that RSP points to a valid page. The write to RSP
+        // is also helpful because it verifies that the stack's write
+        // permissions are intact.
+        MachineInstr *Fence =
+            BuildMI(MBB, MBBI, DebugLoc(), TII->get(X86::LFENCE));
+        addRegOffset(BuildMI(MBB, Fence, DebugLoc(), TII->get(X86::SHL64mi)),
+                     X86::RSP, false, 0)
+            .addImm(0)
+            ->addRegisterDead(X86::EFLAGS, TRI);
+      }
+
+      ++NumFences;
+      Modified = true;
+      break;
     }
-
-    ++NumFences;
-    Modified = true;
   }
 
   if (Modified)

llvm/lib/Target/X86/X86RegisterInfo.cpp

@@ -18,6 +18,7 @@
 #include "X86Subtarget.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
@@ -790,6 +791,55 @@ X86RegisterInfo::eliminateFrameIndex(MachineBasicBlock::iterator II,
   }
 }
 
+unsigned X86RegisterInfo::findDeadCallerSavedReg(
+    MachineBasicBlock &MBB, MachineBasicBlock::iterator &MBBI) const {
+  const MachineFunction *MF = MBB.getParent();
+  if (MF->callsEHReturn())
+    return 0;
+
+  const TargetRegisterClass &AvailableRegs = *getGPRsForTailCall(*MF);
+
+  if (MBBI == MBB.end())
+    return 0;
+
+  switch (MBBI->getOpcode()) {
+  default:
+    return 0;
+  case TargetOpcode::PATCHABLE_RET:
+  case X86::RET:
+  case X86::RETL:
+  case X86::RETQ:
+  case X86::RETIL:
+  case X86::RETIQ:
+  case X86::TCRETURNdi:
+  case X86::TCRETURNri:
+  case X86::TCRETURNmi:
+  case X86::TCRETURNdi64:
+  case X86::TCRETURNri64:
+  case X86::TCRETURNmi64:
+  case X86::EH_RETURN:
+  case X86::EH_RETURN64: {
+    SmallSet<uint16_t, 8> Uses;
+    for (unsigned I = 0, E = MBBI->getNumOperands(); I != E; ++I) {
+      MachineOperand &MO = MBBI->getOperand(I);
+      if (!MO.isReg() || MO.isDef())
+        continue;
+      Register Reg = MO.getReg();
+      if (!Reg)
+        continue;
+      for (MCRegAliasIterator AI(Reg, this, true); AI.isValid(); ++AI)
+        Uses.insert(*AI);
+    }
+
+    for (auto CS : AvailableRegs)
+      if (!Uses.count(CS) && CS != X86::RIP && CS != X86::RSP && CS != X86::ESP)
+        return CS;
+  }
+  }
+
+  return 0;
+}
+
 Register X86RegisterInfo::getFrameRegister(const MachineFunction &MF) const {
   const X86FrameLowering *TFI = getFrameLowering(MF);
   return TFI->hasFP(MF) ? FramePtr : StackPtr;

llvm/lib/Target/X86/X86RegisterInfo.h

@@ -128,6 +128,12 @@ class X86RegisterInfo final : public X86GenRegisterInfo {
                            int SPAdj, unsigned FIOperandNum,
                            RegScavenger *RS = nullptr) const override;
 
+  /// findDeadCallerSavedReg - Return a caller-saved register that isn't live
+  /// when it reaches the "return" instruction. We can then pop a stack object
+  /// to this register without worry about clobbering it.
+  unsigned findDeadCallerSavedReg(MachineBasicBlock &MBB,
+                                  MachineBasicBlock::iterator &MBBI) const;
+
   // Debug information queries.
   Register getFrameRegister(const MachineFunction &MF) const override;
   unsigned getPtrSizedFrameRegister(const MachineFunction &MF) const;

llvm/test/CodeGen/X86/lvi-hardening-ret.ll

@@ -41,9 +41,9 @@ entry:
   %add = add nsw i32 %0, %1
   ret i32 %add
 ; CHECK-NOT:  retq
-; CHECK:      shlq $0, (%{{[^ ]*}})
+; CHECK:      popq %rcx
 ; CHECK-NEXT: lfence
-; CHECK-NEXT: retq
+; CHECK-NEXT: jmpq *%rcx
 }
 
 ; Function Attrs: noinline nounwind optnone uwtable
@@ -52,9 +52,9 @@ define dso_local preserve_mostcc void @preserve_most() #0 {
 entry:
   ret void
 ; CHECK-NOT:  retq
-; CHECK:      popq %r11
+; CHECK:      popq %rax
 ; CHECK-NEXT: lfence
-; CHECK-NEXT: jmpq *%r11
+; CHECK-NEXT: jmpq *%rax
 }
 
 ; Function Attrs: noinline nounwind optnone uwtable
@@ -63,9 +63,18 @@ define dso_local preserve_allcc void @preserve_all() #0 {
 entry:
   ret void
 ; CHECK-NOT:  retq
-; CHECK:      popq %r11
+; CHECK:      popq %rax
 ; CHECK-NEXT: lfence
-; CHECK-NEXT: jmpq *%r11
+; CHECK-NEXT: jmpq *%rax
+}
+
+define { i64, i128 } @ret_i64_i128() #0 {
+; CHECK-LABEL: ret_i64_i128:
+  ret { i64, i128 } { i64 1, i128 36893488147419103235 }
+; CHECK-NOT:  retq
+; CHECK:      popq %rsi
+; CHECK-NEXT: lfence
+; CHECK-NEXT: jmpq *%rsi
 }
 
 attributes #0 = { "target-features"="+lvi-cfi" }