[WinCFG] Handle constant casts carefully in .gfids emission

Summary:
The general Function::hasAddressTaken has two issues that make it
inappropriate for our purposes:
1. it is sensitive to dead constant users (PR43858 / crbug.com/1019970),
   leading to different codegen when debu info is enabled
2. it considers direct calls via a function cast to be address escapes

The first is fixable, but the second is not, because IPO clients rely on
this behavior. They assume this function means that all call sites are
analyzable for IPO purposes.

So, implement our own analysis, which gets closer to finding functions
that may be indirect call targets.

Reviewers: ajpaverd, efriedma, hans

Subscribers: hiraditya, llvm-commits

Tags: #llvm

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

2 files changed, 73 insertions, 1 deletions

diff --git a/llvm/lib/CodeGen/AsmPrinter/WinCFGuard.cpp b/llvm/lib/CodeGen/AsmPrinter/WinCFGuard.cpp
index f071a2583e5..914308d9147 100644
--- a/llvm/lib/CodeGen/AsmPrinter/WinCFGuard.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/WinCFGuard.cpp
@@ -18,6 +18,7 @@
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Metadata.h"
+#include "llvm/IR/Instructions.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCObjectFileInfo.h"
 #include "llvm/MC/MCStreamer.h"
@@ -41,11 +42,47 @@ void WinCFGuard::endFunction(const MachineFunction *MF) {
                         MF->getLongjmpTargets().end());
 }
 
+/// Returns true if this function's address is escaped in a way that might make
+/// it an indirect call target. Function::hasAddressTaken gives different
+/// results when a function is called directly with a function prototype
+/// mismatch, which requires a cast.
+static bool isPossibleIndirectCallTarget(const Function *F) {
+  SmallVector<const Value *, 4> Users{F};
+  while (!Users.empty()) {
+    const Value *FnOrCast = Users.pop_back_val();
+    for (const Use &U : FnOrCast->uses()) {
+      const User *FnUser = U.getUser();
+      if (isa<BlockAddress>(FnUser))
+        continue;
+      if (const auto *Call = dyn_cast<CallBase>(FnUser)) {
+        if (!Call->isCallee(&U))
+          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;
+      }
+    }
+  }
+  return false;
+}
+
 void WinCFGuard::endModule() {
   const Module *M = Asm->MMI->getModule();
   std::vector<const Function *> Functions;
   for (const Function &F : *M)
-    if (F.hasAddressTaken())
+    if (isPossibleIndirectCallTarget(&F))
       Functions.push_back(&F);
   if (Functions.empty() && LongjmpTargets.empty())
     return;
diff --git a/llvm/test/CodeGen/WinCFGuard/cfguard-cast.ll b/llvm/test/CodeGen/WinCFGuard/cfguard-cast.ll
new file mode 100644
index 00000000000..e66345b0d70
--- /dev/null
+++ b/llvm/test/CodeGen/WinCFGuard/cfguard-cast.ll
@@ -0,0 +1,35 @@
+; RUN: llc < %s -mtriple=x86_64-pc-windows-msvc | FileCheck %s
+
+; Check how constant function pointer casts are handled.
+
+declare void @unprototyped(...)
+
+define i32 @call_unprototyped() {
+  call void bitcast (void (...)* @unprototyped to void ()*)()
+  ret i32 0
+}
+
+; CHECK-LABEL: call_unprototyped:
+; CHECK: callq unprototyped
+; CHECK: xorl %eax, %eax
+; CHECK: retq
+
+declare void @escaped_cast()
+
+define i32 @escape_it_with_cast(i8** %p) {
+  store i8* bitcast (void ()* @escaped_cast to i8*), i8** %p
+  ret i32 0
+}
+
+declare void @dead_constant()
+
+!llvm.module.flags = !{!0}
+!0 = !{i32 2, !"cfguard", i32 1}
+
+!dead_constant_root = !{!1}
+!1 = !DITemplateValueParameter(name: "dead_constant", value: i8* bitcast (void ()* @dead_constant to i8*))
+
+; CHECK-LABEL: .section .gfids$y,"dr"
+; CHECK-NEXT:  .symidx escaped_cast
+; CHECK-NOT:   .symidx
+