[IR] Autoupgrade llvm.ptrauth globals to ptrauth constant.

Author: ahmedbougacha

llvm/include/llvm/IR/AutoUpgrade.h

@@ -55,7 +55,7 @@ namespace llvm {
 
   /// This checks for global variables which should be upgraded. If it requires
   /// upgrading, returns a pointer to the upgraded variable.
-  GlobalVariable *UpgradeGlobalVariable(GlobalVariable *GV);
+  bool UpgradeGlobalVariable(GlobalVariable *GV, GlobalVariable *&NewGV);
 
   /// This checks for module flags which should be upgraded. It returns true if
   /// module is modified.

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

@@ -3977,12 +3977,15 @@ Error BitcodeReader::globalCleanup() {
 
   // Look for global variables which need to be renamed.
   std::vector<std::pair<GlobalVariable *, GlobalVariable *>> UpgradedVariables;
-  for (GlobalVariable &GV : TheModule->globals())
-    if (GlobalVariable *Upgraded = UpgradeGlobalVariable(&GV))
+  for (GlobalVariable &GV : TheModule->globals()) {
+    GlobalVariable *Upgraded = nullptr;
+    if (UpgradeGlobalVariable(&GV, Upgraded))
       UpgradedVariables.emplace_back(&GV, Upgraded);
+  }
   for (auto &Pair : UpgradedVariables) {
     Pair.first->eraseFromParent();
-    TheModule->insertGlobalVariable(Pair.second);
+    if (Pair.second)
+      TheModule->insertGlobalVariable(Pair.second);
   }
 
   // Force deallocation of memory for these vectors to favor the client that

llvm/lib/IR/AutoUpgrade.cpp

@@ -1521,17 +1521,62 @@ bool llvm::UpgradeIntrinsicFunction(Function *F, Function *&NewFn,
   return Upgraded;
 }
 
-GlobalVariable *llvm::UpgradeGlobalVariable(GlobalVariable *GV) {
+bool llvm::UpgradeGlobalVariable(GlobalVariable *GV, GlobalVariable *&NewGV) {
+  if (GV->getSection() == "llvm.ptrauth") {
+    auto &Ctx = GV->getContext();
+
+    if (!GV->hasInitializer())
+      return false;
+
+    auto *Init = GV->getInitializer();
+    auto *Ty = dyn_cast<StructType>(GV->getInitializer()->getType());
+    if (!Ty)
+      return false;
+
+    auto *I64Ty = Type::getInt64Ty(Ctx);
+    auto *I32Ty = Type::getInt32Ty(Ctx);
+    auto *PtrTy = PointerType::get(Ctx, 0);
+    // Check that the struct matches its expected shape:
+    //   { ptr, i32, i64, i64 }
+    if (!Ty->isLayoutIdentical(
+            StructType::get(Ctx, {PtrTy, I32Ty, I64Ty, I64Ty})))
+      return false;
+
+    auto Pointer = cast<Constant>(Init->getOperand(0));
+    Pointer = Pointer->stripPointerCasts();
+
+    auto *Key = dyn_cast<ConstantInt>(Init->getOperand(1));
+    if (!Key)
+      return false;
+
+    auto *AddrDiscriminator = cast<Constant>(Init->getOperand(2));
+    if (!isa<ConstantInt>(AddrDiscriminator) &&
+        !isa<ConstantExpr>(AddrDiscriminator))
+      return false;
+
+    auto *Discriminator = dyn_cast<ConstantInt>(Init->getOperand(3));
+    if (!Discriminator)
+      return false;
+
+    AddrDiscriminator = ConstantExpr::getIntToPtr(AddrDiscriminator, PtrTy);
+    Discriminator = ConstantInt::get(I64Ty, Discriminator->getZExtValue());
+    auto *SP =
+        ConstantPtrAuth::get(Pointer, Key, Discriminator, AddrDiscriminator);
+
+    GV->replaceAllUsesWith(ConstantExpr::getBitCast(SP, GV->getType()));
+    return true;
+  }
+
   if (!(GV->hasName() && (GV->getName() == "llvm.global_ctors" ||
                           GV->getName() == "llvm.global_dtors")) ||
       !GV->hasInitializer())
-    return nullptr;
+    return false;
   ArrayType *ATy = dyn_cast<ArrayType>(GV->getValueType());
   if (!ATy)
-    return nullptr;
+    return false;
   StructType *STy = dyn_cast<StructType>(ATy->getElementType());
   if (!STy || STy->getNumElements() != 2)
-    return nullptr;
+    return false;
 
   LLVMContext &C = GV->getContext();
   IRBuilder<> IRB(C);
@@ -1548,8 +1593,9 @@ GlobalVariable *llvm::UpgradeGlobalVariable(GlobalVariable *GV) {
   }
   Constant *NewInit = ConstantArray::get(ArrayType::get(EltTy, N), NewCtors);
 
-  return new GlobalVariable(NewInit->getType(), false, GV->getLinkage(),
-                            NewInit, GV->getName());
+  NewGV = new GlobalVariable(NewInit->getType(), false, GV->getLinkage(),
+                             NewInit, GV->getName());
+  return true;
 }
 
 // Handles upgrading SSE2/AVX2/AVX512BW PSLLDQ intrinsics by converting them

llvm/test/Bitcode/upgrade-ptrauth-global.ll

@@ -0,0 +1,20 @@
+; RUN: llvm-as < %s | llvm-dis | FileCheck %s
+
+@var = global i32 0
+
+@var.ptrauth1 = constant { i8*, i32, i64, i64 } { i8* bitcast(i32* @var to i8*),
+                                               i32 0,
+                                               i64 0,
+                                               i64 1234 }, section "llvm.ptrauth"
+@var_auth1 = global i32* bitcast({i8*, i32, i64, i64}* @var.ptrauth1 to i32*)
+; CHECK: @var_auth1 = global ptr ptrauth (ptr @var, i32 0, i64 1234)
+
+
+@dummy_addrdisc = global i64* null
+
+@var.ptrauth2 = constant { i8*, i32, i64, i64 } { i8* bitcast(i32* @var to i8*),
+                                                  i32 2,
+                                                  i64 ptrtoint(i64** @dummy_addrdisc to i64),
+                                                  i64 5678 }, section "llvm.ptrauth"
+@var_auth2 = global i32* bitcast({i8*, i32, i64, i64}* @var.ptrauth2 to i32*)
+; CHECK: @var_auth2 = global ptr ptrauth (ptr @var, i32 2, i64 5678, ptr @dummy_addrdisc)