[LLD][COFF] Add Support for auxiliary IAT copy (#108610)

In addition to the auxiliary IAT, ARM64EC modules also contain a copy of
it. At runtime, the auxiliary IAT is filled with the addresses of actual
ARM64EC functions when possible. If patching is detected, the OS may use
the IAT copy to revert the auxiliary IAT, ensuring that the call checker
is used for calls to imported functions.

Author: cjacek

lld/COFF/DLL.cpp

@@ -730,13 +730,19 @@ void IdataContents::create(COFFLinkerContext &ctx) {
         auto chunk = make<AuxImportChunk>(s->file);
         auxIat.push_back(chunk);
         s->file->impECSym->setLocation(chunk);
+
+        chunk = make<AuxImportChunk>(s->file);
+        auxIatCopy.push_back(chunk);
+        s->file->auxImpCopySym->setLocation(chunk);
       }
     }
     // Terminate with null values.
     lookups.push_back(make<NullChunk>(ctx.config.wordsize));
     addresses.push_back(make<NullChunk>(ctx.config.wordsize));
-    if (ctx.config.machine == ARM64EC)
+    if (ctx.config.machine == ARM64EC) {
       auxIat.push_back(make<NullChunk>(ctx.config.wordsize));
+      auxIatCopy.push_back(make<NullChunk>(ctx.config.wordsize));
+    }
 
     for (int i = 0, e = syms.size(); i < e; ++i)
       syms[i]->setLocation(addresses[base + i]);

lld/COFF/DLL.h

@@ -32,6 +32,7 @@ class IdataContents {
   std::vector<Chunk *> hints;
   std::vector<Chunk *> dllNames;
   std::vector<Chunk *> auxIat;
+  std::vector<Chunk *> auxIatCopy;
 };
 
 // Windows-specific.

lld/COFF/Driver.cpp

@@ -2448,6 +2448,7 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
     ctx.symtab.addAbsolute("__arm64x_redirection_metadata", 0);
     ctx.symtab.addAbsolute("__arm64x_redirection_metadata_count", 0);
     ctx.symtab.addAbsolute("__hybrid_auxiliary_iat", 0);
+    ctx.symtab.addAbsolute("__hybrid_auxiliary_iat_copy", 0);
     ctx.symtab.addAbsolute("__hybrid_code_map", 0);
     ctx.symtab.addAbsolute("__hybrid_code_map_count", 0);
     ctx.symtab.addAbsolute("__x64_code_ranges_to_entry_points", 0);

lld/COFF/InputFiles.cpp

@@ -1091,6 +1091,12 @@ void ImportFile::parse() {
     }
     if (!impECSym)
       return;
+
+    StringRef auxImpCopyName = saver().save("__auximpcopy_" + name);
+    auxImpCopySym =
+        ctx.symtab.addImportData(auxImpCopyName, this, auxCopyLocation);
+    if (!auxImpCopySym)
+      return;
   }
   // If this was a duplicate, we logged an error but may continue;
   // in this case, impSym is nullptr.

lld/COFF/InputFiles.h

@@ -362,10 +362,12 @@ class ImportFile : public InputFile {
   const coff_import_header *hdr;
   Chunk *location = nullptr;
 
-  // Auxiliary IAT symbol and chunk on ARM64EC.
+  // Auxiliary IAT symbols and chunks on ARM64EC.
   DefinedImportData *impECSym = nullptr;
   Chunk *auxLocation = nullptr;
   Symbol *auxThunkSym = nullptr;
+  DefinedImportData *auxImpCopySym = nullptr;
+  Chunk *auxCopyLocation = nullptr;
 
   // We want to eliminate dllimported symbols if no one actually refers to them.
   // These "Live" bits are used to keep track of which import library members

lld/COFF/Writer.cpp

@@ -916,6 +916,8 @@ void Writer::addSyntheticIdata() {
   add(".idata$7", idata.dllNames);
   if (!idata.auxIat.empty())
     add(".idata$9", idata.auxIat);
+  if (!idata.auxIatCopy.empty())
+    add(".idata$a", idata.auxIatCopy);
 }
 
 void Writer::appendECImportTables() {
@@ -2279,6 +2281,11 @@ void Writer::setECSymbols() {
   replaceSymbol<DefinedSynthetic>(iatSym, "__hybrid_auxiliary_iat",
                                   idata.auxIat.empty() ? nullptr
                                                        : idata.auxIat.front());
+
+  Symbol *iatCopySym = ctx.symtab.findUnderscore("__hybrid_auxiliary_iat_copy");
+  replaceSymbol<DefinedSynthetic>(
+      iatCopySym, "__hybrid_auxiliary_iat_copy",
+      idata.auxIatCopy.empty() ? nullptr : idata.auxIatCopy.front());
 }
 
 // Write section contents to a mmap'ed file.

lld/test/COFF/Inputs/loadconfig-arm64ec.s

@@ -84,7 +84,7 @@ __chpe_metadata:
         .rva __arm64x_extra_rfe_table
         .word __arm64x_extra_rfe_table_size
         .rva __os_arm64x_dispatch_fptr
-        .word 0 // __hybrid_auxiliary_iat_copy
+        .rva __hybrid_auxiliary_iat_copy
         .rva __os_arm64x_helper0
         .rva __os_arm64x_helper1
         .rva __os_arm64x_helper2

lld/test/COFF/arm64ec-import.test

@@ -27,7 +27,7 @@ RUN: llvm-readobj --coff-imports out2.dll | FileCheck --check-prefix=IMPORTS %s
 RUN: llvm-readobj --coff-imports out3.dll | FileCheck -check-prefix=IMPORTS %s
 IMPORTS:      Import {
 IMPORTS-NEXT:   Name: test.dll
-IMPORTS-NEXT:   ImportLookupTableRVA:
+IMPORTS-NEXT:   ImportLookupTableRVA: 0x4230
 IMPORTS-NEXT:   ImportAddressTableRVA: 0x3000
 IMPORTS-NEXT:   Symbol: data (0)
 IMPORTS-NEXT:   Symbol: func (0)
@@ -79,12 +79,13 @@ RUN: llvm-readobj --hex-dump=.test out2.dll | FileCheck --check-prefix=TESTSEC %
 TESTSEC:      0x180007000 08500000 00300000 10500000 20500000
 TESTSEC-NEXT: 0x180007010 08300000 00500000 10300000 20300000
 TESTSEC-NEXT: 0x180007020 14100000 28100000 00200000 08100000
-TESTSEC-NEXT: 0x180007030 3c100000
+TESTSEC-NEXT: 0x180007030 3c100000 a0420000
 
 RUN: llvm-readobj --hex-dump=.test out3.dll | FileCheck -check-prefix=TESTSEC-X64 %s
 TESTSEC-X64:      0x180007000 08300000 00300000 10300000 20300000
 TESTSEC-X64-NEXT: 0x180007010 08300000 00500000 10300000 20300000
 TESTSEC-X64-NEXT: 0x180007020 14100000 28100000 00200000 08100000
+TESTSEC-X64-NEXT: 0x180007030 3c100000 a0420000
 
 RUN: llvm-readobj --headers out.dll | FileCheck -check-prefix=HEADERS %s
 RUN: llvm-readobj --headers out2.dll | FileCheck -check-prefix=HEADERS %s
@@ -97,10 +98,15 @@ RUN: llvm-readobj --coff-load-config out.dll | FileCheck -check-prefix=LOADCONFI
 RUN: llvm-readobj --coff-load-config out2.dll | FileCheck -check-prefix=LOADCONFIG %s
 RUN: llvm-readobj --coff-load-config out3.dll | FileCheck -check-prefix=LOADCONFIG %s
 LOADCONFIG: AuxiliaryIAT: 0x5000
+LOADCONFIG: AuxiliaryIATCopy: 0x4298
 
 RUN: llvm-readobj --hex-dump=.rdata out.dll | FileCheck -check-prefix=RDATA %s
 RUN: llvm-readobj --hex-dump=.rdata out2.dll | FileCheck -check-prefix=RDATA %s
 RUN: llvm-readobj --hex-dump=.rdata out3.dll | FileCheck -check-prefix=RDATA %s
+RDATA:      0x180004290 2e646c6c 00000000 00000000 00000000
+RDATA-NEXT: 0x1800042a0 14100080 01000000 28100080 01000000
+RDATA-NEXT: 0x1800042b0 00000000 00000000 48100080 01000000
+RDATA-NEXT: 0x1800042c0 00000000 00000000 00000000 00000000
 RDATA:      0x180005000 00000000 00000000 14100080 01000000
 RDATA-NEXT: 0x180005010 28100080 01000000 00000000 00000000
 RDATA-NEXT: 0x180005020 48100080 01000000 00000000 00000000
@@ -109,6 +115,18 @@ RUN: llvm-readobj --coff-basereloc out.dll | FileCheck -check-prefix=BASERELOC %
 RUN: llvm-readobj --coff-basereloc out2.dll | FileCheck -check-prefix=BASERELOC %s
 RUN: llvm-readobj --coff-basereloc out3.dll | FileCheck -check-prefix=BASERELOC %s
 BASERELOC:      BaseReloc [
+Aux IAT copy:
+BASERELOC:          Address: 0x42A0
+BASERELOC-NEXT:   }
+BASERELOC-NEXT:   Entry {
+BASERELOC-NEXT:     Type: DIR64
+BASERELOC-NEXT:     Address: 0x42A8
+BASERELOC-NEXT:   }
+BASERELOC-NEXT:   Entry {
+BASERELOC-NEXT:     Type: DIR64
+BASERELOC-NEXT:     Address: 0x42B8
+BASERELOC-NEXT:   }
+Aux IAT:
 BASERELOC-NOT:      Address: 0x5000
 BASERELOC:          Address: 0x5008
 BASERELOC-NEXT:   }
@@ -139,6 +157,7 @@ arm64ec_data_sym:
     .rva func
     .rva "#func"
     .rva "#t2func"
+    .rva __auximpcopy_func
 
 #--- icall.s
     .text