[llvm] [LFI][X86] Add control-flow rewrites (PR #210913)

[Zachary Yedidia via llvm-commits]

================
@@ -0,0 +1,92 @@
+//===- X86LFIRewritePass.cpp - Modify code generation for LFI ---*- C++ -*-===//
+//
+// 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 implements the X86LFIRewritePass, which prepares machine code for
+// LFI sandboxing by making sure that every address which may legitimately be
+// the destination of an indirect branch is aligned to a bundle boundary.
+//
+//===----------------------------------------------------------------------===//
+
+#include "MCTargetDesc/X86MCLFIRewriter.h"
+#include "X86.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineJumpTableInfo.h"
+#include "llvm/Support/Alignment.h"
+#include "llvm/Target/TargetMachine.h"
+
+using namespace llvm;
+
+static constexpr Align BundleAlign = Align::Constant<X86::LFIBundleSize>();
+
+namespace {
+class X86LFIRewriteLegacy : public MachineFunctionPass {
+public:
+  static char ID;
+  X86LFIRewriteLegacy() : MachineFunctionPass(ID) {}
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+
+  StringRef getPassName() const override { return "X86 LFI rewrites"; }
+};
+} // namespace
+
+char X86LFIRewriteLegacy::ID = 0;
+
+static void alignToBundle(MachineBasicBlock &MBB) {
+  MBB.setAlignment(std::max(MBB.getAlignment(), BundleAlign), /*MaxBytes=*/0);
+}
+
+// Returns true if MBB may be reached by an indirect branch (does not include
+// jump table targets).
+static bool isIndirectlyReachable(MachineFunction &MF,
+                                  const MachineBasicBlock &MBB) {
+  if (MBB.hasAddressTaken() || MBB.isEHPad())
----------------
zyedidia wrote:

Yes we do need to consider EHPad since the unwinder runs inside the sandbox and resumes at the landing pad through an indirect jump, which gets the same masking as any other indirect branch.

https://github.com/llvm/llvm-project/pull/210913