COFF: Fix __ImageBase symbol relocation.

__ImageBase is a special symbol whose value is the image base address.
Previously, we handled __ImageBase symbol as an absolute symbol.

Absolute symbols point to specific locations in memory and the locations
never change even if an image is base-relocated. That means that we
don't have base relocation entries for absolute symbols.

This is not a case for __ImageBase. If an image is base-relocated, its
base address changes, and __ImageBase needs to be shifted as well.
So we have to have base relocations for __ImageBase. That means that
__ImageBase is not really an absolute symbol but a different kind of
symbol.

In this patch, I introduced a new type of symbol -- DefinedRelative.
DefinedRelative is similar to DefinedAbsolute, but it has not a VA but RVA
and is a subject of base relocation. Currently only __ImageBase is of
the new symbol type.

llvm-svn: 243176

lld/COFF/Chunks.cpp

@@ -129,14 +129,12 @@ static bool isAbs(const coff_relocation &Rel) {
 // Collect all locations that contain absolute addresses, which need to be
 // fixed by the loader if load-time relocation is needed.
 // Only called when base relocation is enabled.
-void SectionChunk::getBaserels(std::vector<uint32_t> *Res, Defined *ImageBase) {
+void SectionChunk::getBaserels(std::vector<uint32_t> *Res) {
   for (const coff_relocation &Rel : Relocs) {
-    // Symbol __ImageBase is special -- it's an absolute symbol, but its
-    // address never changes even if image is relocated.
     if (!isAbs(Rel))
       continue;
     SymbolBody *Body = File->getSymbolBody(Rel.SymbolTableIndex)->repl();
-    if (Body == ImageBase)
+    if (isa<DefinedAbsolute>(Body))
       continue;
     Res->push_back(RVA + Rel.VirtualAddress);
   }
@@ -252,8 +250,7 @@ ImportThunkChunk::ImportThunkChunk(Defined *S) : ImpSymbol(S) {
   Align = 16;
 }
 
-void ImportThunkChunk::getBaserels(std::vector<uint32_t> *Res,
-                                   Defined *ImageBase) {
+void ImportThunkChunk::getBaserels(std::vector<uint32_t> *Res) {
   if (!Config->is64())
     Res->push_back(getRVA() + 2);
 }
@@ -267,8 +264,7 @@ void ImportThunkChunk::writeTo(uint8_t *Buf) {
   write32le(Buf + FileOff + 2, Operand);
 }
 
-void LocalImportChunk::getBaserels(std::vector<uint32_t> *Res,
-                                   Defined *ImageBase) {
+void LocalImportChunk::getBaserels(std::vector<uint32_t> *Res) {
   Res->push_back(getRVA());
 }
 

lld/COFF/Chunks.h

@@ -83,7 +83,7 @@ public:
 
   // Windows-specific.
   // Collect all locations that contain absolute addresses for base relocations.
-  virtual void getBaserels(std::vector<uint32_t> *Res, Defined *ImageBase) {}
+  virtual void getBaserels(std::vector<uint32_t> *Res) {}
 
   // Returns a human-readable name of this chunk. Chunks are unnamed chunks of
   // bytes, so this is used only for logging or debugging.
@@ -134,7 +134,7 @@ public:
   bool hasData() const override;
   uint32_t getPermissions() const override;
   StringRef getSectionName() const override { return SectionName; }
-  void getBaserels(std::vector<uint32_t> *Res, Defined *ImageBase) override;
+  void getBaserels(std::vector<uint32_t> *Res) override;
   bool isCOMDAT() const;
   void applyRelX64(uint8_t *Off, uint16_t Type, uint64_t S, uint64_t P);
   void applyRelX86(uint8_t *Off, uint16_t Type, uint64_t S, uint64_t P);
@@ -234,7 +234,7 @@ class ImportThunkChunk : public Chunk {
 public:
   explicit ImportThunkChunk(Defined *ImpSymbol);
   size_t getSize() const override { return sizeof(ImportThunkData); }
-  void getBaserels(std::vector<uint32_t> *Res, Defined *ImageBase) override;
+  void getBaserels(std::vector<uint32_t> *Res) override;
   void writeTo(uint8_t *Buf) override;
 
 private:
@@ -247,7 +247,7 @@ class LocalImportChunk : public Chunk {
 public:
   explicit LocalImportChunk(Defined *S) : Sym(S) {}
   size_t getSize() const override;
-  void getBaserels(std::vector<uint32_t> *Res, Defined *ImageBase) override;
+  void getBaserels(std::vector<uint32_t> *Res) override;
   void writeTo(uint8_t *Buf) override;
 
 private:

lld/COFF/DLL.cpp

@@ -319,7 +319,7 @@ public:
     write32le(Buf + FileOff + 13, Helper->getRVA() - RVA - 17);
   }
 
-  void getBaserels(std::vector<uint32_t> *Res, Defined *ImageBase) override {
+  void getBaserels(std::vector<uint32_t> *Res) override {
     Res->push_back(RVA + 3);
     Res->push_back(RVA + 8);
   }
@@ -367,7 +367,7 @@ public:
     write64le(Buf + FileOff, Thunk->getRVA() + Config->ImageBase);
   }
 
-  void getBaserels(std::vector<uint32_t> *Res, Defined *ImageBase) override {
+  void getBaserels(std::vector<uint32_t> *Res) override {
     Res->push_back(RVA);
   }
 

lld/COFF/Driver.cpp

@@ -602,7 +602,7 @@ bool LinkerDriver::link(llvm::ArrayRef<const char *> ArgsArr) {
     }
   }
 
-  Symtab.addAbsolute(mangle("__ImageBase"), Config->ImageBase);
+  Symtab.addRelative(mangle("__ImageBase"), 0);
 
   // Read as much files as we can from directives sections.
   if (auto EC = Symtab.run()) {

lld/COFF/SymbolTable.cpp

@@ -317,8 +317,8 @@ Undefined *SymbolTable::addUndefined(StringRef Name) {
   return New;
 }
 
-void SymbolTable::addAbsolute(StringRef Name, uint64_t VA) {
-  addSymbol(new (Alloc) DefinedAbsolute(Name, VA));
+void SymbolTable::addRelative(StringRef Name, uint64_t VA) {
+  addSymbol(new (Alloc) DefinedRelative(Name, VA));
 }
 
 void SymbolTable::printMap(llvm::raw_ostream &OS) {

lld/COFF/SymbolTable.h

@@ -81,7 +81,7 @@ public:
 
   // Creates an Undefined symbol for a given name.
   Undefined *addUndefined(StringRef Name);
-  void addAbsolute(StringRef Name, uint64_t VA);
+  void addRelative(StringRef Name, uint64_t VA);
 
   // A list of chunks which to be added to .rdata.
   std::vector<Chunk *> LocalImportChunks;

lld/COFF/Symbols.cpp

@@ -143,6 +143,7 @@ int SymbolBody::compare(SymbolBody *Other) {
   case DefinedImportThunkKind:
   case DefinedImportDataKind:
   case DefinedAbsoluteKind:
+  case DefinedRelativeKind:
     // These all simply tie.
     return 0;
   }
@@ -178,6 +179,8 @@ uint64_t Defined::getFileOff() {
     llvm_unreachable("There is no file offset for a bitcode symbol.");
   case DefinedAbsoluteKind:
     llvm_unreachable("Cannot get a file offset for an absolute symbol.");
+  case DefinedRelativeKind:
+    llvm_unreachable("Cannot get a file offset for a relative symbol.");
   case LazyKind:
   case UndefinedKind:
     llvm_unreachable("Cannot get a file offset for an undefined symbol.");

lld/COFF/Symbols.h

@@ -59,6 +59,7 @@ public:
     DefinedImportThunkKind,
     DefinedImportDataKind,
     DefinedAbsoluteKind,
+    DefinedRelativeKind,
     DefinedBitcodeKind,
 
     UndefinedKind,
@@ -222,6 +223,26 @@ private:
   uint64_t VA;
 };
 
+// This is a kind of absolute symbol but relative to the image base.
+// Unlike absolute symbols, relocations referring this kind of symbols
+// are subject of the base relocation. This type is used rarely --
+// mainly for __ImageBase.
+class DefinedRelative : public Defined {
+public:
+  explicit DefinedRelative(StringRef Name, uint64_t V = 0)
+      : Defined(DefinedRelativeKind, Name), RVA(V) {}
+
+  static bool classof(const SymbolBody *S) {
+    return S->kind() == DefinedRelativeKind;
+  }
+
+  uint64_t getRVA() { return RVA; }
+  void setRVA(uint64_t V) { RVA = V; }
+
+private:
+  uint64_t RVA;
+};
+
 // This class represents a symbol defined in an archive file. It is
 // created from an archive file header, and it knows how to load an
 // object file from an archive to replace itself with a defined
@@ -364,6 +385,8 @@ inline uint64_t Defined::getRVA() {
   switch (kind()) {
   case DefinedAbsoluteKind:
     return cast<DefinedAbsolute>(this)->getRVA();
+  case DefinedRelativeKind:
+    return cast<DefinedRelative>(this)->getRVA();
   case DefinedImportDataKind:
     return cast<DefinedImportData>(this)->getRVA();
   case DefinedImportThunkKind:

lld/COFF/Writer.cpp

@@ -597,14 +597,12 @@ OutputSection *Writer::createSection(StringRef Name) {
 // Dest is .reloc section. Add contents to that section.
 void Writer::addBaserels(OutputSection *Dest) {
   std::vector<uint32_t> V;
-  StringRef Name = Config->is64() ? "__ImageBase" : "___ImageBase";
-  Defined *ImageBase = cast<Defined>(Symtab->find(Name)->Body);
   for (OutputSection *Sec : OutputSections) {
     if (Sec == Dest)
       continue;
     // Collect all locations for base relocations.
     for (Chunk *C : Sec->getChunks())
-      C->getBaserels(&V, ImageBase);
+      C->getBaserels(&V);
     // Add the addresses to .reloc section.
     if (!V.empty())
       addBaserelBlocks(Dest, V);

lld/test/COFF/baserel.test

@@ -7,10 +7,6 @@
 # RUN: llvm-readobj -coff-basereloc %t.exe | FileCheck %s -check-prefix=NOBASEREL
 #
 # BASEREL:      BaseReloc [
-# BASEREL-NEXT:   Entry {
-# BASEREL-NEXT:   Type: DIR64
-# BASEREL-NEXT:   Address: 0x2000
-# BASEREL-NEXT: }
 # BASEREL-NEXT: Entry {
 # BASEREL-NEXT:   Type: DIR64
 # BASEREL-NEXT:   Address: 0x2007
@@ -20,12 +16,12 @@
 # BASEREL-NEXT:   Address: 0x200C
 # BASEREL-NEXT: }
 # BASEREL-NEXT: Entry {
-# BASEREL-NEXT:   Type: ABSOLUTE
-# BASEREL-NEXT:   Address: 0x2000
+# BASEREL-NEXT:   Type: DIR64
+# BASEREL-NEXT:   Address: 0x201E
 # BASEREL-NEXT: }
 # BASEREL-NEXT: Entry {
-# BASEREL-NEXT:   Type: DIR64
-# BASEREL-NEXT:   Address: 0x3000
+# BASEREL-NEXT:   Type: ABSOLUTE
+# BASEREL-NEXT:   Address: 0x2000
 # BASEREL-NEXT: }
 # BASEREL-NEXT: Entry {
 # BASEREL-NEXT:   Type: DIR64
@@ -36,6 +32,10 @@
 # BASEREL-NEXT:   Address: 0x300C
 # BASEREL-NEXT: }
 # BASEREL-NEXT: Entry {
+# BASEREL-NEXT:   Type: DIR64
+# BASEREL-NEXT:   Address: 0x301E
+# BASEREL-NEXT: }
+# BASEREL-NEXT: Entry {
 # BASEREL-NEXT:   Type: ABSOLUTE
 # BASEREL-NEXT:   Address: 0x3000
 # BASEREL-NEXT: }