[PAC] Implement function pointer type discrimination (#96992)

Give users an option (-fptrauth-function-pointer-type-discrimination) to
sign a function pointer using a non-zero discriminator based on the
function type.

The discriminator is computed by first translating the function type to
a string and then computing the hash value of the string. Two function
types that are compatible in C must be translated to the same string
with the exception of function types that use typedefs of anonymous
structs in their return type or parameter types.

This patch doesn't have the code to resign function pointers, which is
needed when a function pointer is converted to a different function
type. That will be implemented in another patch.

Co-authored-by: John McCall <rjmccall@apple.com>

---------

Co-authored-by: John McCall <rjmccall@apple.com>

4 files changed, 96 insertions, 16 deletions

diff --git a/clang/lib/CodeGen/CGExprConstant.cpp b/clang/lib/CodeGen/CGExprConstant.cpp
index 1fec587..00a5a7e 100644
--- a/clang/lib/CodeGen/CGExprConstant.cpp
+++ b/clang/lib/CodeGen/CGExprConstant.cpp
@@ -1542,9 +1542,11 @@ ConstantEmitter::emitAbstract(const Expr *E, QualType destType) {
 
 llvm::Constant *
 ConstantEmitter::emitAbstract(SourceLocation loc, const APValue &value,
-                              QualType destType) {
+                              QualType destType,
+                              bool EnablePtrAuthFunctionTypeDiscrimination) {
   auto state = pushAbstract();
-  auto C = tryEmitPrivate(value, destType);
+  auto C =
+      tryEmitPrivate(value, destType, EnablePtrAuthFunctionTypeDiscrimination);
   C = validateAndPopAbstract(C, state);
   if (!C) {
     CGM.Error(loc,
@@ -1938,14 +1940,18 @@ class ConstantLValueEmitter : public ConstStmtVisitor<ConstantLValueEmitter,
   ConstantEmitter &Emitter;
   const APValue &Value;
   QualType DestType;
+  bool EnablePtrAuthFunctionTypeDiscrimination;
 
   // Befriend StmtVisitorBase so that we don't have to expose Visit*.
   friend StmtVisitorBase;
 
 public:
   ConstantLValueEmitter(ConstantEmitter &emitter, const APValue &value,
-                        QualType destType)
-    : CGM(emitter.CGM), Emitter(emitter), Value(value), DestType(destType) {}
+                        QualType destType,
+                        bool EnablePtrAuthFunctionTypeDiscrimination = true)
+      : CGM(emitter.CGM), Emitter(emitter), Value(value), DestType(destType),
+        EnablePtrAuthFunctionTypeDiscrimination(
+            EnablePtrAuthFunctionTypeDiscrimination) {}
 
   llvm::Constant *tryEmit();
 
@@ -2069,7 +2075,10 @@ ConstantLValueEmitter::tryEmitBase(const APValue::LValueBase &base) {
       return CGM.GetWeakRefReference(D).getPointer();
 
     auto PtrAuthSign = [&](llvm::Constant *C) {
-      CGPointerAuthInfo AuthInfo = CGM.getFunctionPointerAuthInfo(DestType);
+      CGPointerAuthInfo AuthInfo;
+
+      if (EnablePtrAuthFunctionTypeDiscrimination)
+        AuthInfo = CGM.getFunctionPointerAuthInfo(DestType);
 
       if (AuthInfo) {
         if (hasNonZeroOffset())
@@ -2220,8 +2229,10 @@ llvm::Constant *ConstantLValueEmitter::emitPointerAuthPointer(const Expr *E) {
 
   // The assertions here are all checked by Sema.
   assert(Result.Val.isLValue());
+  if (isa<FunctionDecl>(Result.Val.getLValueBase().get<const ValueDecl *>()))
+    assert(Result.Val.getLValueOffset().isZero());
   return ConstantEmitter(CGM, Emitter.CGF)
-      .emitAbstract(E->getExprLoc(), Result.Val, E->getType());
+      .emitAbstract(E->getExprLoc(), Result.Val, E->getType(), false);
 }
 
 unsigned ConstantLValueEmitter::emitPointerAuthKey(const Expr *E) {
@@ -2278,15 +2289,18 @@ ConstantLValueEmitter::VisitMaterializeTemporaryExpr(
   return CGM.GetAddrOfGlobalTemporary(E, Inner);
 }
 
-llvm::Constant *ConstantEmitter::tryEmitPrivate(const APValue &Value,
-                                                QualType DestType) {
+llvm::Constant *
+ConstantEmitter::tryEmitPrivate(const APValue &Value, QualType DestType,
+                                bool EnablePtrAuthFunctionTypeDiscrimination) {
   switch (Value.getKind()) {
   case APValue::None:
   case APValue::Indeterminate:
     // Out-of-lifetime and indeterminate values can be modeled as 'undef'.
     return llvm::UndefValue::get(CGM.getTypes().ConvertType(DestType));
   case APValue::LValue:
-    return ConstantLValueEmitter(*this, Value, DestType).tryEmit();
+    return ConstantLValueEmitter(*this, Value, DestType,
+                                 EnablePtrAuthFunctionTypeDiscrimination)
+        .tryEmit();
   case APValue::Int:
     return llvm::ConstantInt::get(CGM.getLLVMContext(), Value.getInt());
   case APValue::FixedPoint:
diff --git a/clang/lib/CodeGen/CGPointerAuth.cpp b/clang/lib/CodeGen/CGPointerAuth.cpp
index 673f6e6..621d567 100644
--- a/clang/lib/CodeGen/CGPointerAuth.cpp
+++ b/clang/lib/CodeGen/CGPointerAuth.cpp
@@ -29,7 +29,9 @@ llvm::ConstantInt *CodeGenModule::getPointerAuthOtherDiscriminator(
     return nullptr;
 
   case PointerAuthSchema::Discrimination::Type:
-    llvm_unreachable("type discrimination not implemented yet");
+    assert(!Type.isNull() && "type not provided for type-discriminated schema");
+    return llvm::ConstantInt::get(
+        IntPtrTy, getContext().getPointerAuthTypeDiscriminator(Type));
 
   case PointerAuthSchema::Discrimination::Decl:
     assert(Decl.getDecl() &&
@@ -43,6 +45,11 @@ llvm::ConstantInt *CodeGenModule::getPointerAuthOtherDiscriminator(
   llvm_unreachable("bad discrimination kind");
 }
 
+uint16_t CodeGen::getPointerAuthTypeDiscriminator(CodeGenModule &CGM,
+                                                  QualType FunctionType) {
+  return CGM.getContext().getPointerAuthTypeDiscriminator(FunctionType);
+}
+
 uint16_t CodeGen::getPointerAuthDeclDiscriminator(CodeGenModule &CGM,
                                                   GlobalDecl Declaration) {
   return CGM.getPointerAuthDeclDiscriminator(Declaration);
@@ -71,12 +78,15 @@ CGPointerAuthInfo CodeGenModule::getFunctionPointerAuthInfo(QualType T) {
   assert(!Schema.isAddressDiscriminated() &&
          "function pointers cannot use address-specific discrimination");
 
-  assert(!Schema.hasOtherDiscrimination() &&
-         "function pointers don't support any discrimination yet");
+  llvm::Constant *Discriminator = nullptr;
+  if (T->isFunctionPointerType() || T->isFunctionReferenceType())
+    T = T->getPointeeType();
+  if (T->isFunctionType())
+    Discriminator = getPointerAuthOtherDiscriminator(Schema, GlobalDecl(), T);
 
   return CGPointerAuthInfo(Schema.getKey(), Schema.getAuthenticationMode(),
                            /*IsaPointer=*/false, /*AuthenticatesNull=*/false,
-                           /*Discriminator=*/nullptr);
+                           Discriminator);
 }
 
 llvm::Value *
@@ -114,6 +124,47 @@ CGPointerAuthInfo CodeGenFunction::EmitPointerAuthInfo(
                            Schema.authenticatesNullValues(), Discriminator);
 }
 
+/// Return the natural pointer authentication for values of the given
+/// pointee type.
+static CGPointerAuthInfo
+getPointerAuthInfoForPointeeType(CodeGenModule &CGM, QualType PointeeType) {
+  if (PointeeType.isNull())
+    return CGPointerAuthInfo();
+
+  // Function pointers use the function-pointer schema by default.
+  if (PointeeType->isFunctionType())
+    return CGM.getFunctionPointerAuthInfo(PointeeType);
+
+  // Normal data pointers never use direct pointer authentication by default.
+  return CGPointerAuthInfo();
+}
+
+CGPointerAuthInfo CodeGenModule::getPointerAuthInfoForPointeeType(QualType T) {
+  return ::getPointerAuthInfoForPointeeType(*this, T);
+}
+
+/// Return the natural pointer authentication for values of the given
+/// pointer type.
+static CGPointerAuthInfo getPointerAuthInfoForType(CodeGenModule &CGM,
+                                                   QualType PointerType) {
+  assert(PointerType->isSignableType());
+
+  // Block pointers are currently not signed.
+  if (PointerType->isBlockPointerType())
+    return CGPointerAuthInfo();
+
+  auto PointeeType = PointerType->getPointeeType();
+
+  if (PointeeType.isNull())
+    return CGPointerAuthInfo();
+
+  return ::getPointerAuthInfoForPointeeType(CGM, PointeeType);
+}
+
+CGPointerAuthInfo CodeGenModule::getPointerAuthInfoForType(QualType T) {
+  return ::getPointerAuthInfoForType(*this, T);
+}
+
 llvm::Constant *
 CodeGenModule::getConstantSignedPointer(llvm::Constant *Pointer, unsigned Key,
                                         llvm::Constant *StorageAddress,
@@ -180,6 +231,14 @@ llvm::Constant *CodeGenModule::getFunctionPointer(GlobalDecl GD,
                                                   llvm::Type *Ty) {
   const auto *FD = cast<FunctionDecl>(GD.getDecl());
   QualType FuncType = FD->getType();
+
+  // Annoyingly, K&R functions have prototypes in the clang AST, but
+  // expressions referring to them are unprototyped.
+  if (!FD->hasPrototype())
+    if (const auto *Proto = FuncType->getAs<FunctionProtoType>())
+      FuncType = Context.getFunctionNoProtoType(Proto->getReturnType(),
+                                                Proto->getExtInfo());
+
   return getFunctionPointer(getRawFunctionPointer(GD, Ty), FuncType);
 }
 
diff --git a/clang/lib/CodeGen/CodeGenModule.h b/clang/lib/CodeGen/CodeGenModule.h
index 8b65348..585c4ea 100644
--- a/clang/lib/CodeGen/CodeGenModule.h
+++ b/clang/lib/CodeGen/CodeGenModule.h
@@ -972,6 +972,10 @@ public:
 
   CGPointerAuthInfo getFunctionPointerAuthInfo(QualType T);
 
+  CGPointerAuthInfo getPointerAuthInfoForPointeeType(QualType type);
+
+  CGPointerAuthInfo getPointerAuthInfoForType(QualType type);
+
   bool shouldSignPointer(const PointerAuthSchema &Schema);
   llvm::Constant *getConstantSignedPointer(llvm::Constant *Pointer,
                                            const PointerAuthSchema &Schema,
diff --git a/clang/lib/CodeGen/ConstantEmitter.h b/clang/lib/CodeGen/ConstantEmitter.h
index eff0a8d..581b05a 100644
--- a/clang/lib/CodeGen/ConstantEmitter.h
+++ b/clang/lib/CodeGen/ConstantEmitter.h
@@ -103,8 +103,9 @@ public:
   /// expression is known to be a constant expression with either a fairly
   /// simple type or a known simple form.
   llvm::Constant *emitAbstract(const Expr *E, QualType T);
-  llvm::Constant *emitAbstract(SourceLocation loc, const APValue &value,
-                               QualType T);
+  llvm::Constant *
+  emitAbstract(SourceLocation loc, const APValue &value, QualType T,
+               bool EnablePtrAuthFunctionTypeDiscrimination = true);
 
   /// Try to emit the result of the given expression as an abstract constant.
   llvm::Constant *tryEmitAbstract(const Expr *E, QualType T);
@@ -138,7 +139,9 @@ public:
   llvm::Constant *tryEmitPrivate(const Expr *E, QualType T);
   llvm::Constant *tryEmitPrivateForMemory(const Expr *E, QualType T);
 
-  llvm::Constant *tryEmitPrivate(const APValue &value, QualType T);
+  llvm::Constant *
+  tryEmitPrivate(const APValue &value, QualType T,
+                 bool EnablePtrAuthFunctionTypeDiscrimination = true);
   llvm::Constant *tryEmitPrivateForMemory(const APValue &value, QualType T);
 
   /// Get the address of the current location.  This is a constant