[SandboxIR] Implement ConstantPtrAuth (#109315)

This patch implements sandboxir::ConstantPtrAuth mirroring
llvm::ConstantPtrAuth

Author: vporpo

llvm/include/llvm/SandboxIR/SandboxIR.h

@@ -134,6 +134,7 @@ class GlobalIFunc;
 class GlobalVariable;
 class GlobalAlias;
 class NoCFIValue;
+class ConstantPtrAuth;
 class Context;
 class Function;
 class Instruction;
@@ -342,6 +343,7 @@ class Value {
   friend class GlobalVariable;        // For `Val`.
   friend class GlobalAlias;           // For `Val`.
   friend class NoCFIValue;            // For `Val`.
+  friend class ConstantPtrAuth;       // For `Val`.
 
   /// All values point to the context.
   Context &Ctx;
@@ -1603,6 +1605,62 @@ class NoCFIValue final : public Constant {
 #endif
 };
 
+class ConstantPtrAuth final : public Constant {
+  ConstantPtrAuth(llvm::ConstantPtrAuth *C, Context &Ctx)
+      : Constant(ClassID::ConstantPtrAuth, C, Ctx) {}
+  friend class Context; // For constructor.
+
+public:
+  /// Return a pointer signed with the specified parameters.
+  static ConstantPtrAuth *get(Constant *Ptr, ConstantInt *Key,
+                              ConstantInt *Disc, Constant *AddrDisc);
+  /// The pointer that is signed in this ptrauth signed pointer.
+  Constant *getPointer() const;
+
+  /// The Key ID, an i32 constant.
+  ConstantInt *getKey() const;
+
+  /// The integer discriminator, an i64 constant, or 0.
+  ConstantInt *getDiscriminator() const;
+
+  /// The address discriminator if any, or the null constant.
+  /// If present, this must be a value equivalent to the storage location of
+  /// the only global-initializer user of the ptrauth signed pointer.
+  Constant *getAddrDiscriminator() const;
+
+  /// Whether there is any non-null address discriminator.
+  bool hasAddressDiscriminator() const {
+    return cast<llvm::ConstantPtrAuth>(Val)->hasAddressDiscriminator();
+  }
+
+  /// Whether the address uses a special address discriminator.
+  /// These discriminators can't be used in real pointer-auth values; they
+  /// can only be used in "prototype" values that indicate how some real
+  /// schema is supposed to be produced.
+  bool hasSpecialAddressDiscriminator(uint64_t Value) const {
+    return cast<llvm::ConstantPtrAuth>(Val)->hasSpecialAddressDiscriminator(
+        Value);
+  }
+
+  /// Check whether an authentication operation with key \p Key and (possibly
+  /// blended) discriminator \p Discriminator is known to be compatible with
+  /// this ptrauth signed pointer.
+  bool isKnownCompatibleWith(const Value *Key, const Value *Discriminator,
+                             const DataLayout &DL) const {
+    return cast<llvm::ConstantPtrAuth>(Val)->isKnownCompatibleWith(
+        Key->Val, Discriminator->Val, DL);
+  }
+
+  /// Produce a new ptrauth expression signing the given value using
+  /// the same schema as is stored in one.
+  ConstantPtrAuth *getWithSameSchema(Constant *Pointer) const;
+
+  /// For isa/dyn_cast.
+  static bool classof(const sandboxir::Value *From) {
+    return From->getSubclassID() == ClassID::ConstantPtrAuth;
+  }
+};
+
 class BlockAddress final : public Constant {
   BlockAddress(llvm::BlockAddress *C, Context &Ctx)
       : Constant(ClassID::BlockAddress, C, Ctx) {}

llvm/include/llvm/SandboxIR/SandboxIRValues.def

@@ -39,6 +39,7 @@ DEF_CONST(GlobalIFunc, GlobalIFunc)
 DEF_CONST(GlobalAlias, GlobalAlias)
 DEF_CONST(BlockAddress, BlockAddress)
 DEF_CONST(NoCFIValue, NoCFIValue)
+DEF_CONST(ConstantPtrAuth, ConstantPtrAuth)
 DEF_CONST(DSOLocalEquivalent, DSOLocalEquivalent)
 DEF_CONST(ConstantTokenNone, ConstantTokenNone)
 

llvm/lib/SandboxIR/SandboxIR.cpp

@@ -2640,6 +2640,40 @@ PointerType *NoCFIValue::getType() const {
   return cast<PointerType>(Ctx.getType(cast<llvm::NoCFIValue>(Val)->getType()));
 }
 
+ConstantPtrAuth *ConstantPtrAuth::get(Constant *Ptr, ConstantInt *Key,
+                                      ConstantInt *Disc, Constant *AddrDisc) {
+  auto *LLVMC = llvm::ConstantPtrAuth::get(
+      cast<llvm::Constant>(Ptr->Val), cast<llvm::ConstantInt>(Key->Val),
+      cast<llvm::ConstantInt>(Disc->Val), cast<llvm::Constant>(AddrDisc->Val));
+  return cast<ConstantPtrAuth>(Ptr->getContext().getOrCreateConstant(LLVMC));
+}
+
+Constant *ConstantPtrAuth::getPointer() const {
+  return Ctx.getOrCreateConstant(
+      cast<llvm::ConstantPtrAuth>(Val)->getPointer());
+}
+
+ConstantInt *ConstantPtrAuth::getKey() const {
+  return cast<ConstantInt>(
+      Ctx.getOrCreateConstant(cast<llvm::ConstantPtrAuth>(Val)->getKey()));
+}
+
+ConstantInt *ConstantPtrAuth::getDiscriminator() const {
+  return cast<ConstantInt>(Ctx.getOrCreateConstant(
+      cast<llvm::ConstantPtrAuth>(Val)->getDiscriminator()));
+}
+
+Constant *ConstantPtrAuth::getAddrDiscriminator() const {
+  return Ctx.getOrCreateConstant(
+      cast<llvm::ConstantPtrAuth>(Val)->getAddrDiscriminator());
+}
+
+ConstantPtrAuth *ConstantPtrAuth::getWithSameSchema(Constant *Pointer) const {
+  auto *LLVMC = cast<llvm::ConstantPtrAuth>(Val)->getWithSameSchema(
+      cast<llvm::Constant>(Pointer->Val));
+  return cast<ConstantPtrAuth>(Ctx.getOrCreateConstant(LLVMC));
+}
+
 BlockAddress *BlockAddress::get(Function *F, BasicBlock *BB) {
   auto *LLVMC = llvm::BlockAddress::get(cast<llvm::Function>(F->Val),
                                         cast<llvm::BasicBlock>(BB->Val));
@@ -2850,6 +2884,10 @@ Value *Context::getOrCreateValueInternal(llvm::Value *LLVMV, llvm::User *U) {
       It->second = std::unique_ptr<NoCFIValue>(
           new NoCFIValue(cast<llvm::NoCFIValue>(C), *this));
       break;
+    case llvm::Value::ConstantPtrAuthVal:
+      It->second = std::unique_ptr<ConstantPtrAuth>(
+          new ConstantPtrAuth(cast<llvm::ConstantPtrAuth>(C), *this));
+      break;
     default:
       It->second = std::unique_ptr<Constant>(new Constant(C, *this));
       break;

llvm/unittests/SandboxIR/SandboxIRTest.cpp

@@ -1142,6 +1142,43 @@ define void @foo() {
   EXPECT_EQ(NoCFI->getType(), F.getType());
 }
 
+TEST_F(SandboxIRTest, ConstantPtrAuth) {
+  parseIR(C, R"IR(
+define ptr @foo() {
+  ret ptr ptrauth (ptr @foo, i32 2, i64 1234)
+}
+)IR");
+  Function &LLVMF = *M->getFunction("foo");
+  auto *LLVMBB = &*LLVMF.begin();
+  auto *LLVMRet = cast<llvm::ReturnInst>(&*LLVMBB->begin());
+  auto *LLVMPtrAuth = cast<llvm::ConstantPtrAuth>(LLVMRet->getReturnValue());
+  sandboxir::Context Ctx(C);
+
+  auto &F = *Ctx.createFunction(&LLVMF);
+  auto *BB = &*F.begin();
+  auto It = BB->begin();
+  auto *Ret = cast<sandboxir::ReturnInst>(&*It++);
+  // Check classof(), creation.
+  auto *PtrAuth = cast<sandboxir::ConstantPtrAuth>(Ret->getReturnValue());
+  // Check get(), getKey(), getDiscriminator(), getAddrDiscriminator().
+  auto *NewPtrAuth = sandboxir::ConstantPtrAuth::get(
+      &F, PtrAuth->getKey(), PtrAuth->getDiscriminator(),
+      PtrAuth->getAddrDiscriminator());
+  EXPECT_EQ(NewPtrAuth, PtrAuth);
+  // Check hasAddressDiscriminator().
+  EXPECT_EQ(PtrAuth->hasAddressDiscriminator(),
+            LLVMPtrAuth->hasAddressDiscriminator());
+  // Check hasSpecialAddressDiscriminator().
+  EXPECT_EQ(PtrAuth->hasSpecialAddressDiscriminator(0u),
+            LLVMPtrAuth->hasSpecialAddressDiscriminator(0u));
+  // Check isKnownCompatibleWith().
+  const DataLayout &DL = M->getDataLayout();
+  EXPECT_TRUE(PtrAuth->isKnownCompatibleWith(PtrAuth->getKey(),
+                                             PtrAuth->getDiscriminator(), DL));
+  // Check getWithSameSchema().
+  EXPECT_EQ(PtrAuth->getWithSameSchema(&F), PtrAuth);
+}
+
 TEST_F(SandboxIRTest, BlockAddress) {
   parseIR(C, R"IR(
 define void @foo(ptr %ptr) {