$OpenBSD: patch-lib_Target_X86_X86ReturnProtectorLowering_cpp,v 1.1 2018/08/21 06:56:09 ajacoutot Exp $

Refactor retguard to make adding additional arches easier.

Index: lib/Target/X86/X86ReturnProtectorLowering.cpp
--- lib/Target/X86/X86ReturnProtectorLowering.cpp.orig
+++ lib/Target/X86/X86ReturnProtectorLowering.cpp
@@ -0,0 +1,118 @@
+//===-- X86ReturnProtectorLowering.cpp - ----------------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains the X86 implementation of ReturnProtectorLowering class.
+//
+//===----------------------------------------------------------------------===//
+
+#include "X86ReturnProtectorLowering.h"
+#include "X86InstrBuilder.h"
+#include "X86InstrInfo.h"
+#include "X86MachineFunctionInfo.h"
+#include "X86Subtarget.h"
+#include "X86TargetMachine.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineInstrBuilder.h"
+#include "llvm/CodeGen/MachineModuleInfo.h"
+#include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/IR/Function.h"
+#include "llvm/MC/MCAsmInfo.h"
+#include "llvm/MC/MCSymbol.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Target/TargetOptions.h"
+#include <cstdlib>
+
+using namespace llvm;
+
+void X86ReturnProtectorLowering::insertReturnProtectorPrologue(
+    MachineFunction &MF, MachineBasicBlock &MBB, GlobalVariable *cookie) const {
+
+  MachineBasicBlock::instr_iterator MI = MBB.instr_begin();
+  DebugLoc MBBDL = MBB.findDebugLoc(MI);
+  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)
+      .addReg(0)
+      .addGlobalAddress(cookie)
+      .addReg(0);
+  addDirectMem(BuildMI(MBB, MI, MBBDL, TII->get(X86::XOR64rm), REG).addReg(REG),
+               X86::RSP);
+}
+
+void X86ReturnProtectorLowering::insertReturnProtectorEpilogue(
+    MachineFunction &MF, MachineInstr &MI, GlobalVariable *cookie) const {
+
+  MachineBasicBlock &MBB = *MI.getParent();
+  DebugLoc MBBDL = MI.getDebugLoc();
+  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))
+      .addReg(REG)
+      .addReg(X86::RIP)
+      .addImm(0)
+      .addReg(0)
+      .addGlobalAddress(cookie)
+      .addReg(0);
+  BuildMI(MBB, MI, MBBDL, TII->get(X86::RETGUARD_JMP_TRAP));
+}
+
+bool X86ReturnProtectorLowering::opcodeIsReturn(unsigned opcode) const {
+  switch (opcode) {
+  case X86::RET:
+  case X86::RETL:
+  case X86::RETQ:
+  case X86::RETW:
+  case X86::RETIL:
+  case X86::RETIQ:
+  case X86::RETIW:
+  case X86::LRETL:
+  case X86::LRETQ:
+  case X86::LRETW:
+  case X86::LRETIL:
+  case X86::LRETIQ:
+  case X86::LRETIW:
+    return true;
+  default:
+    return false;
+  }
+}
+
+void X86ReturnProtectorLowering::fillTempRegisters(
+    MachineFunction &MF, std::vector<unsigned> &TempRegs) const {
+
+  TempRegs.push_back(X86::R11);
+  TempRegs.push_back(X86::R10);
+  const Function &F = MF.getFunction();
+  if (!F.isVarArg()) {
+    // We can use any of the caller saved unused arg registers
+    switch (F.arg_size()) {
+    case 0:
+      TempRegs.push_back(X86::RDI);
+    case 1:
+      TempRegs.push_back(X86::RSI);
+    case 2: // RDX is the 2nd return register
+    case 3:
+      TempRegs.push_back(X86::RCX);
+    case 4:
+      TempRegs.push_back(X86::R8);
+    case 5:
+      TempRegs.push_back(X86::R9);
+    default:
+      break;
+    }
+  }
+}
