Add ehcont section support

In the future Windows will enable Control-flow Enforcement Technology (CET aka shadow stacks). To protect the path where the context is updated during exception handling, the binary is required to enumerate valid unwind entrypoints in a dedicated section which is validated when the context is being set during exception handling.

This change allows llvm to generate the section that contains the appropriate symbol references in the form expected by the msvc linker.

This feature is enabled through a new module flag, ehcontguard, which was modelled on the cfguard flag.

The change includes a test that when the module flag is enabled the section is correctly generated.

The set of exception continuation information includes returns from exceptional control flow (catchret in llvm).

In order to collect catchret we:
1) Includes an additional flag on machine basic blocks to indicate that the given block is the target of a catchret operation,
2) Introduces a new machine function pass to insert and collect symbols at the start of each block, and
3) Combines these targets with the other EHCont targets that were already being collected.

Change originally authored by Daniel Frampton <dframpto@microsoft.com>

For more details, see MSVC documentation for `/guard:ehcont`
  https://docs.microsoft.com/en-us/cpp/build/reference/guard-enable-eh-continuation-metadata

Reviewed By: pengfei

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

include/llvm/CodeGen/MachineBasicBlock.h

@@ -153,6 +153,9 @@ private:
   /// LLVM IR.
   bool IsEHScopeEntry = false;
 
+  /// Indicates if this is a target block of a catchret.
+  bool IsEHCatchretTarget = false;
+
   /// Indicate that this basic block is the entry block of an EH funclet.
   bool IsEHFuncletEntry = false;
 
@@ -175,6 +178,9 @@ private:
   /// is only computed once and is cached.
   mutable MCSymbol *CachedMCSymbol = nullptr;
 
+  /// Cached MCSymbol for this block (used if IsEHCatchRetTarget).
+  mutable MCSymbol *CachedEHCatchretMCSymbol = nullptr;
+
   /// Marks the end of the basic block. Used during basic block sections to
   /// calculate the size of the basic block, or the BB section ending with it.
   mutable MCSymbol *CachedEndMCSymbol = nullptr;
@@ -445,6 +451,12 @@ public:
   /// that used to have a catchpad or cleanuppad instruction in the LLVM IR.
   void setIsEHScopeEntry(bool V = true) { IsEHScopeEntry = V; }
 
+  /// Returns true if this is a target block of a catchret.
+  bool isEHCatchretTarget() const { return IsEHCatchretTarget; }
+
+  /// Indicates if this is a target block of a catchret.
+  void setIsEHCatchretTarget(bool V = true) { IsEHCatchretTarget = V; }
+
   /// Returns true if this is the entry block of an EH funclet.
   bool isEHFuncletEntry() const { return IsEHFuncletEntry; }
 
@@ -910,6 +922,9 @@ public:
   /// Return the MCSymbol for this basic block.
   MCSymbol *getSymbol() const;
 
+  /// Return the EHCatchret Symbol for this basic block.
+  MCSymbol *getEHCatchretSymbol() const;
+
   Optional<uint64_t> getIrrLoopHeaderWeight() const {
     return IrrLoopHeaderWeight;
   }

include/llvm/CodeGen/MachineFunction.h

@@ -321,6 +321,10 @@ class MachineFunction {
   /// construct a table of valid longjmp targets for Windows Control Flow Guard.
   std::vector<MCSymbol *> LongjmpTargets;
 
+  /// List of basic blocks that are the target of catchrets. Used to construct
+  /// a table of valid targets for Windows EHCont Guard.
+  std::vector<MCSymbol *> CatchretTargets;
+
   /// \name Exception Handling
   /// \{
 
@@ -341,6 +345,7 @@ class MachineFunction {
 
   bool CallsEHReturn = false;
   bool CallsUnwindInit = false;
+  bool HasEHCatchret = false;
   bool HasEHScopes = false;
   bool HasEHFunclets = false;
 
@@ -930,6 +935,18 @@ public:
   /// Control Flow Guard.
   void addLongjmpTarget(MCSymbol *Target) { LongjmpTargets.push_back(Target); }
 
+  /// Returns a reference to a list of symbols that we have catchrets.
+  /// Used to construct the catchret target table used by Windows EHCont Guard.
+  const std::vector<MCSymbol *> &getCatchretTargets() const {
+    return CatchretTargets;
+  }
+
+  /// Add the specified symbol to the list of valid catchret targets for Windows
+  /// EHCont Guard.
+  void addCatchretTarget(MCSymbol *Target) {
+    CatchretTargets.push_back(Target);
+  }
+
   /// \name Exception Handling
   /// \{
 
@@ -939,6 +956,9 @@ public:
   bool callsUnwindInit() const { return CallsUnwindInit; }
   void setCallsUnwindInit(bool b) { CallsUnwindInit = b; }
 
+  bool hasEHCatchret() const { return HasEHCatchret; }
+  void setHasEHCatchret(bool V) { HasEHCatchret = V; }
+
   bool hasEHScopes() const { return HasEHScopes; }
   void setHasEHScopes(bool V) { HasEHScopes = V; }
 

include/llvm/CodeGen/Passes.h

@@ -468,6 +468,10 @@ namespace llvm {
   /// \see CFGuardLongjmp.cpp
   FunctionPass *createCFGuardLongjmpPass();
 
+  /// Creates EHContGuard catchret target identification pass.
+  /// \see EHContGuardCatchret.cpp
+  FunctionPass *createEHContGuardCatchretPass();
+
   /// Create Hardware Loop pass. \see HardwareLoops.cpp
   FunctionPass *createHardwareLoopsPass();
 

include/llvm/InitializePasses.h

@@ -148,6 +148,7 @@ void initializeEarlyIfPredicatorPass(PassRegistry &);
 void initializeEarlyMachineLICMPass(PassRegistry&);
 void initializeEarlyTailDuplicatePass(PassRegistry&);
 void initializeEdgeBundlesPass(PassRegistry&);
+void initializeEHContGuardCatchretPass(PassRegistry &);
 void initializeEliminateAvailableExternallyLegacyPassPass(PassRegistry&);
 void initializeEntryExitInstrumenterPass(PassRegistry&);
 void initializeExpandMemCmpPassPass(PassRegistry&);

include/llvm/MC/MCObjectFileInfo.h

@@ -218,6 +218,7 @@ protected:
   MCSection *PDataSection = nullptr;
   MCSection *XDataSection = nullptr;
   MCSection *SXDataSection = nullptr;
+  MCSection *GEHContSection = nullptr;
   MCSection *GFIDsSection = nullptr;
   MCSection *GIATsSection = nullptr;
   MCSection *GLJMPSection = nullptr;
@@ -405,6 +406,7 @@ public:
   MCSection *getPDataSection() const { return PDataSection; }
   MCSection *getXDataSection() const { return XDataSection; }
   MCSection *getSXDataSection() const { return SXDataSection; }
+  MCSection *getGEHContSection() const { return GEHContSection; }
   MCSection *getGFIDsSection() const { return GFIDsSection; }
   MCSection *getGIATsSection() const { return GIATsSection; }
   MCSection *getGLJMPSection() const { return GLJMPSection; }

lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -3197,6 +3197,10 @@ void AsmPrinter::emitBasicBlockStart(const MachineBasicBlock &MBB) {
     }
   }
 
+  if (MBB.isEHCatchretTarget()) {
+    OutStreamer->emitLabel(MBB.getEHCatchretSymbol());
+  }
+
   // With BB sections, each basic block must handle CFI information on its own
   // if it begins a section (Entry block is handled separately by
   // AsmPrinterHandler::beginFunction).

lib/CodeGen/AsmPrinter/WinException.cpp

@@ -55,6 +55,14 @@ void WinException::endModule() {
   for (const Function &F : *M)
     if (F.hasFnAttribute("safeseh"))
       OS.EmitCOFFSafeSEH(Asm->getSymbol(&F));
+
+  if (M->getModuleFlag("ehcontguard") && !EHContTargets.empty()) {
+    // Emit the symbol index of each ehcont target.
+    OS.SwitchSection(Asm->OutContext.getObjectFileInfo()->getGEHContSection());
+    for (const MCSymbol *S : EHContTargets) {
+      OS.EmitCOFFSymbolIndex(S);
+    }
+  }
 }
 
 void WinException::beginFunction(const MachineFunction *MF) {
@@ -164,6 +172,12 @@ void WinException::endFunction(const MachineFunction *MF) {
 
     Asm->OutStreamer->PopSection();
   }
+
+  if (!MF->getCatchretTargets().empty()) {
+    // Copy the function's catchret targets to a module-level list.
+    EHContTargets.insert(EHContTargets.end(), MF->getCatchretTargets().begin(),
+                         MF->getCatchretTargets().end());
+  }
 }
 
 /// Retrieve the MCSymbol for a GlobalValue or MachineBasicBlock.

lib/CodeGen/AsmPrinter/WinException.h

@@ -44,6 +44,9 @@ class LLVM_LIBRARY_VISIBILITY WinException : public EHStreamer {
   /// The section of the last funclet start.
   MCSection *CurrentFuncletTextSection = nullptr;
 
+  /// The list of symbols to add to the ehcont section
+  std::vector<const MCSymbol *> EHContTargets;
+
   void emitCSpecificHandlerTable(const MachineFunction *MF);
 
   void emitSEHActionsForRange(const WinEHFuncInfo &FuncInfo,

lib/CodeGen/CMakeLists.txt

@@ -24,6 +24,7 @@ add_llvm_component_library(LLVMCodeGen
   DwarfEHPrepare.cpp
   EarlyIfConversion.cpp
   EdgeBundles.cpp
+  EHContGuardCatchret.cpp
   ExecutionDomainFix.cpp
   ExpandMemCmp.cpp
   ExpandPostRAPseudos.cpp

lib/CodeGen/EHContGuardCatchret.cpp

@@ -0,0 +1,84 @@
+//===-- EHContGuardCatchret.cpp - Catchret target symbols -------*- 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
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// This file contains a machine function pass to insert a symbol before each
+/// valid catchret target and store this in the MachineFunction's
+/// CatchRetTargets vector. This will be used to emit the table of valid targets
+/// used by EHCont Guard.
+///
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/Statistic.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineInstr.h"
+#include "llvm/CodeGen/MachineModuleInfo.h"
+#include "llvm/CodeGen/MachineOperand.h"
+#include "llvm/CodeGen/Passes.h"
+#include "llvm/InitializePasses.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "ehcontguard-catchret"
+
+STATISTIC(EHContGuardCatchretTargets,
+          "Number of EHCont Guard catchret targets");
+
+namespace {
+
+/// MachineFunction pass to insert a symbol before each valid catchret target
+/// and store these in the MachineFunction's CatchRetTargets vector.
+class EHContGuardCatchret : public MachineFunctionPass {
+public:
+  static char ID;
+
+  EHContGuardCatchret() : MachineFunctionPass(ID) {
+    initializeEHContGuardCatchretPass(*PassRegistry::getPassRegistry());
+  }
+
+  StringRef getPassName() const override {
+    return "EH Cont Guard catchret targets";
+  }
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+};
+
+} // end anonymous namespace
+
+char EHContGuardCatchret::ID = 0;
+
+INITIALIZE_PASS(EHContGuardCatchret, "EHContGuardCatchret",
+                "Insert symbols at valid catchret targets for /guard:ehcont",
+                false, false)
+FunctionPass *llvm::createEHContGuardCatchretPass() {
+  return new EHContGuardCatchret();
+}
+
+bool EHContGuardCatchret::runOnMachineFunction(MachineFunction &MF) {
+
+  // Skip modules for which the ehcontguard flag is not set.
+  if (!MF.getMMI().getModule()->getModuleFlag("ehcontguard"))
+    return false;
+
+  // Skip functions that do not have catchret
+  if (!MF.hasEHCatchret())
+    return false;
+
+  bool Result = false;
+
+  for (MachineBasicBlock &MBB : MF) {
+    if (MBB.isEHCatchretTarget()) {
+      MF.addCatchretTarget(MBB.getEHCatchretSymbol());
+      EHContGuardCatchretTargets++;
+      Result = true;
+    }
+  }
+
+  return Result;
+}

lib/CodeGen/MachineBasicBlock.cpp

@@ -87,6 +87,17 @@ MCSymbol *MachineBasicBlock::getSymbol() const {
   return CachedMCSymbol;
 }
 
+MCSymbol *MachineBasicBlock::getEHCatchretSymbol() const {
+  if (!CachedEHCatchretMCSymbol) {
+    const MachineFunction *MF = getParent();
+    SmallString<128> SymbolName;
+    raw_svector_ostream(SymbolName)
+        << "$ehgcr_" << MF->getFunctionNumber() << '_' << getNumber();
+    CachedEHCatchretMCSymbol = MF->getContext().getOrCreateSymbol(SymbolName);
+  }
+  return CachedEHCatchretMCSymbol;
+}
+
 MCSymbol *MachineBasicBlock::getEndSymbol() const {
   if (!CachedEndMCSymbol) {
     const MachineFunction *MF = getParent();

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

@@ -1593,6 +1593,8 @@ void SelectionDAGBuilder::visitCatchRet(const CatchReturnInst &I) {
   // Update machine-CFG edge.
   MachineBasicBlock *TargetMBB = FuncInfo.MBBMap[I.getSuccessor()];
   FuncInfo.MBB->addSuccessor(TargetMBB);
+  TargetMBB->setIsEHCatchretTarget(true);
+  DAG.getMachineFunction().setHasEHCatchret(true);
 
   auto Pers = classifyEHPersonality(FuncInfo.Fn->getPersonalityFn());
   bool IsSEH = isAsynchronousEHPersonality(Pers);

lib/MC/MCObjectFileInfo.cpp

@@ -753,6 +753,11 @@ void MCObjectFileInfo::initCOFFMCObjectFileInfo(const Triple &T) {
   SXDataSection = Ctx->getCOFFSection(".sxdata", COFF::IMAGE_SCN_LNK_INFO,
                                       SectionKind::getMetadata());
 
+  GEHContSection = Ctx->getCOFFSection(".gehcont$y",
+                                       COFF::IMAGE_SCN_CNT_INITIALIZED_DATA |
+                                           COFF::IMAGE_SCN_MEM_READ,
+                                       SectionKind::getMetadata());
+
   GFIDsSection = Ctx->getCOFFSection(".gfids$y",
                                      COFF::IMAGE_SCN_CNT_INITIALIZED_DATA |
                                          COFF::IMAGE_SCN_MEM_READ,

lib/Target/AArch64/AArch64AsmPrinter.cpp

@@ -191,7 +191,32 @@ private:
 } // end anonymous namespace
 
 void AArch64AsmPrinter::emitStartOfAsmFile(Module &M) {
-  if (!TM.getTargetTriple().isOSBinFormatELF())
+  const Triple &TT = TM.getTargetTriple();
+
+  if (TT.isOSBinFormatCOFF()) {
+    // Emit an absolute @feat.00 symbol.  This appears to be some kind of
+    // compiler features bitfield read by link.exe.
+    MCSymbol *S = MMI->getContext().getOrCreateSymbol(StringRef("@feat.00"));
+    OutStreamer->BeginCOFFSymbolDef(S);
+    OutStreamer->EmitCOFFSymbolStorageClass(COFF::IMAGE_SYM_CLASS_STATIC);
+    OutStreamer->EmitCOFFSymbolType(COFF::IMAGE_SYM_DTYPE_NULL);
+    OutStreamer->EndCOFFSymbolDef();
+    int64_t Feat00Flags = 0;
+
+    if (M.getModuleFlag("cfguard")) {
+      Feat00Flags |= 0x800; // Object is CFG-aware.
+    }
+
+    if (M.getModuleFlag("ehcontguard")) {
+      Feat00Flags |= 0x4000; // Object also has EHCont.
+    }
+
+    OutStreamer->emitSymbolAttribute(S, MCSA_Global);
+    OutStreamer->emitAssignment(
+        S, MCConstantExpr::create(Feat00Flags, MMI->getContext()));
+  }
+
+  if (!TT.isOSBinFormatELF())
     return;
 
   // Assemble feature flags that may require creation of a note section.

lib/Target/AArch64/AArch64TargetMachine.cpp

@@ -682,9 +682,12 @@ void AArch64PassConfig::addPreEmitPass() {
   if (BranchRelaxation)
     addPass(&BranchRelaxationPassID);
 
-  // Identify valid longjmp targets for Windows Control Flow Guard.
-  if (TM->getTargetTriple().isOSWindows())
+  if (TM->getTargetTriple().isOSWindows()) {
+    // Identify valid longjmp targets for Windows Control Flow Guard.
     addPass(createCFGuardLongjmpPass());
+    // Identify valid eh continuation targets for Windows EHCont Guard.
+    addPass(createEHContGuardCatchretPass());
+  }
 
   if (TM->getOptLevel() != CodeGenOpt::None && EnableCompressJumpTables)
     addPass(createAArch64CompressJumpTablesPass());

lib/Target/ARM/ARMTargetMachine.cpp

@@ -564,7 +564,10 @@ void ARMPassConfig::addPreEmitPass2() {
   addPass(createARMConstantIslandPass());
   addPass(createARMLowOverheadLoopsPass());
 
-  // Identify valid longjmp targets for Windows Control Flow Guard.
-  if (TM->getTargetTriple().isOSWindows())
+  if (TM->getTargetTriple().isOSWindows()) {
+    // Identify valid longjmp targets for Windows Control Flow Guard.
     addPass(createCFGuardLongjmpPass());
+    // Identify valid eh continuation targets for Windows EHCont Guard.
+    addPass(createEHContGuardCatchretPass());
+  }
 }

lib/Target/X86/X86AsmPrinter.cpp

@@ -683,8 +683,13 @@ void X86AsmPrinter::emitStartOfAsmFile(Module &M) {
       Feat00Flags |= 1;
     }
 
-    if (M.getModuleFlag("cfguard"))
+    if (M.getModuleFlag("cfguard")) {
       Feat00Flags |= 0x800; // Object is CFG-aware.
+    }
+
+    if (M.getModuleFlag("ehcontguard")) {
+      Feat00Flags |= 0x4000; // Object also has EHCont.
+    }
 
     OutStreamer->emitSymbolAttribute(S, MCSA_Global);
     OutStreamer->emitAssignment(

lib/Target/X86/X86TargetMachine.cpp

@@ -568,9 +568,13 @@ void X86PassConfig::addPreEmitPass2() {
       (!TT.isOSWindows() ||
        MAI->getExceptionHandlingType() == ExceptionHandling::DwarfCFI))
     addPass(createCFIInstrInserter());
-  // Identify valid longjmp targets for Windows Control Flow Guard.
-  if (TT.isOSWindows())
+
+  if (TT.isOSWindows()) {
+    // Identify valid longjmp targets for Windows Control Flow Guard.
     addPass(createCFGuardLongjmpPass());
+    // Identify valid eh continuation targets for Windows EHCont Guard.
+    addPass(createEHContGuardCatchretPass());
+  }
   addPass(createX86LoadValueInjectionRetHardeningPass());
 }
 

test/CodeGen/AArch64/ehcontguard.ll

@@ -0,0 +1,29 @@
+; RUN: llc < %s -mtriple=aarch64-windows | FileCheck %s
+; EHCont Guard is currently only available on Windows
+
+; CHECK: .set @feat.00, 16384
+
+; CHECK: .section .gehcont$y
+
+define dso_local void @"?func1@@YAXXZ"() #0 personality i8* bitcast (i32 (...)* @__CxxFrameHandler3 to i8*) {
+entry:
+  invoke void @"?func2@@YAXXZ"()
+          to label %invoke.cont unwind label %catch.dispatch
+catch.dispatch:                                   ; preds = %entry
+  %0 = catchswitch within none [label %catch] unwind to caller
+catch:                                            ; preds = %catch.dispatch
+  %1 = catchpad within %0 [i8* null, i32 64, i8* null]
+  catchret from %1 to label %catchret.dest
+catchret.dest:                                    ; preds = %catch
+  br label %try.cont
+try.cont:                                         ; preds = %catchret.dest, %invoke.cont
+  ret void
+invoke.cont:                                      ; preds = %entry
+  br label %try.cont
+}
+
+declare dso_local void @"?func2@@YAXXZ"() #1
+declare dso_local i32 @__CxxFrameHandler3(...)
+
+!llvm.module.flags = !{!0}
+!0 = !{i32 1, !"ehcontguard", i32 1}

test/CodeGen/X86/ehcontguard.ll

@@ -0,0 +1,29 @@
+; RUN: llc < %s -mtriple=x86_64-pc-windows-msvc | FileCheck %s
+; EHCont Guard is currently only available on Windows
+
+; CHECK: .set @feat.00, 16384
+
+; CHECK: .section .gehcont$y
+
+define dso_local void @"?func1@@YAXXZ"() #0 personality i8* bitcast (i32 (...)* @__CxxFrameHandler3 to i8*) {
+entry:
+  invoke void @"?func2@@YAXXZ"()
+          to label %invoke.cont unwind label %catch.dispatch
+catch.dispatch:                                   ; preds = %entry
+  %0 = catchswitch within none [label %catch] unwind to caller
+catch:                                            ; preds = %catch.dispatch
+  %1 = catchpad within %0 [i8* null, i32 64, i8* null]
+  catchret from %1 to label %catchret.dest
+catchret.dest:                                    ; preds = %catch
+  br label %try.cont
+try.cont:                                         ; preds = %catchret.dest, %invoke.cont
+  ret void
+invoke.cont:                                      ; preds = %entry
+  br label %try.cont
+}
+
+declare dso_local void @"?func2@@YAXXZ"() #1
+declare dso_local i32 @__CxxFrameHandler3(...)
+
+!llvm.module.flags = !{!0}
+!0 = !{i32 1, !"ehcontguard", i32 1}

test/CodeGen/X86/win-catchpad-nested.ll

@@ -32,8 +32,10 @@ inner.catch:
 ; CHECK-LABEL: test1: # @test1
 ; CHECK: [[Exit:^[^: ]+]]: # Block address taken
 ; CHECK-NEXT:              # %exit
+; CHECK-NEXT: $ehgcr_0_1:
 ; CHECK: [[OuterRet:^[^: ]+]]: # Block address taken
 ; CHECK-NEXT:                  # %outer.ret
+; CHECK-NEXT: $ehgcr_0_3:
 ; CHECK-NEXT: leaq [[Exit]](%rip), %rax
 ; CHECK:      retq   # CATCHRET
 ; CHECK: {{^[^: ]+}}: # %inner.catch

utils/gn/secondary/llvm/lib/CodeGen/BUILD.gn

@@ -42,6 +42,7 @@ static_library("CodeGen") {
     "DwarfEHPrepare.cpp",
     "EarlyIfConversion.cpp",
     "EdgeBundles.cpp",
+    "EHContGuardCatchret.cpp",
     "ExecutionDomainFix.cpp",
     "ExpandMemCmp.cpp",
     "ExpandPostRAPseudos.cpp",