[X86] Alternate implementation of D88194.

[craig.topper]

Summary

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

[xiangzhangllvm]

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

858

Can we replace a constant with a reg without checking the operation?
I am afraid it may cause ISEL failed.

[LuoYuanke]

This patch looks more common than D88194. Is it possible that generate endbr constant after PreprocessISelDAG()?

[spatel]

If this is sufficient, I definitely prefer this implementation - much smaller and easier to understand.

[craig.topper]

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

858

Cases that require constants should be using ISD::TargetConstant instead of ISD::Constant.

I fixed one case I found where it would break in 375849518db97096212f7f2b996b0d15f97be959. There might be others. I'll try to see if I can find a way to check the isel tables for other cases.

[RKSimon]

No objections from me and this looks better than D88194

@craig.topper Is there a better way that we can always check for intrinsics that should be using TargetConstant ?

[craig.topper]

Are we ok with this implementation?

I don't have a great way to protect intrinsics not using TargetConstant properly in the future, but I think we're ok at the moment based on what I see in the isel table.

[xiangzhangllvm]

In D89178#2352272, @craig.topper wrote:

Are we ok with this implementation?

I don't have a great way to protect intrinsics not using TargetConstant properly in the future, but I think we're ok at the moment based on what I see in the isel table.

LGTM, if problem in the future, I'll fix it. Tkx !!

[pengfei]

LGTM too. I think we can let it in. Thanks.

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

 //===- X86ISelDAGToDAG.cpp - A DAG pattern matching inst selector for X86 -===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines a DAG pattern matching instruction selector for X86,
 // converting from a legalized dag to a X86 dag.
 //
 //===----------------------------------------------------------------------===//
 
 #include "X86.h"
 #include "X86MachineFunctionInfo.h"
 #include "X86RegisterInfo.h"
 #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"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/IntrinsicsX86.h"
 #include "llvm/IR/Type.h"

 static cl::opt<bool> AndImmShrink("x86-and-imm-shrink", cl::init(true),
     cl::desc("Enable setting constant bits to reduce size of mask immediates"),
     cl::Hidden);
 
 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
 //===----------------------------------------------------------------------===//
 
 namespace {
   /// This corresponds to X86AddressMode, but uses SDValue's instead of register
   /// numbers for the leaves of the matched tree.
   struct X86ISelAddressMode {
if (Chain.getOperand(0).getNode() == Callee.getNode())
     return true;
   if (Chain.getOperand(0).getOpcode() == ISD::TokenFactor &&
       Callee.getValue(1).isOperandOf(Chain.getOperand(0).getNode()) &&
       Callee.getValue(1).hasOneUse())
     return true;
   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)) {
       SDValue Res = CurDAG->getNode(ISD::AND, SDLoc(N), N->getValueType(0),
                                     N->getOperand(0), N->getOperand(1));
       --I;
       CurDAG->ReplaceAllUsesOfValueWith(SDValue(N, 0), Res);
       ++I;

llvm/lib/Target/X86/X86IndirectBranchTracking.cpp

 //===---- X86IndirectBranchTracking.cpp - Enables CET IBT mechanism -------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //

 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 
 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."));
 
 STATISTIC(NumEndBranchAdded, "Number of ENDBR instructions added");
 
 namespace {
 class X86IndirectBranchTrackingPass : public MachineFunctionPass {
 public:

llvm/test/CodeGen/X86/cet_endbr_imm_enhance.ll

 ; 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
 ; 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:

 ; ~0x000123F32E0F1EFA == -321002333478651 (0XFFFEDC0CD1F0E105)
 
 ; test for MOV64ri
 define dso_local i64 @foo(i64* %azx) #0 {
 ; CHECK-LABEL: foo:
 ; 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
 ; CHECK-NEXT:    retq
 entry:
   %azx.addr = alloca i64*, align 8
   store i64* %azx, i64** %azx.addr, align 8
   %0 = load i64*, i64** %azx.addr, align 8
   %1 = load i64, i64* %0, align 8
   %and = and i64 %1, 321002333478650
   %2 = load i64*, i64** %azx.addr, align 8
   store i64 %and, i64* %2, align 8
   %3 = load i64*, i64** %azx.addr, align 8
   %4 = load i64, i64* %3, align 8
   ret i64 %4
 }
 
 @bzx = dso_local local_unnamed_addr global i32 -217112837, align 4
 
 ; test for AND32ri
 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
   %mul = shl nsw i32 %0, 1
   %and = and i32 %mul, -217112838
   ret i32 %and
 }
 
 
 @czx = dso_local global i32 -217112837, align 4
 
 ; test for AND32mi
 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:
   %0 = load i32, i32* @czx, align 4
   %and = and i32 %0, -217112838
   store i32 %and, i32* @czx, align 4
   ret i32* @czx
 }
 
 ; test for MOV32mi
 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
   store i32 -217112838, i32* %dzx, align 4
   %0 = load i32, i32* %dzx, align 4
   ret i32 %0
 }
 
 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
   store i64 4077854458, i64* %ezx, align 8
   %0 = load i64, i64* %ezx, align 8
   ret i64 %0
 }