[X86] Support the "ms-hotpatch" attribute.

Summary:
Based on two patches by Michael Mueller.

This is a target attribute that causes a function marked with it to be
emitted as "hotpatchable". This particular mechanism was originally
devised by Microsoft for patching their binaries (which they are
constantly updating to stay ahead of crackers, script kiddies, and other
ne'er-do-wells on the Internet), but is now commonly abused by Windows
programs to hook API functions.

This mechanism is target-specific. For x86, a two-byte no-op instruction
is emitted at the function's entry point; the entry point must be
immediately preceded by 64 (32-bit) or 128 (64-bit) bytes of padding.
This padding is where the patch code is written. The two byte no-op is
then overwritten with a short jump into this code. The no-op is usually
a `movl %edi, %edi` instruction; this is used as a magic value
indicating that this is a hotpatchable function.

Reviewers: majnemer, sanjoy, rnk

Subscribers: dberris, llvm-commits

Differential Revision: https://reviews.llvm.org/D19908

llvm-svn: 278048

12 files changed, 168 insertions, 22 deletions

diff --git a/llvm/docs/LangRef.rst b/llvm/docs/LangRef.rst
index 7f69d693294..82f923ee1f7 100644
--- a/llvm/docs/LangRef.rst
+++ b/llvm/docs/LangRef.rst
@@ -1448,7 +1448,7 @@ example:
     generated for this function needs to follow certain conventions that
     make it possible for a runtime function to patch over it later.
     The exact effect of this attribute depends on its string value,
-    for which there currently is one legal possibility:
+    for which there currently are two legal possiblities:
 
      * ``"prologue-short-redirect"`` - This style of patchable
        function is intended to support patching a function prologue to
@@ -1464,6 +1464,24 @@ example:
 
        ``"prologue-short-redirect"`` is currently only supported on
        x86-64.
+     * ``"ms-hotpatch"`` - This style of patchable function is similar to
+       ``"prologue-short-redirect"``, but it also imposes several additional
+       guarantees to support the style of hotpatching used on Windows.  On
+       32-bit x86, the first instruction will be a ``mov %edi, %edi``
+       instruction; this is frequently used as a magic value indicating a
+       hotpatchable function.  On other architectures, however, the first
+       instruction can be anything allowed in a Windows-style prologue;
+       this is because all functions on the non-i386 architectures Windows
+       supports are assumed to be hotpatchable.  Additionally, when not
+       targeting a Visual C++-style toolchain, patch space will be provided
+       prior to the function's entry point of an architecturally specific
+       size.  These sizes are compatible with GCC: on 32-bit x86, the patch
+       space is 64 bytes long; on x86-64, it is 128 bytes long.  The patch
+       space is not provided for MSVC toolchains because the
+       `/FUNCTIONPADMIN <https://msdn.microsoft.com/en-us/library/ms173524.aspx>`_
+       option, which provides this space, is expected to be used there.
+
+       ``"ms-hotpatch"`` is currently only supported on x86 and x86-64.
 
     This attribute by itself does not imply restrictions on
     inter-procedural optimizations.  All of the semantic effects the
diff --git a/llvm/include/llvm/Analysis/TargetTransformInfo.h b/llvm/include/llvm/Analysis/TargetTransformInfo.h
index 214f53e3dc7..f8fc25bd459 100644
--- a/llvm/include/llvm/Analysis/TargetTransformInfo.h
+++ b/llvm/include/llvm/Analysis/TargetTransformInfo.h
@@ -23,6 +23,8 @@
 #define LLVM_ANALYSIS_TARGETTRANSFORMINFO_H
 
 #include "llvm/ADT/Optional.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/Operator.h"
@@ -295,6 +297,18 @@ public:
   /// target-independent defaults.
   void getUnrollingPreferences(Loop *L, UnrollingPreferences &UP) const;
 
+  /// \brief Emit a patchable operation in the given basic block at the
+  /// given insertion point.
+  ///
+  /// Most of the time, this will be a straight-up \c TargetOpcode::PATCHABLE_OP
+  /// instruction, which will be lowered by the target to a no-op that can
+  /// be safely replaced with a short jump. However, some targets under certain
+  /// conditions can have peculiar requirements for this instruction; these
+  /// targets can provide their own implementation of this to emit the correct
+  /// instruction.
+  void emitPatchableOp(StringRef PatchType, MachineBasicBlock &MBB,
+                       MachineBasicBlock::iterator &MBBI) const;
+
   /// @}
 
   /// \name Scalar Target Information
@@ -647,6 +661,9 @@ public:
   virtual bool isSourceOfDivergence(const Value *V) = 0;
   virtual bool isLoweredToCall(const Function *F) = 0;
   virtual void getUnrollingPreferences(Loop *L, UnrollingPreferences &UP) = 0;
+  virtual void emitPatchableOp(StringRef Kind,
+                               MachineBasicBlock &MBB,
+                               MachineBasicBlock::iterator &MBBI) const = 0;
   virtual bool isLegalAddImmediate(int64_t Imm) = 0;
   virtual bool isLegalICmpImmediate(int64_t Imm) = 0;
   virtual bool isLegalAddressingMode(Type *Ty, GlobalValue *BaseGV,
@@ -792,6 +809,10 @@ public:
   void getUnrollingPreferences(Loop *L, UnrollingPreferences &UP) override {
     return Impl.getUnrollingPreferences(L, UP);
   }
+  void emitPatchableOp(StringRef Kind, MachineBasicBlock &MBB,
+                       MachineBasicBlock::iterator &MBBI) const override {
+    return Impl.emitPatchableOp(Kind, MBB, MBBI);
+  }
   bool isLegalAddImmediate(int64_t Imm) override {
     return Impl.isLegalAddImmediate(Imm);
   }
diff --git a/llvm/include/llvm/Analysis/TargetTransformInfoImpl.h b/llvm/include/llvm/Analysis/TargetTransformInfoImpl.h
index 118ca0241f5..80de388b958 100644
--- a/llvm/include/llvm/Analysis/TargetTransformInfoImpl.h
+++ b/llvm/include/llvm/Analysis/TargetTransformInfoImpl.h
@@ -23,6 +23,10 @@
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/Type.h"
 #include "llvm/Analysis/VectorUtils.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
+#include "llvm/CodeGen/MachineInstrBuilder.h"
+#include "llvm/Target/TargetInstrInfo.h"
+#include "llvm/Target/TargetSubtargetInfo.h"
 
 namespace llvm {
 
@@ -206,6 +210,20 @@ public:
 
   void getUnrollingPreferences(Loop *, TTI::UnrollingPreferences &) {}
 
+  void emitPatchableOp(StringRef, MachineBasicBlock &MBB,
+                       MachineBasicBlock::iterator &MBBI) const {
+    auto *TII = MBB.getParent()->getSubtarget().getInstrInfo();
+    auto MIB = BuildMI(MBB, MBBI, MBBI->getDebugLoc(),
+                       TII->get(TargetOpcode::PATCHABLE_OP))
+                   .addImm(2)
+                   .addImm(MBBI->getOpcode());
+
+    for (auto &MO : MBBI->operands())
+      MIB.addOperand(MO);
+
+    MBBI->eraseFromParent();
+  }
+
   bool isLegalAddImmediate(int64_t Imm) { return false; }
 
   bool isLegalICmpImmediate(int64_t Imm) { return false; }
diff --git a/llvm/lib/Analysis/TargetTransformInfo.cpp b/llvm/lib/Analysis/TargetTransformInfo.cpp
index c4aea355a3d..e131a5ccc4c 100644
--- a/llvm/lib/Analysis/TargetTransformInfo.cpp
+++ b/llvm/lib/Analysis/TargetTransformInfo.cpp
@@ -106,6 +106,12 @@ void TargetTransformInfo::getUnrollingPreferences(
   return TTIImpl->getUnrollingPreferences(L, UP);
 }
 
+void TargetTransformInfo::emitPatchableOp(
+    StringRef PatchType, MachineBasicBlock &MBB,
+    MachineBasicBlock::iterator &MBBI) const {
+  return TTIImpl->emitPatchableOp(PatchType, MBB, MBBI);
+}
+
 bool TargetTransformInfo::isLegalAddImmediate(int64_t Imm) const {
   return TTIImpl->isLegalAddImmediate(Imm);
 }
diff --git a/llvm/lib/CodeGen/PatchableFunction.cpp b/llvm/lib/CodeGen/PatchableFunction.cpp
index 32468c90b86..c94307df62b 100644
--- a/llvm/lib/CodeGen/PatchableFunction.cpp
+++ b/llvm/lib/CodeGen/PatchableFunction.cpp
@@ -13,12 +13,11 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/CodeGen/Passes.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/CodeGen/Analysis.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
-#include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/Target/TargetFrameLowering.h"
-#include "llvm/Target/TargetInstrInfo.h"
-#include "llvm/Target/TargetSubtargetInfo.h"
 
 using namespace llvm;
 
@@ -29,8 +28,9 @@ struct PatchableFunction : public MachineFunctionPass {
     initializePatchableFunctionPass(*PassRegistry::getPassRegistry());
   }
 
+  void getAnalysisUsage(AnalysisUsage &AU) const override;
   bool runOnMachineFunction(MachineFunction &F) override;
-   MachineFunctionProperties getRequiredProperties() const override {
+  MachineFunctionProperties getRequiredProperties() const override {
     return MachineFunctionProperties().set(
         MachineFunctionProperties::Property::AllVRegsAllocated);
   }
@@ -53,31 +53,29 @@ static bool doesNotGeneratecode(const MachineInstr &MI) {
   }
 }
 
+void PatchableFunction::getAnalysisUsage(AnalysisUsage &AU) const {
+  MachineFunctionPass::getAnalysisUsage(AU);
+  AU.addRequired<TargetTransformInfoWrapperPass>();
+}
+
 bool PatchableFunction::runOnMachineFunction(MachineFunction &MF) {
   if (!MF.getFunction()->hasFnAttribute("patchable-function"))
     return false;
 
-#ifndef NDEBUG
   Attribute PatchAttr = MF.getFunction()->getFnAttribute("patchable-function");
   StringRef PatchType = PatchAttr.getValueAsString();
-  assert(PatchType == "prologue-short-redirect" && "Only possibility today!");
-#endif
+  assert((PatchType == "prologue-short-redirect" ||
+          PatchType == "ms-hotpatch") && "Only possibilities today!");
 
   auto &FirstMBB = *MF.begin();
   MachineBasicBlock::iterator FirstActualI = FirstMBB.begin();
   for (; doesNotGeneratecode(*FirstActualI); ++FirstActualI)
     assert(FirstActualI != FirstMBB.end());
 
-  auto *TII = MF.getSubtarget().getInstrInfo();
-  auto MIB = BuildMI(FirstMBB, FirstActualI, FirstActualI->getDebugLoc(),
-                     TII->get(TargetOpcode::PATCHABLE_OP))
-                 .addImm(2)
-                 .addImm(FirstActualI->getOpcode());
-
-  for (auto &MO : FirstActualI->operands())
-    MIB.addOperand(MO);
+  const TargetTransformInfo &TTI =
+      getAnalysis<TargetTransformInfoWrapperPass>().getTTI(*MF.getFunction());
+  TTI.emitPatchableOp(PatchType, FirstMBB, FirstActualI);
 
-  FirstActualI->eraseFromParent();
   MF.ensureAlignment(4);
   return true;
 }
diff --git a/llvm/lib/Target/X86/X86AsmPrinter.cpp b/llvm/lib/Target/X86/X86AsmPrinter.cpp
index 67e51f1e919..76241dae069 100644
--- a/llvm/lib/Target/X86/X86AsmPrinter.cpp
+++ b/llvm/lib/Target/X86/X86AsmPrinter.cpp
@@ -76,6 +76,31 @@ bool X86AsmPrinter::runOnMachineFunction(MachineFunction &MF) {
   return false;
 }
 
+void X86AsmPrinter::EmitConstantPool() {
+  if (MF) {
+    // If an MS hotpatch function, we need to ensure 64 (32-bit) or 128 (64-bit)
+    // bytes of padding precede the label.  This is the scratch space used
+    // by the hotpatching mechanism to insert the patch code.  The movl %edi,
+    // %edi instruction emitted as the very first instruction of a hotpatch
+    // function is usually overwritten with a short jump instruction when the
+    // patch is installed, so it will jump directly into this space.  (But
+    // don't add the space when targeting MSVC.  There, the /FUNCTIONPADMIN
+    // option to link.exe is expected to be used.)
+    const Function *Fn = MF->getFunction();
+    if (!Subtarget->isTargetKnownWindowsMSVC() &&
+        Fn->hasFnAttribute("patchable-function") &&
+        Fn->getFnAttribute("patchable-function").getValueAsString() ==
+            "ms-hotpatch") {
+      // Emit INT3 instructions instead of NOPs. If a patch runs off the end,
+      // best to let the patcher know with a crash/debug break than to silently
+      // continue, only to run into the jump back into the patch.
+      OutStreamer->emitFill(Subtarget->is64Bit() ? 128 : 64, 0xcc);
+    }
+  }
+
+  AsmPrinter::EmitConstantPool();
+}
+
 /// printSymbolOperand - Print a raw symbol reference operand.  This handles
 /// jump tables, constant pools, global address and external symbols, all of
 /// which print to a label with various suffixes for relocation types etc.
diff --git a/llvm/lib/Target/X86/X86AsmPrinter.h b/llvm/lib/Target/X86/X86AsmPrinter.h
index dcb7b5a3466..7fc15cfd603 100644
--- a/llvm/lib/Target/X86/X86AsmPrinter.h
+++ b/llvm/lib/Target/X86/X86AsmPrinter.h
@@ -140,6 +140,8 @@ public:
     SMShadowTracker.emitShadowPadding(*OutStreamer, getSubtargetInfo());
   }
 
+  void EmitConstantPool() override;
+
   bool PrintAsmOperand(const MachineInstr *MI, unsigned OpNo,
                        unsigned AsmVariant, const char *ExtraCode,
                        raw_ostream &OS) override;
diff --git a/llvm/lib/Target/X86/X86FrameLowering.cpp b/llvm/lib/Target/X86/X86FrameLowering.cpp
index 850a67a30e0..a3f4d294421 100644
--- a/llvm/lib/Target/X86/X86FrameLowering.cpp
+++ b/llvm/lib/Target/X86/X86FrameLowering.cpp
@@ -928,6 +928,10 @@ void X86FrameLowering::emitPrologue(MachineFunction &MF,
   bool NeedsWinCFI = IsWin64Prologue && Fn->needsUnwindTableEntry();
   bool NeedsDwarfCFI =
       !IsWin64Prologue && (MMI.hasDebugInfo() || Fn->needsUnwindTableEntry());
+  bool IsMSHotpatch =
+      Fn->hasFnAttribute("patchable-function") &&
+      Fn->getFnAttribute("patchable-function").getValueAsString() ==
+          "ms-hotpatch";
   unsigned FramePtr = TRI->getFrameRegister(MF);
   const unsigned MachineFramePtr =
       STI.isTarget64BitILP32()
@@ -1069,7 +1073,9 @@ void X86FrameLowering::emitPrologue(MachineFunction &MF,
     if (!IsWin64Prologue && !IsFunclet) {
       // Update EBP with the new base value.
       BuildMI(MBB, MBBI, DL,
-              TII.get(Uses64BitFramePtr ? X86::MOV64rr : X86::MOV32rr),
+              TII.get(IsMSHotpatch ?
+                      (Uses64BitFramePtr ? X86::MOV64rr_REV : X86::MOV32rr_REV):
+                      (Uses64BitFramePtr ? X86::MOV64rr : X86::MOV32rr)),
               FramePtr)
           .addReg(StackPtr)
           .setMIFlag(MachineInstr::FrameSetup);
diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index c09d76d4845..2c6bd64dbb9 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -2576,10 +2576,13 @@ SDValue X86TargetLowering::LowerFormalArguments(
   X86MachineFunctionInfo *FuncInfo = MF.getInfo<X86MachineFunctionInfo>();
   const TargetFrameLowering &TFI = *Subtarget.getFrameLowering();
 
-  const Function *Fn = MF.getFunction();
-  if (Fn->hasExternalLinkage() &&
-      Subtarget.isTargetCygMing() &&
-      Fn->getName() == "main")
+  const Function* Fn = MF.getFunction();
+  if ((Fn->hasExternalLinkage() &&
+       Subtarget.isTargetCygMing() &&
+       Fn->getName() == "main") ||
+      (!Subtarget.is64Bit() && Fn->hasFnAttribute("patchable-function") &&
+       Fn->getFnAttribute("patchable-function").getValueAsString() ==
+           "ms-hotpatch"))
     FuncInfo->setForceFramePointer(true);
 
   MachineFrameInfo &MFI = MF.getFrameInfo();
diff --git a/llvm/lib/Target/X86/X86TargetTransformInfo.cpp b/llvm/lib/Target/X86/X86TargetTransformInfo.cpp
index 657a0451719..453afc00e9f 100644
--- a/llvm/lib/Target/X86/X86TargetTransformInfo.cpp
+++ b/llvm/lib/Target/X86/X86TargetTransformInfo.cpp
@@ -1681,3 +1681,18 @@ bool X86TTIImpl::areInlineCompatible(const Function *Caller,
   // correct.
   return (CallerBits & CalleeBits) == CalleeBits;
 }
+
+void X86TTIImpl::emitPatchableOp(StringRef PatchType,
+                                 MachineBasicBlock &MBB,
+                                 MachineBasicBlock::iterator &MBBI) const {
+  if (PatchType != "ms-hotpatch" || !ST->is32Bit()) {
+    BaseT::emitPatchableOp(PatchType, MBB, MBBI);
+    return;
+  }
+
+  auto &TII = *MBB.getParent()->getSubtarget().getInstrInfo();
+  BuildMI(MBB, MBBI, MBBI->getDebugLoc(),
+          TII.get(X86::MOV32rr_REV), X86::EDI)
+      .addReg(X86::EDI)
+      .setMIFlag(MachineInstr::FrameSetup);
+}
diff --git a/llvm/lib/Target/X86/X86TargetTransformInfo.h b/llvm/lib/Target/X86/X86TargetTransformInfo.h
index ab8046bb9fd..e8856c77296 100644
--- a/llvm/lib/Target/X86/X86TargetTransformInfo.h
+++ b/llvm/lib/Target/X86/X86TargetTransformInfo.h
@@ -50,6 +50,13 @@ public:
       : BaseT(std::move(static_cast<BaseT &>(Arg))), ST(std::move(Arg.ST)),
         TLI(std::move(Arg.TLI)) {}
 
+  /// \name Generic TTI Implementations
+  /// @{
+  void emitPatchableOp(StringRef PatchType, MachineBasicBlock &MBB,
+                       MachineBasicBlock::iterator &MBBI) const;
+
+  /// @}
+
   /// \name Scalar TTI Implementations
   /// @{
   TTI::PopcntSupportKind getPopcntSupport(unsigned TyWidth);
diff --git a/llvm/test/CodeGen/X86/ms-hotpatch-attr.ll b/llvm/test/CodeGen/X86/ms-hotpatch-attr.ll
new file mode 100644
index 00000000000..defbd1ddf73
--- /dev/null
+++ b/llvm/test/CodeGen/X86/ms-hotpatch-attr.ll
@@ -0,0 +1,27 @@
+; RUN: llc < %s -march=x86 -filetype=asm | FileCheck -check-prefix=CHECK-32 %s
+; RUN: llc < %s -march=x86-64 -filetype=asm | FileCheck -check-prefix=CHECK-64 %s
+; RUN: llc < %s -mtriple=i386-windows-msvc -filetype=asm | FileCheck -check-prefix=MSVC-32 %s
+; RUN: llc < %s -mtriple=x86_64-windows-msvc -filetype=asm | FileCheck -check-prefix=MSVC-64 %s
+
+; CHECK-32: .space 64,204
+; CHECK-32: .p2align 4, 0x90
+; CHECK-32-LABEL: foo:
+; CHECK-32: movl %edi, %edi
+; CHECK-32-NEXT: pushl %ebp
+; CHECK-32-NEXT: movl %esp, %ebp
+; CHECK-64: .space 128,204
+; CHECK-64: .p2align 4, 0x90
+; CHECK-64-LABEL: foo:
+; CHECK-64: xchgw %ax, %ax
+; MSVC-32-NOT: .space 64,204
+; MSVC-32-LABEL: _foo:
+; MSVC-32: movl %edi, %edi
+; MSVC-32-NEXT: pushl %ebp
+; MSVC-32-NEXT: movl %esp, %ebp
+; MSVC-64-NOT: .space 128,204
+; MSVC-64-LABEL: foo:
+; MSVC-64: xchgw %ax, %ax
+define void @foo() nounwind "patchable-function"="ms-hotpatch" {
+entry:
+  ret void
+}