From c46a8cde6ae34903add4dacad5765e467438b78d Mon Sep 17 00:00:00 2001
From: Jacek Caban <jacek@codeweavers.com>
Date: Thu, 7 Sep 2023 01:32:39 +0200
Subject: [PATCH] [LLD][COFF] Always locate the IAT at the beginning of the
 .rdata section and align its size to 4KB on ARM64EC.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This mimics the behavior of MSVC's link.exe. My guess is that the reason for this approach
is to facilitate tracking runtime IAT modifications. An auxiliary IAT allows bypassing
the call checker for imported function calls. It's the OS's responsibility to ensure that,
if runtime patching occurs, the auxiliary IAT is reverted to enable call checking.
Modifying the IAT is a form of runtime patching, and ensuring that it doesn’t share
pages with other data likely helps with tracking accuracy.

Although alignment alone should ensure that the IAT occupies its own pages, placing it
at the beginning of the .rdata section might be an optimization. This way, padding is
only needed after the IAT, not before. The auxiliary IAT seems to follow a similar
idea but is positioned at the end of the .rdata section.
---
 lld/COFF/Writer.cpp               | 29 +++++++++++++++++++++++++++++
 lld/test/COFF/arm64ec-import.test | 15 ++++++++++-----
 2 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/lld/COFF/Writer.cpp b/lld/COFF/Writer.cpp
index 4a0eed4d00997..3cb9b3b512ead 100644
--- a/lld/COFF/Writer.cpp
+++ b/lld/COFF/Writer.cpp
@@ -217,6 +217,7 @@ class Writer {
   void createExportTable();
   void mergeSections();
   void sortECChunks();
+  void appendECImportTables();
   void removeUnusedSections();
   void assignAddresses();
   bool isInRange(uint16_t relType, uint64_t s, uint64_t p, int margin,
@@ -753,6 +754,7 @@ void Writer::run() {
     createExportTable();
     mergeSections();
     sortECChunks();
+    appendECImportTables();
     removeUnusedSections();
     finalizeAddresses();
     removeEmptySections();
@@ -914,6 +916,28 @@ void Writer::addSyntheticIdata() {
   add(".idata$7", idata.dllNames);
 }
 
+void Writer::appendECImportTables() {
+  if (!isArm64EC(ctx.config.machine))
+    return;
+
+  const uint32_t rdata = IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ;
+
+  // IAT is always placed at the begining of .rdata section and its size
+  // is aligned to 4KB. Insert it here, after all merges all done.
+  if (PartialSection *importAddresses = findPartialSection(".idata$5", rdata)) {
+    if (!rdataSec->chunks.empty())
+      rdataSec->chunks.front()->setAlignment(
+          std::max(0x1000u, rdataSec->chunks.front()->getAlignment()));
+    iatSize = alignTo(iatSize, 0x1000);
+
+    rdataSec->chunks.insert(rdataSec->chunks.begin(),
+                            importAddresses->chunks.begin(),
+                            importAddresses->chunks.end());
+    rdataSec->contribSections.insert(rdataSec->contribSections.begin(),
+                                     importAddresses);
+  }
+}
+
 // Locate the first Chunk and size of the import directory list and the
 // IAT.
 void Writer::locateImportTables() {
@@ -1069,6 +1093,11 @@ void Writer::createSections() {
       sortCRTSectionChunks(pSec->chunks);
     }
 
+    // ARM64EC has specific placement and alignment requirements for the IAT.
+    // Delay adding its chunks until appendECImportTables.
+    if (isArm64EC(ctx.config.machine) && pSec->name == ".idata$5")
+      continue;
+
     OutputSection *sec = createSection(name, outChars);
     for (Chunk *c : pSec->chunks)
       sec->addChunk(c);
diff --git a/lld/test/COFF/arm64ec-import.test b/lld/test/COFF/arm64ec-import.test
index 2a80a30910b5c..b1c47d785e445 100644
--- a/lld/test/COFF/arm64ec-import.test
+++ b/lld/test/COFF/arm64ec-import.test
@@ -20,7 +20,7 @@ RUN: llvm-readobj --coff-imports out2.dll | FileCheck --check-prefix=IMPORTS %s
 IMPORTS:      Import {
 IMPORTS-NEXT:   Name: test.dll
 IMPORTS-NEXT:   ImportLookupTableRVA:
-IMPORTS-NEXT:   ImportAddressTableRVA: 0x2258
+IMPORTS-NEXT:   ImportAddressTableRVA: 0x2000
 IMPORTS-NEXT:   Symbol: data (0)
 IMPORTS-NEXT:   Symbol: func (0)
 IMPORTS-NEXT:   Symbol: func2 (0)
@@ -28,7 +28,7 @@ IMPORTS-NEXT: }
 IMPORTS-NEXT: Import {
 IMPORTS-NEXT:   Name: test2.dll
 IMPORTS-NEXT:   ImportLookupTableRVA:
-IMPORTS-NEXT:   ImportAddressTableRVA: 0x2278
+IMPORTS-NEXT:   ImportAddressTableRVA: 0x2020
 IMPORTS-NEXT:   Symbol: t2func (0)
 IMPORTS-NEXT: }
 
@@ -36,12 +36,17 @@ RUN: llvm-objdump -d out.dll | FileCheck --check-prefix=DISASM %s
 RUN: llvm-objdump -d out2.dll | FileCheck --check-prefix=DISASM %s
 
 DISASM:      0000000180001000 <.text>:
-DISASM-NEXT: 180001000: ff 25 5a 12 00 00            jmpq    *0x125a(%rip)           # 0x180002260
+DISASM-NEXT: 180001000: ff 25 02 10 00 00            jmpq    *0x1002(%rip)           # 0x180002008
 
 RUN: llvm-readobj --hex-dump=.test out.dll | FileCheck --check-prefix=TESTSEC %s
 RUN: llvm-readobj --hex-dump=.test out2.dll | FileCheck --check-prefix=TESTSEC %s
-TESTSEC:      0x180004000 60220000 58220000 68220000 78220000
-TESTSEC-NEXT: 0x180004010 00100000
+TESTSEC:      0x180005000 08200000 00200000 10200000 20200000
+TESTSEC-NEXT: 0x180005010 00100000
+
+RUN: llvm-readobj --headers out.dll | FileCheck -check-prefix=HEADERS %s
+HEADERS:  LoadConfigTableRVA: 0x3008
+HEADERS:  IATRVA: 0x2000
+HEADERS:  IATSize: 0x1000
 
 #--- test.s
     .section .test, "r"