[clang] [PAC] Authenticate function pointers in UBSan type checks (PR #99590)

[Akira Hatanaka via cfe-commits]

https://github.com/ahatanak created https://github.com/llvm/llvm-project/pull/99590

The function pointer needs to be authenticated before doing the type checks.

>From c44fbc480f8632a178633637010ab6953ed3e50d Mon Sep 17 00:00:00 2001
From: Akira Hatanaka <ahatanak at gmail.com>
Date: Thu, 18 Jul 2024 16:54:09 -0700
Subject: [PATCH] [PAC] Authenticate function pointers in UBSan type checks

The function pointer needs to be authenticated before doing the type
checks.
---
 clang/lib/CodeGen/Address.h           | 4 +++-
 clang/lib/CodeGen/CGExpr.cpp          | 9 +++++++++
 clang/lib/CodeGen/CGPointerAuth.cpp   | 4 ++++
 clang/test/CodeGen/ubsan-function.cpp | 4 ++++
 4 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/clang/lib/CodeGen/Address.h b/clang/lib/CodeGen/Address.h
index 1c4d2e103b5e7..a18c7169af1eb 100644
--- a/clang/lib/CodeGen/Address.h
+++ b/clang/lib/CodeGen/Address.h
@@ -249,7 +249,9 @@ class Address {
   /// Return the pointer contained in this class after authenticating it and
   /// adding offset to it if necessary.
   llvm::Value *emitRawPointer(CodeGenFunction &CGF) const {
-    return getBasePointer();
+    if (!isSigned())
+      return getBasePointer();
+    return emitRawPointerSlow(CGF);
   }
 
   /// Return address with different pointer, but same element type and
diff --git a/clang/lib/CodeGen/CGExpr.cpp b/clang/lib/CodeGen/CGExpr.cpp
index c5341e4b53651..aa53f96358044 100644
--- a/clang/lib/CodeGen/CGExpr.cpp
+++ b/clang/lib/CodeGen/CGExpr.cpp
@@ -5837,6 +5837,15 @@ RValue CodeGenFunction::EmitCall(QualType CalleeType, const CGCallee &OrigCallee
           CGM.getLLVMContext(), {PrefixSigType, Int32Ty}, /*isPacked=*/true);
 
       llvm::Value *CalleePtr = Callee.getFunctionPointer();
+      if (CGM.getCodeGenOpts().PointerAuth.FunctionPointers) {
+        // Use raw pointer since we are using the callee pointer as data here.
+        Address Addr =
+            Address(CalleePtr, CalleePtr->getType(),
+                    CharUnits::fromQuantity(
+                        CalleePtr->getPointerAlignment(CGM.getDataLayout())),
+                    Callee.getPointerAuthInfo(), nullptr);
+        CalleePtr = Addr.emitRawPointer(*this);
+      }
 
       // On 32-bit Arm, the low bit of a function pointer indicates whether
       // it's using the Arm or Thumb instruction set. The actual first
diff --git a/clang/lib/CodeGen/CGPointerAuth.cpp b/clang/lib/CodeGen/CGPointerAuth.cpp
index 7fe62c0788742..72aed9f24d595 100644
--- a/clang/lib/CodeGen/CGPointerAuth.cpp
+++ b/clang/lib/CodeGen/CGPointerAuth.cpp
@@ -566,6 +566,10 @@ Address Address::getResignedAddress(const CGPointerAuthInfo &NewInfo,
                  /*Offset=*/nullptr, isKnownNonNull());
 }
 
+llvm::Value *Address::emitRawPointerSlow(CodeGenFunction &CGF) const {
+  return CGF.getAsNaturalPointerTo(*this, QualType());
+}
+
 llvm::Value *LValue::getPointer(CodeGenFunction &CGF) const {
   assert(isSimple());
   return emitResignedPointer(getType(), CGF);
diff --git a/clang/test/CodeGen/ubsan-function.cpp b/clang/test/CodeGen/ubsan-function.cpp
index bc37a61c98ee3..8478f05a10b78 100644
--- a/clang/test/CodeGen/ubsan-function.cpp
+++ b/clang/test/CodeGen/ubsan-function.cpp
@@ -4,6 +4,8 @@
 // RUN: %clang_cc1 -triple aarch64_be-linux-gnu -emit-llvm -o - %s -fsanitize=function -fno-sanitize-recover=all | FileCheck %s --check-prefixes=CHECK,GNU,64
 // RUN: %clang_cc1 -triple arm-none-eabi -emit-llvm -o - %s -fsanitize=function -fno-sanitize-recover=all | FileCheck %s --check-prefixes=CHECK,ARM,GNU,32
 
+// RUN: %clang_cc1 -triple arm64e-apple-ios -emit-llvm -o - %s -fsanitize=function -fno-sanitize-recover=all -fptrauth-calls | FileCheck %s --check-prefixes=CHECK,GNU,64,64e
+
 // GNU:  define{{.*}} void @_Z3funv() #0 !func_sanitize ![[FUNCSAN:.*]] {
 // MSVC: define{{.*}} void @"?fun@@YAXXZ"() #0 !func_sanitize ![[FUNCSAN:.*]] {
 void fun() {}
@@ -13,6 +15,8 @@ void fun() {}
 // ARM:   ptrtoint ptr {{.*}} to i32, !nosanitize !5
 // ARM:   and i32 {{.*}}, -2, !nosanitize !5
 // ARM:   inttoptr i32 {{.*}} to ptr, !nosanitize !5
+// 64e:   %[[STRIPPED:.*]] = ptrtoint ptr {{.*}} to i64, !nosanitize
+// 64e:   call i64 @llvm.ptrauth.auth(i64 %[[STRIPPED]], i32 0, i64 0), !nosanitize
 // CHECK: getelementptr <{ i32, i32 }>, ptr {{.*}}, i32 -1, i32 0, !nosanitize
 // CHECK: load i32, ptr {{.*}}, align {{.*}}, !nosanitize
 // CHECK: icmp eq i32 {{.*}}, -1056584962, !nosanitize