[llvm] [ARM64EC] Avoid emitting unnecessary symbol references with /guard:cf. (PR #123235)

[Eli Friedman via llvm-commits]

https://github.com/efriedma-quic created https://github.com/llvm/llvm-project/pull/123235

.gfids$y contains a list of indirect calls for Control Flow Guard. This wasn't working properly for ARM64EC: direct calls were being treated as indirect calls.  Make sure we correctly filter out direct calls.

This improves the protection from Control Flow Guard, and also fixes a link error when using certain functions from oldnames.lib.

>From 5c5b962593d3637eecf7150a93a67fe6136344d8 Mon Sep 17 00:00:00 2001
From: Eli Friedman <efriedma at quicinc.com>
Date: Thu, 16 Jan 2025 11:53:50 -0800
Subject: [PATCH] [ARM64EC] Avoid emitting unnecessary symbol references with
 /guard:cf.

.gfids$y contains a list of indirect calls for Control Flow Guard. This
wasn't working properly for ARM64EC: direct calls were being treated as
indiret calls.  Make sure we correctly filter out direct calls.

This improves the protection from Control Flow Guard, and also fixes a
link error when using certain functions from oldnames.lib.
---
 llvm/lib/CodeGen/AsmPrinter/WinCFGuard.cpp   | 30 ++++++++++++--------
 llvm/test/CodeGen/AArch64/cfguard-arm64ec.ll | 16 +++++++++++
 2 files changed, 34 insertions(+), 12 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/cfguard-arm64ec.ll

diff --git a/llvm/lib/CodeGen/AsmPrinter/WinCFGuard.cpp b/llvm/lib/CodeGen/AsmPrinter/WinCFGuard.cpp
index 1a1e6f0117e2b8..09e4408c9bf234 100644
--- a/llvm/lib/CodeGen/AsmPrinter/WinCFGuard.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/WinCFGuard.cpp
@@ -49,26 +49,32 @@ static bool isPossibleIndirectCallTarget(const Function *F) {
     const Value *FnOrCast = Users.pop_back_val();
     for (const Use &U : FnOrCast->uses()) {
       const User *FnUser = U.getUser();
-      if (isa<BlockAddress>(FnUser))
+      if (isa<BlockAddress>(FnUser)) {
+        // Block addresses are illegal to call.
         continue;
+      }
       if (const auto *Call = dyn_cast<CallBase>(FnUser)) {
-        if (!Call->isCallee(&U))
+        if ((!Call->isCallee(&U) || U != F) &&
+            !Call->getFunction()->getName().ends_with("$exit_thunk")) {
+          // Passing a function pointer to a call may lead to an indirect
+          // call. As an exception, ignore ARM64EC exit thunks.
           return true;
+        }
       } else if (isa<Instruction>(FnUser)) {
         // Consider any other instruction to be an escape. This has some weird
         // consequences like no-op intrinsics being an escape or a store *to* a
         // function address being an escape.
         return true;
-      } else if (const auto *C = dyn_cast<Constant>(FnUser)) {
-        // If this is a constant pointer cast of the function, don't consider
-        // this escape. Analyze the uses of the cast as well. This ensures that
-        // direct calls with mismatched prototypes don't end up in the CFG
-        // table. Consider other constants, such as vtable initializers, to
-        // escape the function.
-        if (C->stripPointerCasts() == F)
-          Users.push_back(FnUser);
-        else
-          return true;
+      } else if (const auto *G = dyn_cast<GlobalValue>(FnUser)) {
+        // Ignore llvm.arm64ec.symbolmap; it doesn't lower to an actual address.
+        if (G->getName() == "llvm.arm64ec.symbolmap")
+          continue;
+        // Globals (for example, vtables) are escapes.
+        return true;
+      } else if (isa<Constant>(FnUser)) {
+        // Constants which aren't a global are intermediate values; recursively
+        // analyze the users to see if they actually escape.
+        Users.push_back(FnUser);
       }
     }
   }
diff --git a/llvm/test/CodeGen/AArch64/cfguard-arm64ec.ll b/llvm/test/CodeGen/AArch64/cfguard-arm64ec.ll
new file mode 100644
index 00000000000000..bdbc99e2d98b0a
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/cfguard-arm64ec.ll
@@ -0,0 +1,16 @@
+; RUN: llc < %s -mtriple=arm64ec-pc-windows-msvc | FileCheck %s
+
+declare void @called()
+declare void @escaped()
+define void @f(ptr %dst) {
+  call void @called()
+  store ptr @escaped, ptr %dst
+  ret void
+}
+
+!llvm.module.flags = !{!0}
+!0 = !{i32 2, !"cfguard", i32 1}
+
+; CHECK-LABEL: .section .gfids$y,"dr"
+; CHECK-NEXT:  .symidx escaped
+; CHECK-NOT:   .symidx