Implement mangling/demangling __ptrauth on Windows

ahatanak · ahatanak · commit 45e5b93a556c · 2025-04-14T18:21:30.000-07:00
Add mangling/demangling tests.
diff --git a/clang/lib/AST/MicrosoftMangle.cpp b/clang/lib/AST/MicrosoftMangle.cpp
@@ -430,6 +430,7 @@ class MicrosoftCXXNameMangler {
   void mangleRefQualifier(RefQualifierKind RefQualifier);
   void manglePointerCVQualifiers(Qualifiers Quals);
   void manglePointerExtQualifiers(Qualifiers Quals, QualType PointeeType);
+  void manglePointerAuthQualifier(Qualifiers Quals);
 
   void mangleUnscopedTemplateName(GlobalDecl GD);
   void
@@ -2334,6 +2335,17 @@ void MicrosoftCXXNameMangler::manglePointerExtQualifiers(Qualifiers Quals,
     Out << 'F';
 }
 
+void MicrosoftCXXNameMangler::manglePointerAuthQualifier(Qualifiers Quals) {
+  PointerAuthQualifier PointerAuth = Quals.getPointerAuth();
+  if (!PointerAuth)
+    return;
+
+  Out << "__ptrauth";
+  mangleNumber(PointerAuth.getKey());
+  mangleNumber(PointerAuth.isAddressDiscriminated());
+  mangleNumber(PointerAuth.getExtraDiscriminator());
+}
+
 void MicrosoftCXXNameMangler::manglePointerCVQualifiers(Qualifiers Quals) {
   // <pointer-cv-qualifiers> ::= P  # no qualifiers
   //                         ::= Q  # const
@@ -3366,6 +3378,7 @@ void MicrosoftCXXNameMangler::mangleType(const PointerType *T, Qualifiers Quals,
   QualType PointeeType = T->getPointeeType();
   manglePointerCVQualifiers(Quals);
   manglePointerExtQualifiers(Quals, PointeeType);
+  manglePointerAuthQualifier(Quals);
 
   // For pointer size address spaces, go down the same type mangling path as
   // non address space types.
diff --git a/clang/test/CodeGenCXX/mangle-itanium-ptrauth.cpp b/clang/test/CodeGenCXX/mangle-itanium-ptrauth.cpp
@@ -0,0 +1,12 @@
+// RUN: %clang_cc1 -std=c++11 -fptrauth-intrinsics -fptrauth-calls -emit-llvm -o - -triple=arm64-apple-ios %s | FileCheck %s
+// RUN: %clang_cc1 -std=c++11 -fptrauth-intrinsics -fptrauth-calls -emit-llvm -o - -triple=aarch64-linux-gnu %s | FileCheck %s
+
+// CHECK: define {{.*}}void @_Z3fooPU9__ptrauthILj3ELb1ELj234EEPi(
+void foo(int * __ptrauth(3, 1, 234) *) {}
+
+template <class T>
+void foo(T t) {}
+
+// CHECK: define weak_odr void @_Z3fooIPU9__ptrauthILj1ELb0ELj64EEPiEvT_(
+template void foo<int * __ptrauth(1, 0, 64) *>(int * __ptrauth(1, 0, 64) *);
+
diff --git a/clang/test/CodeGenCXX/mangle-ms-ptrauth.cpp b/clang/test/CodeGenCXX/mangle-ms-ptrauth.cpp
@@ -0,0 +1,17 @@
+// RUN: %clang_cc1 -std=c++11 -fptrauth-intrinsics -fptrauth-calls -emit-llvm -o - -triple=aarch64-windows-msvc %s | FileCheck %s
+
+template <class T>
+struct S {};
+
+// CHECK: @"?s@@3U?$S@PE__ptrauth1A@ENC@AH@@A" =
+S<int * __ptrauth(2, 0, 1234)> s;
+
+// CHECK: define dso_local void @"?foo@@YAXPEAPE__ptrauth20OK@AH@Z"(
+void foo(int * __ptrauth(3, 1, 234) *) {}
+
+template <class T>
+void foo(T t) {}
+
+// CHECK: define weak_odr dso_local void @"??$foo@PEAPE__ptrauth0A@EA@AH@@YAXPEAPE__ptrauth0A@EA@AH@Z"(
+template void foo<int * __ptrauth(1, 0, 64) *>(int * __ptrauth(1, 0, 64) *);
+
diff --git a/libcxxabi/test/test_demangle.pass.cpp b/libcxxabi/test/test_demangle.pass.cpp
@@ -30247,6 +30247,9 @@ const char* cases[][2] = {
     {"_Z1fDSDRm", "f(_Sat unsigned long _Fract)"},
 
     {"_Z11bfloat16addDF16bDF16b", "bfloat16add(std::bfloat16_t, std::bfloat16_t)"},
+
+    {"_Z3fooPU9__ptrauthILj3ELb1ELj234EEPi", "foo(int* __ptrauth<3u, true, 234u>*)"},
+    {"_Z3fooIPU9__ptrauthILj1ELb0ELj64EEPiEvT_", "void foo<int* __ptrauth<1u, false, 64u>*>(int* __ptrauth<1u, false, 64u>*)"},
     // clang-format on
 };
 
diff --git a/llvm/include/llvm/Demangle/MicrosoftDemangle.h b/llvm/include/llvm/Demangle/MicrosoftDemangle.h
@@ -173,6 +173,14 @@ class Demangler {
 
   Qualifiers demanglePointerExtQualifiers(std::string_view &MangledName);
 
+  bool isMemberPointer(std::string_view MangledName, bool &Error);
+
+  std::optional<PointerAuthQualifierNode::ArgArray>
+  demanglePointerAuthQualifier(std::string_view &MangledName);
+
+  PointerAuthQualifierNode *
+  createPointerAuthQualifier(std::string_view &MangledName);
+
   // Parser functions. This is a recursive-descent parser.
   TypeNode *demangleType(std::string_view &MangledName,
                          QualifierMangleMode QMM);
diff --git a/llvm/include/llvm/Demangle/MicrosoftDemangleNodes.h b/llvm/include/llvm/Demangle/MicrosoftDemangleNodes.h
@@ -253,7 +253,8 @@ enum class NodeKind {
   LocalStaticGuardVariable,
   FunctionSymbol,
   VariableSymbol,
-  SpecialTableSymbol
+  SpecialTableSymbol,
+  PointerAuthQualifier,
 };
 
 struct Node {
@@ -295,6 +296,7 @@ struct SymbolNode;
 struct FunctionSymbolNode;
 struct VariableSymbolNode;
 struct SpecialTableSymbolNode;
+struct PointerAuthQualifierNode;
 
 struct TypeNode : public Node {
   explicit TypeNode(NodeKind K) : Node(K) {}
@@ -467,6 +469,8 @@ struct PointerTypeNode : public TypeNode {
   // If this is a member pointer, this is the class that the member is in.
   QualifiedNameNode *ClassParent = nullptr;
 
+  PointerAuthQualifierNode *PointerAuthQualifier = nullptr;
+
   // Represents a type X in "a pointer to X", "a reference to X", or
   // "rvalue-reference to X"
   TypeNode *Pointee = nullptr;
@@ -625,6 +629,22 @@ struct FunctionSymbolNode : public SymbolNode {
   FunctionSignatureNode *Signature = nullptr;
 };
 
+struct PointerAuthQualifierNode : public Node {
+  PointerAuthQualifierNode() : Node(NodeKind::PointerAuthQualifier) {}
+
+  // __ptrauth takes three arguments:
+  //  - key
+  //  - isAddressDiscriminated
+  //  - extra discriminator
+  static constexpr unsigned NumArgs = 3;
+  typedef std::array<uint64_t, NumArgs> ArgArray;
+
+  void output(OutputBuffer &OB, OutputFlags Flags) const override;
+
+  // List of arguments.
+  NodeArrayNode *Components = nullptr;
+};
+
 } // namespace ms_demangle
 } // namespace llvm
 
diff --git a/llvm/lib/Demangle/MicrosoftDemangle.cpp b/llvm/lib/Demangle/MicrosoftDemangle.cpp
@@ -66,7 +66,7 @@ static bool startsWith(std::string_view S, std::string_view PrefixA,
   return llvm::itanium_demangle::starts_with(S, Prefix);
 }
 
-static bool isMemberPointer(std::string_view MangledName, bool &Error) {
+bool Demangler::isMemberPointer(std::string_view MangledName, bool &Error) {
   Error = false;
   const char F = MangledName.front();
   MangledName.remove_prefix(1);
@@ -107,6 +107,7 @@ static bool isMemberPointer(std::string_view MangledName, bool &Error) {
   consumeFront(MangledName, 'E'); // 64-bit
   consumeFront(MangledName, 'I'); // restrict
   consumeFront(MangledName, 'F'); // unaligned
+  demanglePointerAuthQualifier(MangledName);
 
   if (MangledName.empty()) {
     Error = true;
@@ -2094,6 +2095,8 @@ PointerTypeNode *Demangler::demanglePointerType(std::string_view &MangledName) {
   Qualifiers ExtQuals = demanglePointerExtQualifiers(MangledName);
   Pointer->Quals = Qualifiers(Pointer->Quals | ExtQuals);
 
+  Pointer->PointerAuthQualifier = createPointerAuthQualifier(MangledName);
+
   Pointer->Pointee = demangleType(MangledName, QualifierMangleMode::Mangle);
   return Pointer;
 }
@@ -2142,6 +2145,49 @@ Demangler::demanglePointerExtQualifiers(std::string_view &MangledName) {
   return Quals;
 }
 
+std::optional<PointerAuthQualifierNode::ArgArray>
+Demangler::demanglePointerAuthQualifier(std::string_view &MangledName) {
+  if (!consumeFront(MangledName, "__ptrauth"))
+    return std::nullopt;
+
+  constexpr unsigned NumArgs = PointerAuthQualifierNode::NumArgs;
+  PointerAuthQualifierNode::ArgArray Array;
+
+  for (unsigned I = 0; I < NumArgs; ++I) {
+    bool IsNegative = false;
+    uint64_t Value = 0;
+    std::tie(Value, IsNegative) = demangleNumber(MangledName);
+    if (IsNegative)
+      return std::nullopt;
+
+    Array[I] = Value;
+  }
+
+  return Array;
+}
+
+PointerAuthQualifierNode *
+Demangler::createPointerAuthQualifier(std::string_view &MangledName) {
+  constexpr unsigned NumArgs = PointerAuthQualifierNode::NumArgs;
+  std::optional<PointerAuthQualifierNode::ArgArray> Vals =
+      demanglePointerAuthQualifier(MangledName);
+
+  if (!Vals)
+    return nullptr;
+
+  PointerAuthQualifierNode *PtrAuthQual =
+      Arena.alloc<PointerAuthQualifierNode>();
+  NodeArrayNode *Array = Arena.alloc<NodeArrayNode>();
+  PtrAuthQual->Components = Array;
+  Array->Count = NumArgs;
+  Array->Nodes = Arena.allocArray<Node *>(NumArgs);
+
+  for (unsigned I = 0; I < NumArgs; ++I)
+    Array->Nodes[I] = Arena.alloc<IntegerLiteralNode>((*Vals)[I], false);
+
+  return PtrAuthQual;
+}
+
 ArrayTypeNode *Demangler::demangleArrayType(std::string_view &MangledName) {
   assert(MangledName.front() == 'Y');
   MangledName.remove_prefix(1);
diff --git a/llvm/lib/Demangle/MicrosoftDemangleNodes.cpp b/llvm/lib/Demangle/MicrosoftDemangleNodes.cpp
@@ -521,6 +521,9 @@ void PointerTypeNode::outputPre(OutputBuffer &OB, OutputFlags Flags) const {
     assert(false);
   }
   outputQualifiers(OB, Quals, false, false);
+
+  if (PointerAuthQualifier)
+    PointerAuthQualifier->output(OB, Flags);
 }
 
 void PointerTypeNode::outputPost(OutputBuffer &OB, OutputFlags Flags) const {
@@ -591,6 +594,13 @@ void FunctionSymbolNode::output(OutputBuffer &OB, OutputFlags Flags) const {
   Signature->outputPost(OB, Flags);
 }
 
+void PointerAuthQualifierNode::output(OutputBuffer &OB,
+                                      OutputFlags Flags) const {
+  OB << "__ptrauth(";
+  Components->output(OB, Flags);
+  OB << ")";
+}
+
 void VariableSymbolNode::output(OutputBuffer &OB, OutputFlags Flags) const {
   const char *AccessSpec = nullptr;
   bool IsStatic = true;
diff --git a/llvm/test/Demangle/ms-ptrauth.test b/llvm/test/Demangle/ms-ptrauth.test
@@ -0,0 +1,12 @@
+; RUN: llvm-undname < %s | FileCheck %s
+
+; CHECK-NOT: Invalid mangled name
+
+?s@@3U?$S@PE__ptrauth1A@ENC@AH@@A
+; CHECK: struct S<int *__ptrauth(2, 0, 1234)> s
+
+?foo@@YAXPEAPE__ptrauth20OK@AH@Z
+; CHECK: void __cdecl foo(int *__ptrauth(3, 1, 234)*)
+
+??$foo@PEAPE__ptrauth0A@EA@AH@@YAXPEAPE__ptrauth0A@EA@AH@Z
+; CHECK: void __cdecl foo<int *__ptrauth(1, 0, 64)*>(int *__ptrauth(1, 0, 64)*)
