[PAC] Make ValueMapper handle ConstantPtrAuth values

[atrosinenko]

Fix assertion failure when building PAuth-hardened code with LTO.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Anatoly Trosinenko (atrosinenko)

Changes

Fix assertion failure when building PAuth-hardened code with LTO.

---

Full diff: https://github.com/llvm/llvm-project/pull/129088.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/Utils/ValueMapper.cpp (+3)
• (modified) llvm/unittests/Transforms/Utils/ValueMapperTest.cpp (+30)

diff --git a/llvm/lib/Transforms/Utils/ValueMapper.cpp b/llvm/lib/Transforms/Utils/ValueMapper.cpp
index b8569454379bf..9882f81b759ac 100644
--- a/llvm/lib/Transforms/Utils/ValueMapper.cpp
+++ b/llvm/lib/Transforms/Utils/ValueMapper.cpp
@@ -525,6 +525,9 @@ Value *Mapper::mapValue(const Value *V) {
     return getVM()[V] = ConstantStruct::get(cast<StructType>(NewTy), Ops);
   if (isa<ConstantVector>(C))
     return getVM()[V] = ConstantVector::get(Ops);
+  if (isa<ConstantPtrAuth>(C))
+    return getVM()[V] = ConstantPtrAuth::get(Ops[0], cast<ConstantInt>(Ops[1]),
+                                             cast<ConstantInt>(Ops[2]), Ops[3]);
   // If this is a no-operand constant, it must be because the type was remapped.
   if (isa<PoisonValue>(C))
     return getVM()[V] = PoisonValue::get(NewTy);
diff --git a/llvm/unittests/Transforms/Utils/ValueMapperTest.cpp b/llvm/unittests/Transforms/Utils/ValueMapperTest.cpp
index fb1b4edf25328..86ad41fa7ad50 100644
--- a/llvm/unittests/Transforms/Utils/ValueMapperTest.cpp
+++ b/llvm/unittests/Transforms/Utils/ValueMapperTest.cpp
@@ -436,4 +436,34 @@ TEST(ValueMapperTest, mapValueConstantTargetNoneToLayoutTypeNullValue) {
   EXPECT_EQ(NewConstant, Mapper.mapValue(*OldConstant));
 }
 
+TEST(ValueMapperTest, mapValuePtrAuth) {
+  LLVMContext C;
+  Type *PtrTy = PointerType::get(C, 0);
+  IntegerType *Int32Ty = Type::getInt32Ty(C);
+  IntegerType *Int64Ty = Type::getInt64Ty(C);
+
+  std::unique_ptr<GlobalVariable> Var0 = std::make_unique<GlobalVariable>(
+      PtrTy, false, GlobalValue::ExternalLinkage, nullptr, "Var0");
+  std::unique_ptr<GlobalVariable> Var1 = std::make_unique<GlobalVariable>(
+      PtrTy, false, GlobalValue::ExternalLinkage, nullptr, "Var1");
+  std::unique_ptr<GlobalVariable> Storage0 = std::make_unique<GlobalVariable>(
+      PtrTy, false, GlobalValue::ExternalLinkage, nullptr, "Storage0");
+  std::unique_ptr<GlobalVariable> Storage1 = std::make_unique<GlobalVariable>(
+      PtrTy, false, GlobalValue::ExternalLinkage, nullptr, "Storage1");
+
+  ConstantInt *ConstKey = ConstantInt::get(Int32Ty, 1);
+  ConstantInt *ConstDisc = ConstantInt::get(Int64Ty, 1234);
+
+  ValueToValueMapTy VM;
+  VM[Var0.get()] = Var1.get();
+  VM[Storage0.get()] = Storage1.get();
+
+  ConstantPtrAuth *Value =
+      ConstantPtrAuth::get(Var0.get(), ConstKey, ConstDisc, Storage0.get());
+  ConstantPtrAuth *MappedValue =
+      ConstantPtrAuth::get(Var1.get(), ConstKey, ConstDisc, Storage1.get());
+
+  EXPECT_EQ(ValueMapper(VM).mapValue(*Value), MappedValue);
+}
+
 } // end namespace

[asl]

I already happen to have this patch locally, but for bugpoint / llvm-reduce purposes.

However, it turned out that that it is required in LTO as mapper is used there. When assertions are enabled we're seeing assert being triggered at https://github.com/llvm/llvm-project/pull/129088/files#diff-8f546b2d38006292655c6fd7287a101d758f8f12be381f0180354431e387d3ceR540

But, in the release build w/o assert we're ending with constant zero instead producing the wrong code. The reproducer is as easy as:

// m.h:
// A caller may actually call this
void m2Export(const char* input);

// Called by m2
int export_foo(const char* str);

// Only called within m1
__attribute__((noinline)) int foo(int(*printfn)(const char* const, ...), const char* str);
int bar(const char* const, ...);

// m1.c
#include "m.h"

__attribute__((noinline)) int foo(int(*printfn)(const char* const, ...), const char* str) {
    printfn("Called foo");
    printfn(str);
    return 0;
}

int export_foo(const char* str) {
    return foo(bar, str);
}

// m2.c
#include <stdio.h>

#include "m.h"

void m2Export(const char* input) {
    printf("I will call export_foo\n");
    export_foo(input);
    printf("I have called export_foo\n");
}

// m3.c
#include <stdio.h>

int bar(const char* const str, ...) {
    printf("Bar is called with input: ");
    printf("%s", str);
    return 0;
}

m2Export (m2.c) calls export_foo (m1.c) which calls foo (m1.c, noinline attribute) which passes a function pointer argument bar (m3.c)

export_foo is inlined and instead foo is called directly, however instead of bar's address, a nullptr would be passed