1 files changed, 68 insertions, 32 deletions

diff --git a/tools/llvm-objcopy/Object.cpp b/tools/llvm-objcopy/Object.cpp
index 28323f47500..a8cd321086b 100644
--- a/tools/llvm-objcopy/Object.cpp
+++ b/tools/llvm-objcopy/Object.cpp
@@ -49,6 +49,7 @@ void Segment::writeSegment(FileOutputBuffer &Out) const {
   std::copy(std::begin(Contents), std::end(Contents), Buf);
 }
 
+void SectionBase::initialize(SectionTableRef SecTable) {}
 void SectionBase::finalize() {}
 
 template <class ELFT>
@@ -149,6 +150,16 @@ void SymbolTableSection::addSymbol(StringRef Name, uint8_t Bind, uint8_t Type,
   Size += this->EntrySize;
 }
 
+void SymbolTableSection::initialize(SectionTableRef SecTable) {
+  Size = 0;
+  setStrTab(SecTable.getSectionOfType<StringTableSection>(
+      Link,
+      "Symbol table has link index of " + Twine(Link) +
+          " which is not a valid index",
+      "Symbol table has link index of " + Twine(Link) +
+          " which is not a string table"));
+}
+
 void SymbolTableSection::finalize() {
   // Make sure SymbolNames is finalized before getting name indexes.
   SymbolNames->finalize();
@@ -195,6 +206,19 @@ void SymbolTableSectionImpl<ELFT>::writeSection(
   }
 }
 
+template <class ELFT>
+void RelocationSection<ELFT>::initialize(SectionTableRef SecTable) {
+  setSymTab(SecTable.getSectionOfType<SymbolTableSection>(
+      Link,
+      "Link field value " + Twine(Link) + " in section " + Name + " is invalid",
+      "Link field value " + Twine(Link) + " in section " + Name +
+          " is not a symbol table"));
+
+  setSection(SecTable.getSection(Info,
+                                 "Info field value " + Twine(Info) +
+                                     " in section " + Name + " is invalid"));
+}
+
 template <class ELFT> void RelocationSection<ELFT>::finalize() {
   this->Link = Symbols->Index;
   this->Info = SecToApplyRel->Index;
@@ -232,6 +256,14 @@ bool SectionWithStrTab::classof(const SectionBase *S) {
   return isa<DynamicSymbolTableSection>(S) || isa<DynamicSection>(S);
 }
 
+void SectionWithStrTab::initialize(SectionTableRef SecTable) {
+  setStrTab(SecTable.getSectionOfType<StringTableSection>(
+      Link,
+      "Link field value " + Twine(Link) + " in section " + Name + " is invalid",
+      "Link field value " + Twine(Link) + " in section " + Name +
+          " is not a string table"));
+}
+
 void SectionWithStrTab::finalize() { this->Link = StrTab->Index; }
 
 // Returns true IFF a section is wholly inside the range of a segment
@@ -311,15 +343,8 @@ void Object<ELFT>::readProgramHeaders(const ELFFile<ELFT> &ElfFile) {
 
 template <class ELFT>
 void Object<ELFT>::initSymbolTable(const llvm::object::ELFFile<ELFT> &ElfFile,
-                                   SymbolTableSection *SymTab) {
-
-  SymTab->Size = 0;
-  SymTab->setStrTab(getSectionOfType<StringTableSection>(
-      SymbolTable->Link,
-      "Symbol table has link index of " + Twine(SymTab->Link) +
-          " which is not a valid index",
-      "Symbol table has link index of " + Twine(SymTab->Link) +
-          " which is not a string table"));
+                                   SymbolTableSection *SymTab,
+                                   SectionTableRef SecTable) {
 
   const Elf_Shdr &Shdr = *unwrapOrError(ElfFile.getSection(SymTab->Index));
   StringRef StrTabData = unwrapOrError(ElfFile.getStringTableForSymtab(Shdr));
@@ -336,7 +361,7 @@ void Object<ELFT>::initSymbolTable(const llvm::object::ELFFile<ELFT> &ElfFile,
             Twine(Sym.st_shndx));
       }
     } else if (Sym.st_shndx != SHN_UNDEF) {
-      DefSection = getSection(
+      DefSection = SecTable.getSection(
           Sym.st_shndx,
           "Symbol '" + Name + "' is defined in invalid section with index " +
               Twine(Sym.st_shndx));
@@ -368,6 +393,20 @@ void initRelocations(RelocationSection<ELFT> *Relocs,
   }
 }
 
+SectionBase *SectionTableRef::getSection(uint16_t Index, Twine ErrMsg) {
+  if (Index == SHN_UNDEF || Index > Sections.size())
+    error(ErrMsg);
+  return Sections[Index - 1].get();
+}
+
+template <class T>
+T *SectionTableRef::getSectionOfType(uint16_t Index, Twine IndexErrMsg,
+                                  Twine TypeErrMsg) {
+  if (T *Sec = llvm::dyn_cast<T>(getSection(Index, IndexErrMsg)))
+    return Sec;
+  error(TypeErrMsg);
+}
+
 template <class ELFT>
 SectionBase *Object<ELFT>::getSection(uint16_t Index, Twine ErrMsg) {
   if (Index == SHN_UNDEF || Index > Sections.size())
@@ -428,7 +467,7 @@ Object<ELFT>::makeSection(const llvm::object::ELFFile<ELFT> &ElfFile,
 }
 
 template <class ELFT>
-void Object<ELFT>::readSectionHeaders(const ELFFile<ELFT> &ElfFile) {
+SectionTableRef Object<ELFT>::readSectionHeaders(const ELFFile<ELFT> &ElfFile) {
   uint32_t Index = 0;
   for (const auto &Shdr : unwrapOrError(ElfFile.sections())) {
     if (Index == 0) {
@@ -451,35 +490,30 @@ void Object<ELFT>::readSectionHeaders(const ELFFile<ELFT> &ElfFile) {
     Sections.push_back(std::move(Sec));
   }
 
+  SectionTableRef SecTable(Sections);
+
   // Now that all of the sections have been added we can fill out some extra
-  // details about symbol tables.
-  if (SymbolTable)
-    initSymbolTable(ElfFile, SymbolTable);
+  // details about symbol tables. We need the symbol table filled out before
+  // any relocations.
+  if (SymbolTable) {
+    SymbolTable->initialize(SecTable);
+    initSymbolTable(ElfFile, SymbolTable, SecTable);
+  }
 
   // Now that all sections and symbols have been added we can add
   // relocations that reference symbols and set the link and info fields for
   // relocation sections.
   for (auto &Section : Sections) {
+    if (Section.get() == SymbolTable)
+      continue;
+    Section->initialize(SecTable);
     if (auto RelSec = dyn_cast<RelocationSection<ELFT>>(Section.get())) {
-
-      auto SymTab = getSectionOfType<SymbolTableSection>(
-          RelSec->Link,
-          "Link field value " + Twine(RelSec->Link) + " in section " +
-              RelSec->Name + " is invalid",
-          "Link field value " + Twine(RelSec->Link) + " in section " +
-              RelSec->Name + " is not a symbol table");
-      RelSec->setSymTab(SymTab);
-
-      RelSec->setSection(getSection(RelSec->Info,
-                                    "Info field value " + Twine(RelSec->Link) +
-                                        " in section " + RelSec->Name +
-                                        " is invalid"));
-
       auto Shdr = unwrapOrError(ElfFile.sections()).begin() + RelSec->Index;
       if (RelSec->Type == SHT_REL)
-        initRelocations(RelSec, SymTab, unwrapOrError(ElfFile.rels(Shdr)));
+        initRelocations(RelSec, SymbolTable, unwrapOrError(ElfFile.rels(Shdr)));
       else
-        initRelocations(RelSec, SymTab, unwrapOrError(ElfFile.relas(Shdr)));
+        initRelocations(RelSec, SymbolTable,
+                        unwrapOrError(ElfFile.relas(Shdr)));
     }
 
     if (auto Sec = dyn_cast<SectionWithStrTab>(Section.get())) {
@@ -491,6 +525,8 @@ void Object<ELFT>::readSectionHeaders(const ELFFile<ELFT> &ElfFile) {
               " is not a string table"));
     }
   }
+
+  return SecTable;
 }
 
 template <class ELFT> Object<ELFT>::Object(const ELFObjectFile<ELFT> &Obj) {
@@ -504,10 +540,10 @@ template <class ELFT> Object<ELFT>::Object(const ELFObjectFile<ELFT> &Obj) {
   Entry = Ehdr.e_entry;
   Flags = Ehdr.e_flags;
 
-  readSectionHeaders(ElfFile);
+  SectionTableRef SecTable = readSectionHeaders(ElfFile);
   readProgramHeaders(ElfFile);
 
-  SectionNames = getSectionOfType<StringTableSection>(
+  SectionNames = SecTable.getSectionOfType<StringTableSection>(
       Ehdr.e_shstrndx,
       "e_shstrndx field value " + Twine(Ehdr.e_shstrndx) + " in elf header " +
           " is invalid",