[X86] CET endbr enhance

[xiangzhangllvm]

Summary

This patch is for CET (Control-flow Enforcement Technology) 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.

The goal of this patch is not to 100% remove the unintended ENDBR-IMM.
Theoretically, it can occurrence in address info, and even between 2 instructions.
In fact, All the probability of its occurrence is very small.
The idea of this patch tend to “Greatly reduce the probability of ENDBR-IMM occurrence” by handling the most comment instructions with imm32/64.

[craig.topper]

Should we only be doing this when CET protections are enabled?

[craig.topper]

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

1344

This can't be an ArrayRef. The RHS is a temporary initializer_list that's lifetime ends at the end of the statement.

Use uint8_t Prefixes[] = { ... }

[MaskRay]

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

1338

Newer function should stick with the coding standard and use lowerCase function names.

[xiangzhangllvm]

Done, TKS !!

[craig.topper]

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

1352

You don't need an ArrayRef. you can just use std::begin() and std::end() on Bytes

[xiangzhangllvm]

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

1352

Done, thank you !!

[craig.topper]

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

1352

We shouldn't hard code the size. Its not friendly to change multiple places if the array size changes which is why I suggested std::begin/std::end

But looking again I forgot we have llvm::is_contained in STLExtras.h which would be even better here.

[xiangzhangllvm]

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

1352

Yes, It wraps std::find, that's better, thank you !!

[xiangzhangllvm]

Hello Craig, could you help accept, TKS!

[craig.topper]

@spatel or @RKSimon i'd appreciate if you could look at this as well.

@xiangzhangllvm do we care about -O0 builds for this? Those won't always go through SelectionDAG. Only when the "fast" instruction selector find something it can't handle.

[xiangzhangllvm]

I tend not do it for -O0, most projects/libs not use -O0 to build.
This is not a bug, we just refine the CET.

[spatel]

Warning: I had no idea what "CET" means before seeing this patch, so feel free to discount my feedback.
The patch title and description are not clear to me.
Is the goal of this patch to change *the compiler* binary itself? (rather than an application binary)
It would be good to pre-commit a minimal test at least so we can see exactly what asm *change* is introduced by this patch.

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

209

Formatting: function name should be 'lowerCamelCase'

1344

Bytes -> OptionalPrefixBytes ?

1374

The function name does not describe what it actually does.
Would "obscureEndbrOpcodeImmediate" be more accurate?

1386–1387

This code structure with fallthrough is difficult to follow. Create a separate static helper function that just returns the constant operand index for a given opcode?

[xiangzhangllvm]

@spatel thank you very much for your review! and very sorry for the later update!
CET stands for Intel Control-flow Enforcement Technology, which was designed by H.J. and has already implemented in llvm and gcc. Now it can be googled out.
CET adds an Indirect Branch Tracking capability (by inserting Endbr instruction at the destination of the indirect branch) to provide software the ability to restrict COP/JOP attacks.

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

1386–1387

I really thought using a function and refined here before commit it to phabricator, try to simplify the code I used the LLVM_FALLTHROUGH.
Here not just need 'return' constant operand index, please refer line 1412-1414.

[spatel]

In D88194#2318848, @xiangzhangllvm wrote:

CET adds an Indirect Branch Tracking capability (by inserting Endbr instruction at the destination of the indirect branch) to provide software the ability to restrict COP/JOP attacks.

Ok, I understand that much. But it is still not clear to me why this patch is re-using the "-x86-indirect-branch-tracking" option. Is this obscuring of the opcode not independent of producing the instruction? I would again suggest that we pre-commit some tests, so we can see only the diffs in the codegen (and if inserting "endb64" itself is a requirement to obscure the opcode, please add comments in the test file to explain that).

[craig.topper]

In D88194#2321872, @spatel wrote:

In D88194#2318848, @xiangzhangllvm wrote:

CET adds an Indirect Branch Tracking capability (by inserting Endbr instruction at the destination of the indirect branch) to provide software the ability to restrict COP/JOP attacks.

Ok, I understand that much. But it is still not clear to me why this patch is re-using the "-x86-indirect-branch-tracking" option. Is this obscuring of the opcode not independent of producing the instruction? I would again suggest that we pre-commit some tests, so we can see only the diffs in the codegen (and if inserting "endb64" itself is a requirement to obscure the opcode, please add comments in the test file to explain that).

I suggested adding that the -x86-indirect-branch-tracking check because it seemed like we'd only want to do this in a binary that was expected to contain endbr opcodes. Putting things that could be confused for endbr opcodes in a binary that doesn't enable CET should be harmless.

[spatel]

In D88194#2322158, @craig.topper wrote:

In D88194#2321872, @spatel wrote:

In D88194#2318848, @xiangzhangllvm wrote:

CET adds an Indirect Branch Tracking capability (by inserting Endbr instruction at the destination of the indirect branch) to provide software the ability to restrict COP/JOP attacks.

Ok, I understand that much. But it is still not clear to me why this patch is re-using the "-x86-indirect-branch-tracking" option. Is this obscuring of the opcode not independent of producing the instruction? I would again suggest that we pre-commit some tests, so we can see only the diffs in the codegen (and if inserting "endb64" itself is a requirement to obscure the opcode, please add comments in the test file to explain that).

I suggested adding that the -x86-indirect-branch-tracking check because it seemed like we'd only want to do this in a binary that was expected to contain endbr opcodes. Putting things that could be confused for endbr opcodes in a binary that doesn't enable CET should be harmless.

Is this patch designed to change the compiler itself? Ie, we would enable this obfuscation mechanism while building clang, then use the newly built clang with the existing cl::opt flag to build an app that contains endbr instructions that are constructed in the convoluted way? We don't care if the compiler is littered with endbr itself?

[craig.topper]

In D88194#2322322, @spatel wrote:

In D88194#2322158, @craig.topper wrote:

In D88194#2321872, @spatel wrote:

In D88194#2318848, @xiangzhangllvm wrote:

CET adds an Indirect Branch Tracking capability (by inserting Endbr instruction at the destination of the indirect branch) to provide software the ability to restrict COP/JOP attacks.

Ok, I understand that much. But it is still not clear to me why this patch is re-using the "-x86-indirect-branch-tracking" option. Is this obscuring of the opcode not independent of producing the instruction? I would again suggest that we pre-commit some tests, so we can see only the diffs in the codegen (and if inserting "endb64" itself is a requirement to obscure the opcode, please add comments in the test file to explain that).

I suggested adding that the -x86-indirect-branch-tracking check because it seemed like we'd only want to do this in a binary that was expected to contain endbr opcodes. Putting things that could be confused for endbr opcodes in a binary that doesn't enable CET should be harmless.

Is this patch designed to change the compiler itself? Ie, we would enable this obfuscation mechanism while building clang, then use the newly built clang with the existing cl::opt flag to build an app that contains endbr instructions that are constructed in the convoluted way? We don't care if the compiler is littered with endbr itself?

No its not designed to change the compiler itself. It is just supposed to make it less likely that a binary containing real endbr opcodes will have bytes that look like endbr opcodes.

[spatel]

In D88194#2322331, @craig.topper wrote:

In D88194#2322322, @spatel wrote:

In D88194#2322158, @craig.topper wrote:

In D88194#2321872, @spatel wrote:

In D88194#2318848, @xiangzhangllvm wrote:

CET adds an Indirect Branch Tracking capability (by inserting Endbr instruction at the destination of the indirect branch) to provide software the ability to restrict COP/JOP attacks.

Ok, I understand that much. But it is still not clear to me why this patch is re-using the "-x86-indirect-branch-tracking" option. Is this obscuring of the opcode not independent of producing the instruction? I would again suggest that we pre-commit some tests, so we can see only the diffs in the codegen (and if inserting "endb64" itself is a requirement to obscure the opcode, please add comments in the test file to explain that).

I suggested adding that the -x86-indirect-branch-tracking check because it seemed like we'd only want to do this in a binary that was expected to contain endbr opcodes. Putting things that could be confused for endbr opcodes in a binary that doesn't enable CET should be harmless.

Is this patch designed to change the compiler itself? Ie, we would enable this obfuscation mechanism while building clang, then use the newly built clang with the existing cl::opt flag to build an app that contains endbr instructions that are constructed in the convoluted way? We don't care if the compiler is littered with endbr itself?

No its not designed to change the compiler itself. It is just supposed to make it less likely that a binary containing real endbr opcodes will have bytes that look like endbr opcodes.

Ok, I don't have any other questions/comments (other than it would still be helpful to reduce/pre-commit the tests).

[xiangzhangllvm]

Hello @spatel, Sorrry, I don't much understand the "reduce/pre-commit tests",
did you mean the test here "cet_endbr_imm_enhance.ll" should be merge into other existing test ?

[xiangzhangllvm]

In D88194#2322366, @spatel wrote:

Ok, I don't have any other questions/comments (other than it would still be helpful to reduce/pre-commit the tests).

Done, then update the tests.

[craig.topper]

llvm/test/CodeGen/X86/cet_endbr_imm_enhance.ll

2

Don't we need to test the endbr32 case?

[spatel]

In D88194#2323095, @xiangzhangllvm wrote:

In D88194#2322366, @spatel wrote:

Ok, I don't have any other questions/comments (other than it would still be helpful to reduce/pre-commit the tests).

Done, then update the tests.

Thanks! That makes it easier to see how things are changing with this patch.
In addition to @craig.topper 's question about endbr32, shouldn't we have a test that checks for multiple and different prefix bytes? If I am seeing correctly, there is only 1 test checking for a single 0x2E optional prefix.

[spatel]

In D88194#2323400, @spatel wrote:

In D88194#2323095, @xiangzhangllvm wrote:

In D88194#2322366, @spatel wrote:

Ok, I don't have any other questions/comments (other than it would still be helpful to reduce/pre-commit the tests).

Done, then update the tests.

Thanks! That makes it easier to see how things are changing with this patch.
In addition to @craig.topper 's question about endbr32, shouldn't we have a test that checks for multiple and different prefix bytes? If I am seeing correctly, there is only 1 test checking for a single 0x2E optional prefix.

Thinking about this a bit more: if the attacker can recognize multiple different optional prefix bytes, then is there much value in complementing those bytes? They are already searching for a match from some dictionary of byte strings, so inverting the bits just means they need to increase their search by 2x?

[MaskRay]

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

20

Which function uses llvm/CodeGen/MachineModuleInfo.h ?

1341

Please don't mix lower-case and upper-case hexadecimal literals.

1344

constexpr uint8_t OptionalPrefixBytes[]

Are these prefix bytes tested?

1347

You can use Imm >>= i; while (Imm != 0) { ... Imm >>= 8; }

1375

Add const if appropriate.
ditto below

1383

switch (Opc)

Please clang-format the patch (git diff -U0 --no-color 'HEAD^' | llvm-project/clang/tools/clang-format/clang-format-diff.py -i -p1)

1470

There is a space before //

[craig.topper]

In D88194#2323402, @spatel wrote:

In D88194#2323400, @spatel wrote:

In D88194#2323095, @xiangzhangllvm wrote:

In D88194#2322366, @spatel wrote:

Ok, I don't have any other questions/comments (other than it would still be helpful to reduce/pre-commit the tests).

Done, then update the tests.

Thanks! That makes it easier to see how things are changing with this patch.
In addition to @craig.topper 's question about endbr32, shouldn't we have a test that checks for multiple and different prefix bytes? If I am seeing correctly, there is only 1 test checking for a single 0x2E optional prefix.

Thinking about this a bit more: if the attacker can recognize multiple different optional prefix bytes, then is there much value in complementing those bytes? They are already searching for a match from some dictionary of byte strings, so inverting the bits just means they need to increase their search by 2x?

I'm not sure I follow. If the bytes are inverted in the binary then they are no longer useful to the attacker. The goal here is to make sure that if an attacker gains control of the register or memory location used by an indirect call or jump that they can't jump to an immediate in the binary that matches the encoding for ENDBR. As that would convince the CET protection that this a place that an indirect jump was expected to be able to jump to.

[xiangzhangllvm]

Thanks for your reviews! Please let me suspend this patch first.

duo to Craig's new patch https://reviews.llvm.org/D89178.
(Reason: Every instruction contain I32/64 targetconstant operand has a corresponding reg version.
So, we can replace the I32/64 targetconstant without checking the operation before ISel.)

[craig.topper]

Abandon this?

[xiangzhangllvm]

OK

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

 //===----------------------------------------------------------------------===//
 
 #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 {
public:
     void PostprocessISelDAG() override;
 
 // Include the pieces autogenerated from the target description.
 #include "X86GenDAGISel.inc"
 
   private:
     void Select(SDNode *N) override;
 
+    bool EndbrImmAntiAttack(SDNode *N);
     bool foldOffsetIntoAddress(uint64_t Offset, X86ISelAddressMode &AM);
     bool matchLoadInAddress(LoadSDNode *N, X86ISelAddressMode &AM);
     bool matchWrapper(SDValue N, X86ISelAddressMode &AM);
     bool matchAddress(SDValue N, X86ISelAddressMode &AM);
     bool matchVectorAddress(SDValue N, X86ISelAddressMode &AM);
     bool matchAdd(SDValue &N, X86ISelAddressMode &AM, unsigned Depth);
     bool matchAddressRecursively(SDValue N, X86ISelAddressMode &AM,
                                  unsigned Depth);
bool X86DAGToDAGISel::tryOptimizeRem8Extend(SDNode *N) {
   } else {
     // Ok we can drop this extend and just use the original extend.
     ReplaceUses(N, N00.getNode());
   }
 
   return true;
 }
 
+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 Bytes[10]= {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 (std::find(Bytes, Bytes + 10, Byte) == Bytes + 10)
+      return false;
+    i += 8;
+  }
+
+  return false;
+}
+
+// This function 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.
+bool X86DAGToDAGISel::EndbrImmAntiAttack(SDNode *N) {
+  unsigned Opc = N->getMachineOpcode();
+  MachineSDNode *N0 = nullptr;
+  MachineSDNode *N1 = nullptr;
+  int Idx = -1;
+  unsigned MOVOpc = X86::MOV32ri64; // Mov i64i32imm to 64-bit reg.
+  unsigned NOTOpc = X86::NOT64r;
+  EVT VT = MVT::i64;
+  SDLoc dl(N);
+  switch(Opc) {
+    default: break;
+    case X86::MOV32ri: // *32ri
+      Idx = 0;
+      LLVM_FALLTHROUGH;
+    case X86::ADC32ri:
+    case X86::ADD32ri:
+    case X86::AND32ri:
+    case X86::CMP32ri:
+    case X86::OR32ri:
+    case X86::SBB32ri:
+    case X86::SUB32ri:
+    case X86::TEST32ri:
+    case X86::XOR32ri:
+      if (Idx < 0)
+        Idx = 1;
+      LLVM_FALLTHROUGH;
+    case X86::MOV32mi: // *32mi
+    case X86::ADC32mi:
+    case X86::ADD32mi:
+    case X86::AND32mi:
+    case X86::CMP32mi:
+    case X86::OR32mi:
+    case X86::SBB32mi:
+    case X86::SUB32mi:
+    case X86::TEST32mi:
+    case X86::XOR32mi:
+      if (Idx < 0)
+        Idx = 5;
+      MOVOpc = X86::MOV32ri;
+      NOTOpc = X86::NOT32r;
+      VT = MVT::i32;
+      LLVM_FALLTHROUGH;
+    case X86::MOV64ri32: // *64ri32
+      if (Idx < 0)
+        Idx = 0;
+      LLVM_FALLTHROUGH;
+    case X86::ADC64ri32:
+    case X86::ADD64ri32:
+    case X86::AND64ri32:
+    case X86::CMP64ri32:
+    case X86::OR64ri32:
+    case X86::SBB64ri32:
+    case X86::SUB64ri32:
+    case X86::TEST64ri32:
+    case X86::XOR64ri32:
+      if (Idx < 0)
+        Idx = 1;
+      LLVM_FALLTHROUGH;
+    case X86::MOV64mi32: // *64mi32
+    case X86::ADC64mi32:
+    case X86::ADD64mi32:
+    case X86::AND64mi32:
+    case X86::CMP64mi32:
+    case X86::OR64mi32:
+    case X86::SBB64mi32:
+    case X86::SUB64mi32:
+    case X86::TEST64mi32:
+    case X86::XOR64mi32: {
+      if (Idx < 0)
+        Idx = 5;
+      assert(Idx >= 0);
+      if (!isa<llvm::ConstantSDNode>(N->getOperand(Idx)))
+        return false;
+      uint32_t Imm = N->getConstantOperandVal(Idx);
+      uint32_t EndbrImm = Subtarget->is64Bit() ? 0xF30F1EFA : 0xF30F1EFB;
+      if (Imm != EndbrImm)
+        return false;
+
+      SDValue Complement = MOVOpc == X86::MOV32ri ?
+                           getI32Imm(~EndbrImm, dl) :
+                           getI64Imm((uint64_t)(~EndbrImm), dl);
+
+      // ~0xF30F1EFA = 0x0CF0E105
+      // ~0xF30F1EFB = 0x0CF0E104
+      // 1st Move the complement of endbr-imm into a reg.
+      N0 = CurDAG->getMachineNode(MOVOpc, dl, VT, Complement);
+
+      // ~~0xF30F1EFA  = 0xF30F1EFA
+      // ~~0xF30F1EFB  = 0xF30F1EFB
+      // 2nd The complement of endbr-imm's complement equal to the old value.
+      // The NOT operation has no-effect to status flags.
+      N0 = CurDAG->getMachineNode(NOTOpc, dl, VT, SDValue(N0, 0));
+
+      unsigned NewOpc;
+      switch (Opc) {
+      default: llvm_unreachable("Unexpected opcode!");
+      case X86::ADC32ri: NewOpc = X86::ADC32rr; break;// *32ri
+      case X86::ADD32ri: NewOpc = X86::ADD32rr; break;
+      case X86::AND32ri: NewOpc = X86::AND32rr; break;
+      case X86::CMP32ri: NewOpc = X86::CMP32rr; break;
+      case X86::MOV32ri: NewOpc = X86::MOV32rr; break;
+      case X86::OR32ri:  NewOpc = X86::OR32rr;  break;
+      case X86::SBB32ri: NewOpc = X86::SBB32rr; break;
+      case X86::SUB32ri: NewOpc = X86::SUB32rr; break;
+      case X86::TEST32ri: NewOpc = X86::TEST32rr; break;
+      case X86::XOR32ri: NewOpc = X86::XOR32rr; break;
+      case X86::ADC32mi: NewOpc = X86::ADC32mr; break;// *32mi
+      case X86::ADD32mi: NewOpc = X86::ADD32mr; break;
+      case X86::AND32mi: NewOpc = X86::AND32mr; break;
+      case X86::CMP32mi: NewOpc = X86::CMP32mr; break;
+      case X86::MOV32mi: NewOpc = X86::MOV32mr; break;
+      case X86::OR32mi:  NewOpc = X86::OR32mr;  break;
+      case X86::SBB32mi: NewOpc = X86::SBB32mr; break;
+      case X86::SUB32mi: NewOpc = X86::SUB32mr; break;
+      case X86::TEST32mi: NewOpc = X86::TEST32mr; break;
+      case X86::XOR32mi: NewOpc = X86::XOR32mr; break;
+      case X86::ADC64ri32: NewOpc = X86::ADC64rr; break;// *64ri32
+      case X86::ADD64ri32: NewOpc = X86::ADD64rr; break;
+      case X86::AND64ri32: NewOpc = X86::AND64rr; break;
+      case X86::CMP64ri32: NewOpc = X86::CMP64rr; break;
+      case X86::MOV64ri32: NewOpc = X86::MOV64rr; break;
+      case X86::OR64ri32:  NewOpc = X86::OR64rr;  break;
+      case X86::SBB64ri32: NewOpc = X86::SBB64rr; break;
+      case X86::SUB64ri32: NewOpc = X86::SUB64rr; break;
+      case X86::TEST64ri32: NewOpc = X86::TEST64rr; break;
+      case X86::XOR64ri32: NewOpc = X86::XOR64rr; break;
+      case X86::ADC64mi32: NewOpc = X86::ADC64mr; break;// *64mi32
+      case X86::ADD64mi32: NewOpc = X86::ADD64mr; break;
+      case X86::AND64mi32: NewOpc = X86::AND64mr; break;
+      case X86::CMP64mi32: NewOpc = X86::CMP64mr; break;
+      case X86::MOV64mi32: NewOpc = X86::MOV64mr; break;
+      case X86::OR64mi32:  NewOpc = X86::OR64mr;  break;
+      case X86::SBB64mi32: NewOpc = X86::SBB64mr; break;
+      case X86::SUB64mi32: NewOpc = X86::SUB64mr; break;
+      case X86::TEST64mi32: NewOpc = X86::TEST64mr; break;
+      case X86::XOR64mi32: NewOpc = X86::XOR64mr; break;
+      }
+      if (Idx == 0)
+        N1 = CurDAG->getMachineNode(NewOpc, dl, VT, SDValue(N0, 0));
+      else if (Idx == 1)
+        N1 = CurDAG->getMachineNode(NewOpc, dl, VT,
+                                    {N->getOperand(0), SDValue(N0, 0)});
+      else if (Idx == 5) { // *mi
+        // Mem operand should always has a chain.
+        N1 = CurDAG->getMachineNode(NewOpc, dl, N->getVTList(),
+                                    {N->getOperand(0),
+                                    N->getOperand(1),
+                                    N->getOperand(2),
+                                    N->getOperand(3),
+                                    N->getOperand(4),
+                                    SDValue(N0, 0),
+                                    N->getOperand(6) /*chain*/ });
+        CurDAG->setNodeMemRefs(N1, cast<MachineSDNode>(N)->memoperands());
+      } else
+        llvm_unreachable("Unexpected Index!");
+      ReplaceUses(N, N1);
+      return true;
+    }
+    case X86::MOV32ri64:
+    case X86::MOV64ri:{
+      // There maybe some address operation using *ri* opcode, since address
+      // may changed after assembling and linking, we don't handle it here.
+      if (!isa<llvm::ConstantSDNode>(N->getOperand(0)))
+        return false;
+      uint64_t Imm = N->getConstantOperandVal(0);
+      if (!isEndbrImm64(Imm))
+        return false;
+      N0 = CurDAG->getMachineNode(X86::MOV64ri, dl, VT, getI64Imm(~Imm, dl));
+      N1 = CurDAG->getMachineNode(X86::NOT64r, dl, VT, SDValue(N0, 0));
+      ReplaceUses(N, N1);
+      return true;
+    }
+  }
+  return false;
+}
+
 void X86DAGToDAGISel::PostprocessISelDAG() {
   // Skip peepholes at -O0.
   if (TM.getOptLevel() == CodeGenOpt::None)
     return;
 
   SelectionDAG::allnodes_iterator Position = CurDAG->allnodes_end();
 
   bool MadeChange = false;
if ((Opc == X86::KORTESTBrr || Opc == X86::KORTESTWrr ||
                                                       And.getOperand(0),
                                                       And.getOperand(1));
         ReplaceUses(N, KTest);
         MadeChange = true;
         continue;
       }
     }
 
+    // Check that the cf-protection-branch is enabled.
+    Metadata *CFProtectionBranch =
+      MF->getMMI().getModule()->getModuleFlag("cf-protection-branch");
+    if ((CFProtectionBranch || IndirectBranchTracking) &&
+        EndbrImmAntiAttack(N)) {
+      MadeChange = true;
+      continue;
+    }
+
     // Attempt to remove vectors moves that were inserted to zero upper bits.
     if (Opc != TargetOpcode::SUBREG_TO_REG)
       continue;
 
     unsigned SubRegIdx = N->getConstantOperandVal(2);
     if (SubRegIdx != X86::sub_xmm && SubRegIdx != X86::sub_ymm)
       continue;
 

llvm/lib/Target/X86/X86IndirectBranchTracking.cpp

 #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
+
+; 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:
+; 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.
+
+; 0xF30F1EFA == -217112838  ~0xF30F1EFA == 217112837 (0xCF0E105)
+; 0x000123F32E0F1EFA == 321002333478650
+; ~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 $-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), %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:    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 $217112837, %eax # imm = 0xCF0E105
+; CHECK-NEXT:    notl %eax
+; CHECK-NEXT:    movl %eax, -{{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movl %eax, %eax
+; 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:    movabsq $-4077854459, %rax # imm = 0xFFFFFFFF0CF0E105
+; CHECK-NEXT:    notq %rax
+; CHECK-NEXT:    movq %rax, -{{[0-9]+}}(%rsp)
+; 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
+}