[X86] Alternate implementation of D88194.

This uses PreprocessISelDAG to replace the constant before
instruction selection instead of matching opcodes after.

Reviewed By: pengfei

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

diff --git a/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp b/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp
index 97c0615..9a16dc1 100644 (file)
--- a/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp
+++ b/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp
@@ -17,6 +17,7 @@
 #include "X86Subtarget.h"
 #include "X86TargetMachine.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/SelectionDAGISel.h"
 #include "llvm/Config/llvm-config.h"
 #include "llvm/IR/ConstantRange.h"
@@ -44,6 +45,8 @@ static cl::opt<bool> EnablePromoteAnyextLoad(
     "x86-promote-anyext-load", cl::init(true),
     cl::desc("Enable promoting aligned anyext load to wider load"), cl::Hidden);
 
+extern cl::opt<bool> IndirectBranchTracking;
+
 //===----------------------------------------------------------------------===//
 //                      Pattern Matcher Implementation
 //===----------------------------------------------------------------------===//
@@ -797,12 +800,69 @@ static bool isCalleeLoad(SDValue Callee, SDValue &Chain, bool HasCallSeq) {
   return false;
 }
 
+static bool isEndbrImm64(uint64_t Imm) {
+// There may be some other prefix bytes between 0xF3 and 0x0F1EFA.
+// i.g: 0xF3660F1EFA, 0xF3670F1EFA
+  if ((Imm & 0x00FFFFFF) != 0x0F1EFA)
+    return false;
+
+  uint8_t OptionalPrefixBytes [] = {0x26, 0x2e, 0x36, 0x3e, 0x64,
+                                    0x65, 0x66, 0x67, 0xf0, 0xf2};
+  int i = 24; // 24bit 0x0F1EFA has matched
+  while (i < 64) {
+    uint8_t Byte = (Imm >> i) & 0xFF;
+    if (Byte == 0xF3)
+      return true;
+    if (!llvm::is_contained(OptionalPrefixBytes, Byte))
+      return false;
+    i += 8;
+  }
+
+  return false;
+}
+
 void X86DAGToDAGISel::PreprocessISelDAG() {
   bool MadeChange = false;
   for (SelectionDAG::allnodes_iterator I = CurDAG->allnodes_begin(),
        E = CurDAG->allnodes_end(); I != E; ) {
     SDNode *N = &*I++; // Preincrement iterator to avoid invalidation issues.
 
+    // This is for CET enhancement.
+    //
+    // ENDBR32 and ENDBR64 have specific opcodes:
+    // ENDBR32: F3 0F 1E FB
+    // ENDBR64: F3 0F 1E FA
+    // And we want that attackers won’t find unintended ENDBR32/64
+    // opcode matches in the binary
+    // Here’s an example:
+    // If the compiler had to generate asm for the following code:
+    // a = 0xF30F1EFA
+    // it could, for example, generate:
+    // mov 0xF30F1EFA, dword ptr[a]
+    // In such a case, the binary would include a gadget that starts
+    // with a fake ENDBR64 opcode. Therefore, we split such generation
+    // into multiple operations, let it not shows in the binary
+    if (N->getOpcode() == ISD::Constant) {
+      MVT VT = N->getSimpleValueType(0);
+      int64_t Imm = cast<ConstantSDNode>(N)->getSExtValue();
+      int32_t EndbrImm = Subtarget->is64Bit() ? 0xF30F1EFA : 0xF30F1EFB;
+      if (Imm == EndbrImm || isEndbrImm64(Imm)) {
+        // Check that the cf-protection-branch is enabled.
+        Metadata *CFProtectionBranch =
+          MF->getMMI().getModule()->getModuleFlag("cf-protection-branch");
+        if (CFProtectionBranch || IndirectBranchTracking) {
+          SDLoc dl(N);
+          SDValue Complement = CurDAG->getConstant(~Imm, dl, VT, false, true);
+          Complement = CurDAG->getNOT(dl, Complement, VT);
+          --I;
+          CurDAG->ReplaceAllUsesOfValueWith(SDValue(N, 0), Complement);
+          ++I;
+          MadeChange = true;
+          continue;
+        }
+      }
+    }
+
     // If this is a target specific AND node with no flag usages, turn it back
     // into ISD::AND to enable test instruction matching.
     if (N->getOpcode() == X86ISD::AND && !N->hasAnyUseOfValue(1)) {

diff --git a/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp b/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp
index 1628f85..85410c5 100644 (file)
--- a/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp
+++ b/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp
@@ -28,7 +28,7 @@ using namespace llvm;
 
 #define DEBUG_TYPE "x86-indirect-branch-tracking"
 
-static cl::opt<bool> IndirectBranchTracking(
+cl::opt<bool> IndirectBranchTracking(
     "x86-indirect-branch-tracking", cl::init(false), cl::Hidden,
     cl::desc("Enable X86 indirect branch tracking pass."));
 

diff --git a/llvm/test/CodeGen/X86/cet_endbr_imm_enhance.ll b/llvm/test/CodeGen/X86/cet_endbr_imm_enhance.ll
index 8d0f9c7..1c98738 100644 (file)
--- a/llvm/test/CodeGen/X86/cet_endbr_imm_enhance.ll
+++ b/llvm/test/CodeGen/X86/cet_endbr_imm_enhance.ll
@@ -1,7 +1,7 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -O2 -mtriple=x86_64-unknown-unknown -x86-indirect-branch-tracking | FileCheck %s
 
-; TBD: This test is for CET enhancement, we should replace the endbr imm.
+; This test is for CET enhancement.
 ;
 ; ENDBR32 and ENDBR64 have specific opcodes:
 ; ENDBR32: F3 0F 1E FB
@@ -27,7 +27,8 @@ define dso_local i64 @foo(i64* %azx) #0 {
 ; CHECK:       # %bb.0: # %entry
 ; CHECK-NEXT:    endbr64
 ; CHECK-NEXT:    movq %rdi, -{{[0-9]+}}(%rsp)
-; CHECK-NEXT:    movabsq $321002333478650, %rax # imm = 0x123F32E0F1EFA
+; CHECK-NEXT:    movabsq $-321002333478651, %rax # imm = 0xFFFEDC0CD1F0E105
+; CHECK-NEXT:    notq %rax
 ; CHECK-NEXT:    andq %rax, (%rdi)
 ; CHECK-NEXT:    movq -{{[0-9]+}}(%rsp), %rax
 ; CHECK-NEXT:    movq (%rax), %rax
@@ -52,9 +53,11 @@ define dso_local i32 @foo2() local_unnamed_addr #0 {
 ; CHECK-LABEL: foo2:
 ; CHECK:       # %bb.0: # %entry
 ; CHECK-NEXT:    endbr64
-; CHECK-NEXT:    movl {{.*}}(%rip), %eax
-; CHECK-NEXT:    addl %eax, %eax
-; CHECK-NEXT:    andl $-217112838, %eax # imm = 0xF30F1EFA
+; CHECK-NEXT:    movl {{.*}}(%rip), %ecx
+; CHECK-NEXT:    addl %ecx, %ecx
+; CHECK-NEXT:    movl $217112837, %eax # imm = 0xCF0E105
+; CHECK-NEXT:    notl %eax
+; CHECK-NEXT:    andl %ecx, %eax
 ; CHECK-NEXT:    retq
 entry:
   %0 = load i32, i32* @bzx, align 4
@@ -71,7 +74,9 @@ define dso_local nonnull i32* @foo3() local_unnamed_addr #0 {
 ; CHECK-LABEL: foo3:
 ; CHECK:       # %bb.0: # %entry
 ; CHECK-NEXT:    endbr64
-; CHECK-NEXT:    andl $-217112838, {{.*}}(%rip) # imm = 0xF30F1EFA
+; CHECK-NEXT:    movl $217112837, %eax # imm = 0xCF0E105
+; CHECK-NEXT:    notl %eax
+; CHECK-NEXT:    andl %eax, {{.*}}(%rip)
 ; CHECK-NEXT:    movl $czx, %eax
 ; CHECK-NEXT:    retq
 entry:
@@ -86,8 +91,9 @@ define dso_local i32 @foo4() #0 {
 ; CHECK-LABEL: foo4:
 ; CHECK:       # %bb.0: # %entry
 ; CHECK-NEXT:    endbr64
-; CHECK-NEXT:    movl $-217112838, -{{[0-9]+}}(%rsp) # imm = 0xF30F1EFA
-; CHECK-NEXT:    movl $-217112838, %eax # imm = 0xF30F1EFA
+; CHECK-NEXT:    movl $217112837, %eax # imm = 0xCF0E105
+; CHECK-NEXT:    notl %eax
+; CHECK-NEXT:    movl %eax, -{{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    retq
 entry:
   %dzx = alloca i32, align 4
@@ -100,9 +106,9 @@ define dso_local i64 @foo5() #0 {
 ; CHECK-LABEL: foo5:
 ; CHECK:       # %bb.0: # %entry
 ; CHECK-NEXT:    endbr64
-; CHECK-NEXT:    movl $4077854458, %eax # imm = 0xF30F1EFA
+; CHECK-NEXT:    movabsq $-4077854459, %rax # imm = 0xFFFFFFFF0CF0E105
+; CHECK-NEXT:    notq %rax
 ; CHECK-NEXT:    movq %rax, -{{[0-9]+}}(%rsp)
-; CHECK-NEXT:    movl $4077854458, %eax # imm = 0xF30F1EFA
 ; CHECK-NEXT:    retq
 entry:
   %ezx = alloca i64, align 8