[clang] Implement pointer authentication for C++ member function pointers.

Author: ahmedbougacha

clang/include/clang/Basic/Features.def

@@ -109,6 +109,7 @@ FEATURE(ptrauth_objc_method_list_pointer, LangOpts.PointerAuthCalls)
 FEATURE(ptrauth_vtable_pointer_address_discrimination, LangOpts.PointerAuthVTPtrAddressDiscrimination)
 FEATURE(ptrauth_vtable_pointer_type_discrimination, LangOpts.PointerAuthVTPtrTypeDiscrimination)
 FEATURE(ptrauth_returns, LangOpts.PointerAuthReturns)
+FEATURE(ptrauth_member_function_pointer_type_discrimination, LangOpts.PointerAuthCalls)
 FEATURE(ptrauth_function_pointer_type_discrimination, LangOpts.FunctionPointerTypeDiscrimination)
 FEATURE(ptrauth_signed_block_descriptors, LangOpts.PointerAuthBlockDescriptorPointers)
 FEATURE(swiftasynccc,

clang/include/clang/Basic/PointerAuthOptions.h

@@ -200,6 +200,9 @@ struct PointerAuthOptions {
 
   /// The ABI for variadic C++ virtual function pointers.
   PointerAuthSchema CXXVirtualVariadicFunctionPointers;
+
+  /// The ABI for C++ member function pointers.
+  PointerAuthSchema CXXMemberFunctionPointers;
 };
 
 }  // end namespace clang

clang/lib/CodeGen/CGCall.cpp

@@ -4980,7 +4980,8 @@ RValue CodeGenFunction::EmitCall(const CGFunctionInfo &CallInfo,
                                  ReturnValueSlot ReturnValue,
                                  const CallArgList &CallArgs,
                                  llvm::CallBase **callOrInvoke, bool IsMustTail,
-                                 SourceLocation Loc) {
+                                 SourceLocation Loc,
+                                 bool IsVirtualFunctionPointerThunk) {
   // FIXME: We no longer need the types from CallArgs; lift up and simplify.
 
   assert(Callee.isOrdinary() || Callee.isVirtual());
@@ -5043,8 +5044,12 @@ RValue CodeGenFunction::EmitCall(const CGFunctionInfo &CallInfo,
   RawAddress SRetAlloca = RawAddress::invalid();
   llvm::Value *UnusedReturnSizePtr = nullptr;
   if (RetAI.isIndirect() || RetAI.isInAlloca() || RetAI.isCoerceAndExpand()) {
-    if (!ReturnValue.isNull()) {
-      SRetPtr = ReturnValue.getValue();
+    if (IsVirtualFunctionPointerThunk && RetAI.isIndirect()) {
+      SRetPtr = makeNaturalAddressForPointer(CurFn->arg_begin() +
+                                                 IRFunctionArgs.getSRetArgNo(),
+                                             RetTy, CharUnits::fromQuantity(1));
+    } else if (!ReturnValue.isNull()) {
+      SRetPtr = ReturnValue.getAddress();
     } else {
       SRetPtr = CreateMemTemp(RetTy, "tmp", &SRetAlloca);
       if (HaveInsertPoint() && ReturnValue.isUnused()) {
@@ -5776,7 +5781,14 @@ RValue CodeGenFunction::EmitCall(const CGFunctionInfo &CallInfo,
   CallArgs.freeArgumentMemory(*this);
 
   // Extract the return value.
-  RValue Ret = [&] {
+  RValue Ret;
+
+  // If the current function is a virtual function pointer thunk, avoid copying
+  // the return value of the musttail call to a temporary.
+  if (IsVirtualFunctionPointerThunk)
+    Ret = RValue::get(CI);
+  else
+    Ret = [&] {
     switch (RetAI.getKind()) {
     case ABIArgInfo::CoerceAndExpand: {
       auto coercionType = RetAI.getCoerceAndExpandType();

clang/lib/CodeGen/CGPointerAuth.cpp

@@ -106,6 +106,23 @@ CodeGenModule::getFunctionPointerAuthInfo(QualType T) {
                            /* authenticatesNullValues */ false, Discriminator);
 }
 
+CGPointerAuthInfo
+CodeGenModule::getMemberFunctionPointerAuthInfo(QualType functionType) {
+  assert(functionType->getAs<MemberPointerType>() &&
+         "MemberPointerType expected");
+  auto &schema = getCodeGenOpts().PointerAuth.CXXMemberFunctionPointers;
+  if (!schema)
+    return CGPointerAuthInfo();
+
+  assert(!schema.isAddressDiscriminated() &&
+         "function pointers cannot use address-specific discrimination");
+
+  auto discriminator =
+      getPointerAuthOtherDiscriminator(schema, GlobalDecl(), functionType);
+  return CGPointerAuthInfo(schema.getKey(), schema.getAuthenticationMode(),
+                           /* authenticatesNullValues */ false, discriminator);
+}
+
 /// Return the natural pointer authentication for values of the given
 /// pointee type.
 static CGPointerAuthInfo
@@ -812,6 +829,7 @@ void ConstantAggregateBuilderBase::addSignedPointer(
 void CodeGenModule::destroyConstantSignedPointerCaches() {
   destroyCache<ByConstantCacheTy>(ConstantSignedPointersByConstant);
   destroyCache<FunctionPointerCacheTy>(SignedFunctionPointersByDeclAndType);
+  destroyCache<ByDeclCacheTy>(SignedThunkPointers);
 }
 
 /// If applicable, sign a given constant function pointer with the ABI rules for
@@ -869,6 +887,40 @@ llvm::Constant *CodeGenModule::getFunctionPointer(GlobalDecl GD,
   return getFunctionPointer(getRawFunctionPointer(GD, Ty), FuncType, GD);
 }
 
+llvm::Constant *
+CodeGenModule::getMemberFunctionPointer(llvm::Constant *pointer,
+                                        QualType functionType,
+                                        const FunctionDecl *FD) {
+  if (auto pointerAuth = getMemberFunctionPointerAuthInfo(functionType)) {
+    llvm::Constant **entry = nullptr;
+    if (FD) {
+      auto &cache =
+          getOrCreateCache<ByDeclCacheTy>(SignedThunkPointers);
+      entry = &cache[FD->getCanonicalDecl()];
+      if (*entry)
+        return llvm::ConstantExpr::getBitCast(*entry, pointer->getType());
+    }
+
+    pointer = getConstantSignedPointer(
+        pointer, pointerAuth.getKey(), nullptr,
+        cast_or_null<llvm::Constant>(pointerAuth.getDiscriminator()));
+
+    if (entry)
+      *entry = pointer;
+  }
+
+  return pointer;
+}
+
+llvm::Constant *
+CodeGenModule::getMemberFunctionPointer(const FunctionDecl *FD, llvm::Type *Ty) {
+  QualType functionType = FD->getType();
+  functionType = getContext().getMemberPointerType(
+      functionType, cast<CXXMethodDecl>(FD)->getParent()->getTypeForDecl());
+  return getMemberFunctionPointer(getRawFunctionPointer(FD, Ty), functionType,
+                                  FD);
+}
+
 llvm::Value *CodeGenFunction::AuthPointerToPointerCast(llvm::Value *ResultPtr,
                                                        QualType SourceType,
                                                        QualType DestType) {

clang/lib/CodeGen/CodeGenFunction.h

@@ -4173,7 +4173,8 @@ class CodeGenFunction : public CodeGenTypeCache {
   RValue EmitCall(const CGFunctionInfo &CallInfo, const CGCallee &Callee,
                   ReturnValueSlot ReturnValue, const CallArgList &Args,
                   llvm::CallBase **callOrInvoke, bool IsMustTail,
-                  SourceLocation Loc);
+                  SourceLocation Loc,
+                  bool IsVirtualFunctionPointerThunk = false);
   RValue EmitCall(const CGFunctionInfo &CallInfo, const CGCallee &Callee,
                   ReturnValueSlot ReturnValue, const CallArgList &Args,
                   llvm::CallBase **callOrInvoke = nullptr,

clang/lib/CodeGen/CodeGenModule.h

@@ -441,6 +441,7 @@ class CodeGenModule : public CodeGenTypeCache {
 
   /// Signed constant pointers.
   void *SignedFunctionPointersByDeclAndType = nullptr;
+  void *SignedThunkPointers = nullptr;
   void *ConstantSignedPointersByConstant = nullptr;
 
   llvm::StringMap<llvm::GlobalVariable *> CFConstantStringMap;
@@ -964,6 +965,13 @@ class CodeGenModule : public CodeGenTypeCache {
                                      QualType functionType,
                                      GlobalDecl GD = GlobalDecl());
 
+  llvm::Constant *getMemberFunctionPointer(const FunctionDecl *FD,
+                                           llvm::Type *Ty = nullptr);
+
+  llvm::Constant *getMemberFunctionPointer(llvm::Constant *pointer,
+                                           QualType functionType,
+                                           const FunctionDecl *FD = nullptr);
+
   CGPointerAuthInfo getFunctionPointerAuthInfo(QualType functionType);
 
   CGPointerAuthInfo getMemberFunctionPointerAuthInfo(QualType functionType);