[LLD][COFF] Add basic ARM64X dynamic relocations support

cjacek · cjacek · commit 51a677242843 · 2024-11-28T22:48:54.000+01:00
This update modifies the machine field in the hybrid view to be AMD64, aligning it with
expectations from ARM64EC modules. While this provides initial support, additional
relocations will be necessary for full functionality. Many of these cases depend on
implementing separate namespace support first.
diff --git a/lld/COFF/COFFLinkerContext.h b/lld/COFF/COFFLinkerContext.h
@@ -88,6 +88,8 @@ class COFFLinkerContext : public CommonLinkerContext {
   Timer diskCommitTimer;
 
   Configuration config;
+
+  DynamicRelocsChunk *dynamicRelocs = nullptr;
 };
 
 } // namespace lld::coff
diff --git a/lld/COFF/Chunks.cpp b/lld/COFF/Chunks.cpp
@@ -25,6 +25,7 @@
 
 using namespace llvm;
 using namespace llvm::object;
+using namespace llvm::support;
 using namespace llvm::support::endian;
 using namespace llvm::COFF;
 using llvm::support::ulittle32_t;
@@ -1147,4 +1148,85 @@ uint32_t ImportThunkChunkARM64EC::extendRanges() {
   return sizeof(arm64Thunk) - sizeof(uint32_t);
 }
 
+size_t Arm64XDynamicRelocEntry::getSize() const {
+  switch (type) {
+  case IMAGE_DVRT_ARM64X_FIXUP_TYPE_VALUE:
+    return sizeof(uint16_t) + size; // A header and a payload.
+  case IMAGE_DVRT_ARM64X_FIXUP_TYPE_DELTA:
+  case IMAGE_DVRT_ARM64X_FIXUP_TYPE_ZEROFILL:
+    llvm_unreachable("unsupported type");
+  }
+}
+
+void Arm64XDynamicRelocEntry::writeTo(uint8_t *buf) const {
+  auto out = reinterpret_cast<ulittle16_t *>(buf);
+  *out = (offset & 0xfff) | (type << 12);
+
+  switch (type) {
+  case IMAGE_DVRT_ARM64X_FIXUP_TYPE_VALUE:
+    *out |= ((bit_width(size) - 1) << 14); // Encode the size.
+    switch (size) {
+    case 2:
+      out[1] = value;
+      break;
+    case 4:
+      *reinterpret_cast<ulittle32_t *>(out + 1) = value;
+      break;
+    case 8:
+      *reinterpret_cast<ulittle64_t *>(out + 1) = value;
+      break;
+    default:
+      llvm_unreachable("invalid size");
+    }
+    break;
+  case IMAGE_DVRT_ARM64X_FIXUP_TYPE_DELTA:
+  case IMAGE_DVRT_ARM64X_FIXUP_TYPE_ZEROFILL:
+    llvm_unreachable("unsupported type");
+  }
+}
+
+void DynamicRelocsChunk::finalize() {
+  llvm::stable_sort(arm64xRelocs, [=](const Arm64XDynamicRelocEntry &a,
+                                      const Arm64XDynamicRelocEntry &b) {
+    return a.offset < b.offset;
+  });
+
+  size = sizeof(coff_dynamic_reloc_table) + sizeof(coff_dynamic_relocation64) +
+         sizeof(coff_base_reloc_block_header);
+
+  for (const Arm64XDynamicRelocEntry &entry : arm64xRelocs) {
+    assert(!(entry.offset & ~0xfff)); // Not yet supported.
+    size += entry.getSize();
+  }
+
+  size = alignTo(size, sizeof(uint32_t));
+}
+
+void DynamicRelocsChunk::writeTo(uint8_t *buf) const {
+  auto table = reinterpret_cast<coff_dynamic_reloc_table *>(buf);
+  table->Version = 1;
+  table->Size = sizeof(coff_dynamic_relocation64);
+  buf += sizeof(*table);
+
+  auto header = reinterpret_cast<coff_dynamic_relocation64 *>(buf);
+  header->Symbol = IMAGE_DYNAMIC_RELOCATION_ARM64X;
+  buf += sizeof(*header);
+
+  auto pageHeader = reinterpret_cast<coff_base_reloc_block_header *>(buf);
+  pageHeader->BlockSize = sizeof(*pageHeader);
+  size_t relocSize = sizeof(*pageHeader);
+  for (const Arm64XDynamicRelocEntry &entry : arm64xRelocs) {
+    entry.writeTo(buf + relocSize);
+    size_t entrySize = entry.getSize();
+    pageHeader->BlockSize += entrySize;
+    relocSize += entrySize;
+  }
+  pageHeader->BlockSize = alignTo(pageHeader->BlockSize, sizeof(uint32_t));
+  relocSize = alignTo(relocSize, sizeof(uint32_t));
+
+  header->BaseRelocSize = relocSize;
+  table->Size += relocSize;
+  assert(size == sizeof(*table) + sizeof(*header) + relocSize);
+}
+
 } // namespace lld::coff
diff --git a/lld/COFF/Chunks.h b/lld/COFF/Chunks.h
@@ -835,6 +835,43 @@ class ECExportThunkChunk : public NonSectionCodeChunk {
   Defined *target;
 };
 
+// ARM64X entry for dynamic relocations.
+class Arm64XDynamicRelocEntry {
+public:
+  Arm64XDynamicRelocEntry(llvm::COFF::Arm64XFixupType type, uint8_t size,
+                          uint32_t offset, uint64_t value)
+      : offset(offset), value(value), type(type), size(size) {}
+
+  size_t getSize() const;
+  void writeTo(uint8_t *buf) const;
+
+  uint32_t offset;
+  uint64_t value;
+
+private:
+  llvm::COFF::Arm64XFixupType type;
+  uint8_t size;
+};
+
+// Dynamic relocation chunk containing ARM64X relocations for the hybrid image.
+class DynamicRelocsChunk : public NonSectionChunk {
+public:
+  DynamicRelocsChunk() {}
+  size_t getSize() const override { return size; }
+  void writeTo(uint8_t *buf) const override;
+  void finalize();
+
+  uint32_t add(llvm::COFF::Arm64XFixupType type, uint8_t size, uint32_t offset,
+               uint64_t value) {
+    arm64xRelocs.emplace_back(type, size, offset, value);
+    return arm64xRelocs.size() - 1;
+  }
+
+private:
+  std::vector<Arm64XDynamicRelocEntry> arm64xRelocs;
+  size_t size;
+};
+
 // MinGW specific, for the "automatic import of variables from DLLs" feature.
 // This provides the table of runtime pseudo relocations, for variable
 // references that turned out to need to be imported from a DLL even though
diff --git a/lld/COFF/Writer.cpp b/lld/COFF/Writer.cpp
@@ -272,12 +272,12 @@ class Writer {
   OutputSection *findSection(StringRef name);
   void addBaserels();
   void addBaserelBlocks(std::vector<Baserel> &v);
+  void createDynamicRelocs();
 
   uint32_t getSizeOfInitializedData();
 
   void prepareLoadConfig();
   template <typename T> void prepareLoadConfig(T *loadConfig);
-  template <typename T> void checkLoadConfigGuardData(const T *loadConfig);
 
   std::unique_ptr<FileOutputBuffer> &buffer;
   std::map<PartialSectionKey, PartialSection *> partialSections;
@@ -754,6 +754,8 @@ void Writer::run() {
     llvm::TimeTraceScope timeScope("Write PE");
     ScopedTimer t1(ctx.codeLayoutTimer);
 
+    if (ctx.config.machine == ARM64X)
+      ctx.dynamicRelocs = make<DynamicRelocsChunk>();
     createImportTables();
     createSections();
     appendImportThunks();
@@ -764,6 +766,7 @@ void Writer::run() {
     mergeSections();
     sortECChunks();
     appendECImportTables();
+    createDynamicRelocs();
     removeUnusedSections();
     finalizeAddresses();
     removeEmptySections();
@@ -1596,8 +1599,13 @@ void Writer::assignAddresses() {
 
   for (OutputSection *sec : ctx.outputSections) {
     llvm::TimeTraceScope timeScope("Section: ", sec->name);
-    if (sec == relocSec)
+    if (sec == relocSec) {
       addBaserels();
+      if (ctx.dynamicRelocs) {
+        ctx.dynamicRelocs->finalize();
+        relocSec->addChunk(ctx.dynamicRelocs);
+      }
+    }
     uint64_t rawSize = 0, virtualSize = 0;
     sec->header.VirtualAddress = rva;
 
@@ -1798,9 +1806,12 @@ template <typename PEHeaderTy> void Writer::writeHeader() {
                                 exceptionTable.last->getSize() -
                                 exceptionTable.first->getRVA();
   }
-  if (relocSec->getVirtualSize()) {
+  size_t relocSize = relocSec->getVirtualSize();
+  if (ctx.dynamicRelocs)
+    relocSize -= ctx.dynamicRelocs->getSize();
+  if (relocSize) {
     dir[BASE_RELOCATION_TABLE].RelativeVirtualAddress = relocSec->getRVA();
-    dir[BASE_RELOCATION_TABLE].Size = relocSec->getVirtualSize();
+    dir[BASE_RELOCATION_TABLE].Size = relocSize;
   }
   if (Symbol *sym = ctx.symtab.findUnderscore("_tls_used")) {
     if (Defined *b = dyn_cast<Defined>(sym)) {
@@ -2555,6 +2566,33 @@ void Writer::addBaserelBlocks(std::vector<Baserel> &v) {
   relocSec->addChunk(make<BaserelChunk>(page, &v[i], &v[0] + j));
 }
 
+void Writer::createDynamicRelocs() {
+  if (!ctx.dynamicRelocs)
+    return;
+
+  const uint32_t coffHeaderOffset = dosStubSize + sizeof(PEMagic);
+  const uint32_t peHeaderOffset = coffHeaderOffset + sizeof(coff_file_header);
+  const uint32_t dataDirOffset = peHeaderOffset + sizeof(pe32plus_header);
+
+  // Adjust the Machine field in the COFF header to AMD64.
+  ctx.dynamicRelocs->add(IMAGE_DVRT_ARM64X_FIXUP_TYPE_VALUE, sizeof(uint16_t),
+                         coffHeaderOffset + offsetof(coff_file_header, Machine),
+                         AMD64);
+
+  // Adjust the load config directory.
+  // FIXME: Use the hybrid load config value instead.
+  ctx.dynamicRelocs->add(IMAGE_DVRT_ARM64X_FIXUP_TYPE_VALUE, sizeof(uint32_t),
+                         dataDirOffset +
+                             LOAD_CONFIG_TABLE * sizeof(data_directory) +
+                             offsetof(data_directory, RelativeVirtualAddress),
+                         0);
+  ctx.dynamicRelocs->add(IMAGE_DVRT_ARM64X_FIXUP_TYPE_VALUE, sizeof(uint32_t),
+                         dataDirOffset +
+                             LOAD_CONFIG_TABLE * sizeof(data_directory) +
+                             offsetof(data_directory, Size),
+                         0);
+}
+
 PartialSection *Writer::createPartialSection(StringRef name,
                                              uint32_t outChars) {
   PartialSection *&pSec = partialSections[{name, outChars}];
@@ -2634,11 +2672,6 @@ template <typename T> void Writer::prepareLoadConfig(T *loadConfig) {
   if (ctx.config.dependentLoadFlags)
     loadConfig->DependentLoadFlags = ctx.config.dependentLoadFlags;
 
-  checkLoadConfigGuardData(loadConfig);
-}
-
-template <typename T>
-void Writer::checkLoadConfigGuardData(const T *loadConfig) {
   size_t loadConfigSize = loadConfig->Size;
 
 #define RETURN_IF_NOT_CONTAINS(field)                                          \
@@ -2660,6 +2693,18 @@ void Writer::checkLoadConfigGuardData(const T *loadConfig) {
     if (loadConfig->field != s->getVA())                                       \
       warn(#field " not set correctly in '_load_config_used'");
 
+  if (ctx.dynamicRelocs) {
+    IF_CONTAINS(DynamicValueRelocTableSection) {
+      loadConfig->DynamicValueRelocTableSection = relocSec->sectionIndex;
+      loadConfig->DynamicValueRelocTableOffset =
+          ctx.dynamicRelocs->getRVA() - relocSec->getRVA();
+    }
+    else {
+      warn("'_load_config_used' structure too small to include dynamic "
+           "relocations");
+    }
+  }
+
   if (ctx.config.guardCF == GuardCFLevel::Off)
     return;
   RETURN_IF_NOT_CONTAINS(GuardFlags)
diff --git a/lld/test/COFF/arm64x-loadconfig.s b/lld/test/COFF/arm64x-loadconfig.s
@@ -0,0 +1,68 @@
+// REQUIRES: aarch64
+// RUN: split-file %s %t.dir && cd %t.dir
+
+// RUN: llvm-mc -filetype=obj -triple=aarch64-windows test.s -o test.obj
+// RUN: llvm-mc -filetype=obj -triple=aarch64-windows loadconfig.s -o loadconfig.obj
+// RUN: llvm-mc -filetype=obj -triple=aarch64-windows loadconfig-short.s -o loadconfig-short.obj
+
+// RUN: lld-link -machine:arm64x -out:out.dll -dll -noentry loadconfig.obj test.obj
+
+// RUN: llvm-readobj --coff-load-config out.dll | FileCheck -check-prefix=DYNRELOCS %s
+// DYNRELOCS:      DynamicValueRelocTableOffset: 0xC
+// DYNRELOCS-NEXT: DynamicValueRelocTableSection: 4
+// DYNRELOCS:      DynamicRelocations [
+// DYNRELOCS-NEXT:   Version: 0x1
+// DYNRELOCS-NEXT:   Arm64X [
+// DYNRELOCS-NEXT:     Entry [
+// DYNRELOCS-NEXT:       RVA: 0x7C
+// DYNRELOCS-NEXT:       Type: VALUE
+// DYNRELOCS-NEXT:       Size: 0x2
+// DYNRELOCS-NEXT:       Value: 0x8664
+// DYNRELOCS-NEXT:     ]
+// DYNRELOCS-NEXT:     Entry [
+// DYNRELOCS-NEXT:       RVA: 0x150
+// DYNRELOCS-NEXT:       Type: VALUE
+// DYNRELOCS-NEXT:       Size: 0x4
+// DYNRELOCS-NEXT:       Value: 0x0
+// DYNRELOCS-NEXT:     ]
+// DYNRELOCS-NEXT:     Entry [
+// DYNRELOCS-NEXT:       RVA: 0x154
+// DYNRELOCS-NEXT:       Type: VALUE
+// DYNRELOCS-NEXT:       Size: 0x4
+// DYNRELOCS-NEXT:       Value: 0x0
+// DYNRELOCS-NEXT:     ]
+// DYNRELOCS-NEXT:   ]
+// DYNRELOCS-NEXT: ]
+
+// RUN: llvm-readobj --headers out.dll | FileCheck -check-prefix=HEADERS %s
+// HEADERS:      BaseRelocationTableRVA: 0x4000
+// HEADERS-NEXT: BaseRelocationTableSize: 0xC
+// HEADERS:       LoadConfigTableRVA: 0x1000
+// HEADERS-NEXT:  LoadConfigTableSize: 0x140
+// HEADERS:       Name: .reloc (2E 72 65 6C 6F 63 00 00)
+// HEADERS-NEXT:  VirtualSize: 0x38
+
+// RUN: lld-link -machine:arm64x -out:out-short.dll -dll -noentry loadconfig-short.obj 2>&1 | FileCheck --check-prefix=WARN-RELOC-SIZE %s
+// WARN-RELOC-SIZE: lld-link: warning: '_load_config_used' structure too small to include dynamic relocations
+
+#--- test.s
+        .data
+sym:
+        // Emit a basereloc to make the loadconfig test more meaningful.
+        .xword sym
+
+#--- loadconfig.s
+        .section .rdata,"dr"
+        .globl _load_config_used
+        .p2align 3, 0
+_load_config_used:
+        .word 0x140
+        .fill 0x13c,1,0
+
+#--- loadconfig-short.s
+        .section .rdata,"dr"
+        .globl _load_config_used
+        .p2align 3, 0
+_load_config_used:
+        .word 0xe4
+        .fill 0xe0,1,0
