[clang][ptrauth] Add support for querying the ptrauth schema of a type

This adds a number of builtins to query the ptrauth schema of an arbitrary
type in a way that can be fed into other ptrauth qualifiers.

Author: ojhunt

clang/docs/PointerAuthentication.rst

@@ -499,6 +499,24 @@ type.  Implementations are not required to make all bits of the result equally
 significant; in particular, some implementations are known to not leave
 meaningful data in the low bits.
 
+``__ptrauth type queries``
+^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+There are a number of builtins that can be used to query the ptrauth qualifier
+parameters of a type, including those configured implicitly. These are:
+
+.. code-block:: c
+__builtin_ptrauth_has_authentication(type)
+__builtin_ptrauth_schema_key(type)
+__builtin_ptrauth_schema_is_address_discriminated(type)
+__builtin_ptrauth_schema_extra_discriminator(type)
+__builtin_ptrauth_schema_options(type)
+
+All these builtins are compile time constants. The schema queries are only valid
+on types that have some form of pointer authentication, including implicit
+authentication as is present of function pointers. Each schema query returns a
+value of the appropriate type for the relevant parameter to the __ptrauth
+qualifier.
 
 
 Alternative Implementations

clang/include/clang/AST/ASTContext.h

@@ -1358,6 +1358,11 @@ class ASTContext : public RefCountedBase<ASTContext> {
   /// Return the "other" type-specific discriminator for the given type.
   uint16_t getPointerAuthTypeDiscriminator(QualType T);
 
+  /// Produces the canonical "options" string for the given PointerAuthQualifier
+  /// such that using it explicitly in a __ptrauth qualifier would result in as
+  /// identical configuration
+  std::string getPointerAuthOptionsString(const PointerAuthQualifier &PAQ);
+
   /// Apply Objective-C protocol qualifiers to the given type.
   /// \param allowOnPointerType specifies if we can apply protocol
   /// qualifiers on ObjCObjectPointerType. It can be set to true when
@@ -1696,6 +1701,13 @@ class ASTContext : public RefCountedBase<ASTContext> {
 
   QualType adjustStringLiteralBaseType(QualType StrLTy) const;
 
+  /// Synthesizes a PointerAuthQualifier representing the actual authentication
+  /// policy for the given type. This may be either the schema specified
+  /// explicitly via the __ptrauth qualified in the source, or the implicit
+  /// schema associated with function pointers and similar.
+  std::optional<PointerAuthQualifier>
+  getExplicitOrImplicitPointerAuth(QualType T);
+
 private:
   /// Return a normal function type with a typed argument list.
   QualType getFunctionTypeInternal(QualType ResultTy, ArrayRef<QualType> Args,

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -1039,6 +1039,8 @@ def err_ptrauth_address_discrimination_invalid : Error<
   "invalid address discrimination flag '%0'; '__ptrauth' requires '0' or '1'">;
 def err_ptrauth_extra_discriminator_invalid : Error<
   "invalid extra discriminator flag '%0'; '__ptrauth' requires a value between '0' and '%1'">;
+def err_ptrauth_query_type_has_no_pointer_authentication
+    : Error<"argument to %0 parameter is not an authenticated value">;
 
 /// main()
 // static main() is not an error in C, just in C++.

clang/include/clang/Basic/LangOptions.def

@@ -177,6 +177,7 @@ LANGOPT(PointerAuthInitFiniAddressDiscrimination, 1, 0,
         "incorporate address discrimination in authenticated function pointers in init/fini arrays")
 LANGOPT(PointerAuthELFGOT, 1, 0, "authenticate pointers from GOT")
 LANGOPT(AArch64JumpTableHardening, 1, 0, "use hardened lowering for jump-table dispatch")
+LANGOPT(PointerAuthFunctionKey, 16, 0, "authentication key for function pointers")
 
 LANGOPT(DoubleSquareBracketAttributes, 1, 0, "'[[]]' attributes extension for all language standard modes")
 LANGOPT(ExperimentalLateParseAttributes, 1, 0, "experimental late parsing of attributes")

clang/include/clang/Basic/LangOptions.h

@@ -65,6 +65,17 @@ enum class PointerAuthenticationMode : unsigned {
   SignAndAuth
 };
 
+static constexpr llvm::StringLiteral PointerAuthenticationOptionStrip = "strip";
+static constexpr llvm::StringLiteral PointerAuthenticationOptionSignAndStrip =
+    "sign-and-strip";
+static constexpr llvm::StringLiteral PointerAuthenticationOptionSignAndAuth =
+    "sign-and-auth";
+static constexpr llvm::StringLiteral PointerAuthenticationOptionIsaPointer =
+    "isa-pointer";
+static constexpr llvm::StringLiteral
+    PointerAuthenticationOptionAuthenticatesNullValues =
+        "authenticates-null-values";
+
 /// Bitfields of LangOptions, split out from LangOptions in order to ensure that
 /// this large collection of bitfields is a trivial class type.
 class LangOptionsBase {

clang/include/clang/Basic/TokenKinds.def

@@ -605,6 +605,11 @@ KEYWORD(__private_extern__          , KEYALL)
 KEYWORD(__module_private__          , KEYALL)
 
 UNARY_EXPR_OR_TYPE_TRAIT(__builtin_ptrauth_type_discriminator, PtrAuthTypeDiscriminator, KEYALL)
+UNARY_EXPR_OR_TYPE_TRAIT(__builtin_ptrauth_has_authentication, PtrAuthHasAuthentication, KEYALL)
+UNARY_EXPR_OR_TYPE_TRAIT(__builtin_ptrauth_schema_key, PtrAuthSchemaKey, KEYALL)
+UNARY_EXPR_OR_TYPE_TRAIT(__builtin_ptrauth_schema_is_address_discriminated, PtrAuthSchemaIsAddressDiscriminated, KEYALL)
+UNARY_EXPR_OR_TYPE_TRAIT(__builtin_ptrauth_schema_extra_discriminator, PtrAuthSchemaExtraDiscriminator, KEYALL)
+UNARY_EXPR_OR_TYPE_TRAIT(__builtin_ptrauth_schema_options, PtrAuthSchemaOptions, KEYALL)
 
 // Extension that will be enabled for Microsoft, Borland and PS4, but can be
 // disabled via '-fno-declspec'.

clang/include/clang/Parse/Parser.h

@@ -3981,7 +3981,9 @@ class Parser : public CodeCompletionHandler {
   ExprResult ParseArrayTypeTrait();
   ExprResult ParseExpressionTrait();
 
+  ExprResult ParseBuiltinUnaryExprOrTypeTrait(UnaryExprOrTypeTrait ExprKind);
   ExprResult ParseBuiltinPtrauthTypeDiscriminator();
+  ExprResult ParseBuiltinPtrauthQuery(tok::TokenKind Token);
 
   //===--------------------------------------------------------------------===//
   // Preprocessor code-completion pass-through

clang/lib/AST/ASTContext.cpp

@@ -96,6 +96,7 @@
 #include <map>
 #include <memory>
 #include <optional>
+#include <sstream>
 #include <string>
 #include <tuple>
 #include <utility>
@@ -9619,6 +9620,65 @@ ObjCInterfaceDecl *ASTContext::getObjCProtocolDecl() const {
   return ObjCProtocolClassDecl;
 }
 
+std::optional<PointerAuthQualifier>
+ASTContext::getExplicitOrImplicitPointerAuth(QualType T) {
+  auto ExplicitQualifier = T.getPointerAuth();
+  if (ExplicitQualifier.isPresent())
+    return ExplicitQualifier;
+  if (T->isDependentType()) {
+    return std::nullopt;
+  }
+  // FIXME: The more we expand pointer auth support, the more it becomes clear
+  // the codegen option's PointerAuth info belongs in LangOpts
+  if (!LangOpts.PointerAuthCalls)
+    return PointerAuthQualifier();
+  if (T->isFunctionPointerType() || T->isFunctionReferenceType())
+    T = T->getPointeeType();
+  if (!T->isFunctionType())
+    return PointerAuthQualifier();
+  int ExtraDiscriminator = 0;
+  if (LangOpts.PointerAuthFunctionTypeDiscrimination) {
+    ExtraDiscriminator = getPointerAuthTypeDiscriminator(T);
+  }
+  return PointerAuthQualifier::Create(
+      LangOpts.PointerAuthFunctionKey, false, ExtraDiscriminator,
+      PointerAuthenticationMode::SignAndAuth,
+      /*isIsaPointer=*/false, /*authenticatesNullValues=*/false);
+}
+
+std::string
+ASTContext::getPointerAuthOptionsString(const PointerAuthQualifier &PAQ) {
+  llvm::SmallVector<llvm::StringLiteral, 4> Options;
+  switch (PAQ.getAuthenticationMode()) {
+  case PointerAuthenticationMode::Strip:
+    Options.push_back(PointerAuthenticationOptionStrip);
+    break;
+  case PointerAuthenticationMode::SignAndStrip:
+    Options.push_back(PointerAuthenticationOptionSignAndStrip);
+    break;
+  case PointerAuthenticationMode::SignAndAuth:
+    // Just default to not listing this explicitly
+    break;
+  default:
+    llvm_unreachable("Invalid authentication mode");
+  }
+  if (PAQ.isIsaPointer())
+    Options.push_back(PointerAuthenticationOptionIsaPointer);
+  if (PAQ.authenticatesNullValues())
+    Options.push_back(PointerAuthenticationOptionAuthenticatesNullValues);
+  if (Options.empty())
+    return std::string();
+  if (Options.size() == 1)
+    return Options[0].str();
+  std::ostringstream Buffer;
+  Buffer << Options[0].str();
+  for (unsigned i = 1; i < Options.size(); i++) {
+    Buffer << ',';
+    Buffer << Options[i].str();
+  }
+  return Buffer.str();
+}
+
 //===----------------------------------------------------------------------===//
 // __builtin_va_list Construction Functions
 //===----------------------------------------------------------------------===//

clang/lib/AST/ExprConstant.cpp

@@ -9561,6 +9561,41 @@ class PointerExprEvaluator
     return true;
   }
 
+  bool VisitUnaryExprOrTypeTraitExpr(const UnaryExprOrTypeTraitExpr *E) {
+    // This is the only UETT we evaluate here.
+    assert(E->getKind() == UETT_PtrAuthSchemaOptions &&
+           "Unknown UnaryExprOrTypeTraitExpr");
+
+    // Note for review: there are other UETTs down the road
+    // that make a switch make sense, but for now this is the only
+    // one should this just be an
+    //   if (E->getKind() != UETT_PtrAuthSchemaOptions)
+    //     return false;
+    ASTContext &Ctx = Info.Ctx;
+    switch (E->getKind()) {
+    case UETT_PtrAuthSchemaOptions: {
+      auto ArgumentType = E->getArgumentType();
+      auto Qualifier = Ctx.getExplicitOrImplicitPointerAuth(ArgumentType);
+      if (!Qualifier)
+        return false;
+      if (!Qualifier->isPresent())
+        return false;
+      auto OptionsString = Ctx.getPointerAuthOptionsString(*Qualifier);
+      QualType StrTy =
+          Ctx.getStringLiteralArrayType(Ctx.CharTy, OptionsString.length());
+      StringLiteral *OptionsLit =
+          StringLiteral::Create(Ctx, OptionsString, StringLiteralKind::Ordinary,
+                                /*Pascal=*/false, StrTy, SourceLocation());
+      APValue OptionsVal(OptionsLit, CharUnits::Zero(),
+                         {APValue::LValuePathEntry::ArrayIndex(0)},
+                         /*OnePastTheEnd=*/false);
+      return Success(OptionsVal, E);
+    }
+    default:
+      return false;
+    }
+  }
+
   bool VisitSYCLUniqueStableNameExpr(const SYCLUniqueStableNameExpr *E) {
     std::string ResultStr = E->ComputeName(Info.Ctx);
 
@@ -14878,6 +14913,43 @@ bool IntExprEvaluator::VisitUnaryExprOrTypeTraitExpr(
     return Success(
         Info.Ctx.getPointerAuthTypeDiscriminator(E->getArgumentType()), E);
   }
+  case UETT_PtrAuthHasAuthentication: {
+    auto ArgumentType = E->getArgumentType();
+    auto Qualifier = Info.Ctx.getExplicitOrImplicitPointerAuth(ArgumentType);
+    if (!Qualifier)
+      return false;
+    return Success(Qualifier->isPresent(), E);
+  }
+  case UETT_PtrAuthSchemaKey: {
+    auto ArgumentType = E->getArgumentType();
+    auto Qualifier = Info.Ctx.getExplicitOrImplicitPointerAuth(ArgumentType);
+    if (!Qualifier)
+      return false;
+    if (!Qualifier->isPresent())
+      return false;
+    return Success(Qualifier->getKey(), E);
+  }
+  case UETT_PtrAuthSchemaIsAddressDiscriminated: {
+    auto ArgumentType = E->getArgumentType();
+    auto Qualifier = Info.Ctx.getExplicitOrImplicitPointerAuth(ArgumentType);
+    if (!Qualifier)
+      return false;
+    if (!Qualifier->isPresent())
+      return false;
+    return Success(Qualifier->isAddressDiscriminated(), E);
+  }
+  case UETT_PtrAuthSchemaExtraDiscriminator: {
+    auto ArgumentType = E->getArgumentType();
+    auto Qualifier = Info.Ctx.getExplicitOrImplicitPointerAuth(ArgumentType);
+    if (!Qualifier)
+      return false;
+    if (!Qualifier->isPresent())
+      return false;
+    return Success(Qualifier->getExtraDiscriminator(), E);
+  }
+  case UETT_PtrAuthSchemaOptions:
+    llvm_unreachable(
+        "UETT_PtrAuthSchemaOptions should be evaluated as a pointer");
   case UETT_VecStep: {
     QualType Ty = E->getTypeOfArgument();
 

clang/lib/AST/ItaniumMangle.cpp

@@ -5426,6 +5426,14 @@ void CXXNameMangler::mangleExpression(const Expr *E, unsigned Arity,
       Diags.Report(E->getExprLoc(), DiagID) << getTraitSpelling(SAE->getKind());
       return;
     }
+    case UETT_PtrAuthHasAuthentication:
+    case UETT_PtrAuthSchemaKey:
+    case UETT_PtrAuthSchemaIsAddressDiscriminated:
+    case UETT_PtrAuthSchemaExtraDiscriminator:
+    case UETT_PtrAuthSchemaOptions: {
+      MangleExtensionBuiltin(SAE);
+      break;
+    }
     }
     break;
   }

clang/lib/CodeGen/CGExprScalar.cpp

@@ -3586,6 +3586,16 @@ ScalarExprEmitter::VisitUnaryExprOrTypeTraitExpr(
   } else if (E->getKind() == UETT_VectorElements) {
     auto *VecTy = cast<llvm::VectorType>(ConvertType(E->getTypeOfArgument()));
     return Builder.CreateElementCount(CGF.SizeTy, VecTy->getElementCount());
+  } else if (E->getKind() == clang::UETT_PtrAuthSchemaOptions) {
+    auto PointerAuth =
+        CGF.getContext().getExplicitOrImplicitPointerAuth(E->getArgumentType());
+    assert(PointerAuth);
+    auto OptionsString =
+        CGF.getContext().getPointerAuthOptionsString(*PointerAuth);
+    ConstantAddress C =
+        CGF.CGM.GetAddrOfConstantCString(OptionsString, OptionsString.c_str());
+    return CGF.Builder.CreateBitCast(C.getPointer(),
+                                     CGF.getTypes().ConvertType(E->getType()));
   }
 
   // If this isn't sizeof(vla), the result must be constant; use the constant

clang/lib/Frontend/CompilerInvocation.cpp

@@ -1530,7 +1530,7 @@ void CompilerInvocation::setDefaultPointerAuthOptions(
     using Discrimination = PointerAuthSchema::Discrimination;
     // If you change anything here, be sure to update <ptrauth.h>.
     Opts.FunctionPointers = PointerAuthSchema(
-        Key::ASIA, false,
+        (Key)LangOpts.PointerAuthFunctionKey, false,
         LangOpts.PointerAuthFunctionTypeDiscrimination ? Discrimination::Type
                                                        : Discrimination::None);
 
@@ -3581,6 +3581,8 @@ static void ParsePointerAuthArgs(LangOptions &Opts, ArgList &Args,
   Opts.PointerAuthELFGOT = Args.hasArg(OPT_fptrauth_elf_got);
   Opts.AArch64JumpTableHardening =
       Args.hasArg(OPT_faarch64_jump_table_hardening);
+  Opts.PointerAuthFunctionKey =
+      static_cast<unsigned>(PointerAuthSchema::ARM8_3Key::ASIA);
 }
 
 /// Check if input file kind and language standard are compatible.

clang/lib/Parse/ParseExpr.cpp

@@ -861,7 +861,14 @@ bool Parser::isRevertibleTypeTrait(const IdentifierInfo *II,
   return false;
 }
 
-ExprResult Parser::ParseBuiltinPtrauthTypeDiscriminator() {
+ExprResult
+Parser::ParseBuiltinUnaryExprOrTypeTrait(UnaryExprOrTypeTrait ExprKind) {
+  assert(ExprKind == UETT_PtrAuthTypeDiscriminator ||
+         ExprKind == UETT_PtrAuthHasAuthentication ||
+         ExprKind == UETT_PtrAuthSchemaKey ||
+         ExprKind == UETT_PtrAuthSchemaIsAddressDiscriminated ||
+         ExprKind == UETT_PtrAuthSchemaExtraDiscriminator ||
+         ExprKind == UETT_PtrAuthSchemaOptions);
   SourceLocation Loc = ConsumeToken();
 
   BalancedDelimiterTracker T(*this, tok::l_paren);
@@ -877,10 +884,33 @@ ExprResult Parser::ParseBuiltinPtrauthTypeDiscriminator() {
   SourceLocation EndLoc = Tok.getLocation();
   T.consumeClose();
   return Actions.ActOnUnaryExprOrTypeTraitExpr(
-      Loc, UETT_PtrAuthTypeDiscriminator,
+      Loc, ExprKind,
       /*isType=*/true, Ty.get().getAsOpaquePtr(), SourceRange(Loc, EndLoc));
 }
 
+ExprResult Parser::ParseBuiltinPtrauthTypeDiscriminator() {
+  return ParseBuiltinUnaryExprOrTypeTrait(UETT_PtrAuthTypeDiscriminator);
+}
+
+ExprResult Parser::ParseBuiltinPtrauthQuery(tok::TokenKind Token) {
+  switch (Token) {
+  case tok::kw___builtin_ptrauth_has_authentication:
+    return ParseBuiltinUnaryExprOrTypeTrait(UETT_PtrAuthHasAuthentication);
+  case tok::kw___builtin_ptrauth_schema_key:
+    return ParseBuiltinUnaryExprOrTypeTrait(UETT_PtrAuthSchemaKey);
+  case tok::kw___builtin_ptrauth_schema_is_address_discriminated:
+    return ParseBuiltinUnaryExprOrTypeTrait(
+        UETT_PtrAuthSchemaIsAddressDiscriminated);
+  case tok::kw___builtin_ptrauth_schema_extra_discriminator:
+    return ParseBuiltinUnaryExprOrTypeTrait(
+        UETT_PtrAuthSchemaExtraDiscriminator);
+  case tok::kw___builtin_ptrauth_schema_options:
+    return ParseBuiltinUnaryExprOrTypeTrait(UETT_PtrAuthSchemaOptions);
+  default:
+    llvm_unreachable("Unexpected token");
+  }
+}
+
 /// Parse a cast-expression, or, if \pisUnaryExpression is true, parse
 /// a unary-expression.
 ///
@@ -1858,6 +1888,13 @@ ExprResult Parser::ParseCastExpression(CastParseKind ParseKind,
   case tok::kw___builtin_ptrauth_type_discriminator:
     return ParseBuiltinPtrauthTypeDiscriminator();
 
+  case tok::kw___builtin_ptrauth_has_authentication:
+  case tok::kw___builtin_ptrauth_schema_key:
+  case tok::kw___builtin_ptrauth_schema_is_address_discriminated:
+  case tok::kw___builtin_ptrauth_schema_extra_discriminator:
+  case tok::kw___builtin_ptrauth_schema_options:
+    return ParseBuiltinPtrauthQuery(SavedKind);
+
   case tok::kw___is_lvalue_expr:
   case tok::kw___is_rvalue_expr:
     if (NotPrimaryExpression)