Reland "[llvm] Add symbol table support to llvm-objcopy"

This change adds support for SHT_SYMTAB sections.

Patch by Jake Ehrlich

Differential Revision: https://reviews.llvm.org/D34167

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@311974 91177308-0d34-0410-b5e6-96231b3b80d8

2 files changed, 147 insertions, 0 deletions

diff --git a/tools/llvm-objcopy/Object.cpp b/tools/llvm-objcopy/Object.cpp
index 55f8855f50f..7738171d4e8 100644
--- a/tools/llvm-objcopy/Object.cpp
+++ b/tools/llvm-objcopy/Object.cpp
@@ -90,6 +90,70 @@ void StringTableSection::writeSection(FileOutputBuffer &Out) const {
   StrTabBuilder.write(Out.getBufferStart() + Offset);
 }
 
+void SymbolTableSection::addSymbol(StringRef Name, uint8_t Bind, uint8_t Type,
+                                   SectionBase *DefinedIn, uint64_t Value,
+                                   uint64_t Sz) {
+  Symbol Sym;
+  Sym.Name = Name;
+  Sym.Binding = Bind;
+  Sym.Type = Type;
+  Sym.DefinedIn = DefinedIn;
+  Sym.Value = Value;
+  Sym.Size = Sz;
+  Sym.Index = Symbols.size();
+  Symbols.emplace_back(llvm::make_unique<Symbol>(Sym));
+  Size += this->EntrySize;
+}
+
+void SymbolTableSection::finalize() {
+  // Make sure SymbolNames is finalized before getting name indexes.
+  SymbolNames->finalize();
+
+  uint32_t MaxLocalIndex = 0;
+  for (auto &Sym : Symbols) {
+    Sym->NameIndex = SymbolNames->findIndex(Sym->Name);
+    if (Sym->Binding == STB_LOCAL)
+      MaxLocalIndex = std::max(MaxLocalIndex, Sym->Index);
+  }
+  // Now we need to set the Link and Info fields.
+  Link = SymbolNames->Index;
+  Info = MaxLocalIndex + 1;
+}
+
+void SymbolTableSection::addSymbolNames() {
+  // Add all of our strings to SymbolNames so that SymbolNames has the right
+  // size before layout is decided.
+  for (auto &Sym : Symbols)
+    SymbolNames->addString(Sym->Name);
+}
+
+const Symbol *SymbolTableSection::getSymbolByIndex(uint32_t Index) const {
+  if (Symbols.size() <= Index)
+    error("Invalid symbol index: " + Twine(Index));
+  return Symbols[Index].get();
+}
+
+template <class ELFT>
+void SymbolTableSectionImpl<ELFT>::writeSection(
+    llvm::FileOutputBuffer &Out) const {
+  uint8_t *Buf = Out.getBufferStart();
+  Buf += Offset;
+  typename ELFT::Sym *Sym = reinterpret_cast<typename ELFT::Sym *>(Buf);
+  // Loop though symbols setting each entry of the symbol table.
+  for (auto &Symbol : Symbols) {
+    Sym->st_name = Symbol->NameIndex;
+    Sym->st_value = Symbol->Value;
+    Sym->st_size = Symbol->Size;
+    Sym->setBinding(Symbol->Binding);
+    Sym->setType(Symbol->Type);
+    if (Symbol->DefinedIn)
+      Sym->st_shndx = Symbol->DefinedIn->Index;
+    else
+      Sym->st_shndx = SHN_UNDEF;
+    ++Sym;
+  }
+}
+
 // Returns true IFF a section is wholly inside the range of a segment
 static bool sectionWithinSegment(const SectionBase &Section,
                                  const Segment &Segment) {
@@ -133,6 +197,40 @@ 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;
+  if (SymbolTable->Link - 1 >= Sections.size())
+    error("Symbol table has link index of " + Twine(SymbolTable->Link) +
+          " which is not a valid index");
+
+  if (auto StrTab =
+          dyn_cast<StringTableSection>(Sections[SymbolTable->Link - 1].get()))
+    SymTab->setStrTab(StrTab);
+  else
+    error("Symbol table has link index of " + Twine(SymbolTable->Link) +
+          "which is not a string table");
+
+  const Elf_Shdr &Shdr = *unwrapOrError(ElfFile.getSection(SymTab->Index));
+  StringRef StrTabData = unwrapOrError(ElfFile.getStringTableForSymtab(Shdr));
+
+  for (const auto &Sym : unwrapOrError(ElfFile.symbols(&Shdr))) {
+    SectionBase *DefSection = nullptr;
+    StringRef Name = unwrapOrError(Sym.getName(StrTabData));
+    if (Sym.st_shndx != SHN_UNDEF) {
+      if (Sym.st_shndx >= Sections.size())
+        error("Symbol '" + Name +
+              "' is defined in invalid section with index " +
+              Twine(Sym.st_shndx));
+      DefSection = Sections[Sym.st_shndx - 1].get();
+    }
+    SymTab->addSymbol(Name, Sym.getBinding(), Sym.getType(), DefSection,
+                      Sym.getValue(), Sym.st_size);
+  }
+}
+
+template <class ELFT>
 std::unique_ptr<SectionBase>
 Object<ELFT>::makeSection(const llvm::object::ELFFile<ELFT> &ElfFile,
                           const Elf_Shdr &Shdr) {
@@ -140,6 +238,11 @@ Object<ELFT>::makeSection(const llvm::object::ELFFile<ELFT> &ElfFile,
   switch (Shdr.sh_type) {
   case SHT_STRTAB:
     return llvm::make_unique<StringTableSection>();
+  case SHT_SYMTAB: {
+    auto SymTab = llvm::make_unique<SymbolTableSectionImpl<ELFT>>();
+    SymbolTable = SymTab.get();
+    return std::move(SymTab);
+  }
   case SHT_NOBITS:
     return llvm::make_unique<Section>(Data);
   default:
@@ -171,6 +274,11 @@ void Object<ELFT>::readSectionHeaders(const ELFFile<ELFT> &ElfFile) {
     Sec->Index = Index++;
     Sections.push_back(std::move(Sec));
   }
+
+  // Now that all of the sections have been added we can fill out some extra
+  // details about symbol tables.
+  if (SymbolTable)
+    initSymbolTable(ElfFile, SymbolTable);
 }
 
 template <class ELFT> Object<ELFT>::Object(const ELFObjectFile<ELFT> &Obj) {
@@ -315,9 +423,12 @@ template <class ELFT> void ELFObject<ELFT>::write(FileOutputBuffer &Out) const {
 }
 
 template <class ELFT> void ELFObject<ELFT>::finalize() {
+  // Make sure we add the names of all the sections.
   for (const auto &Section : this->Sections) {
     this->SectionNames->addString(Section->Name);
   }
+  // Make sure we add the names of all the symbols.
+  this->SymbolTable->addSymbolNames();
 
   sortSections();
   assignOffsets();
diff --git a/tools/llvm-objcopy/Object.h b/tools/llvm-objcopy/Object.h
index 1a3777bbd72..c41bd36490c 100644
--- a/tools/llvm-objcopy/Object.h
+++ b/tools/llvm-objcopy/Object.h
@@ -113,6 +113,39 @@ public:
   }
 };
 
+struct Symbol {
+  uint8_t Binding;
+  SectionBase *DefinedIn;
+  uint32_t Index;
+  llvm::StringRef Name;
+  uint32_t NameIndex;
+  uint64_t Size;
+  uint8_t Type;
+  uint64_t Value;
+};
+
+class SymbolTableSection : public SectionBase {
+protected:
+  std::vector<std::unique_ptr<Symbol>> Symbols;
+  StringTableSection *SymbolNames;
+
+public:
+  void setStrTab(StringTableSection *StrTab) { SymbolNames = StrTab; }
+  void addSymbol(llvm::StringRef Name, uint8_t Bind, uint8_t Type,
+                 SectionBase *DefinedIn, uint64_t Value, uint64_t Sz);
+  void addSymbolNames();
+  const Symbol *getSymbolByIndex(uint32_t Index) const;
+  void finalize() override;
+  static bool classof(const SectionBase *S) {
+    return S->Type == llvm::ELF::SHT_SYMTAB;
+  }
+};
+
+// Only writeSection depends on the ELF type so we implement it in a subclass.
+template <class ELFT> class SymbolTableSectionImpl : public SymbolTableSection {
+  void writeSection(llvm::FileOutputBuffer &Out) const override;
+};
+
 template <class ELFT> class Object {
 private:
   typedef std::unique_ptr<SectionBase> SecPtr;
@@ -122,6 +155,8 @@ private:
   typedef typename ELFT::Ehdr Elf_Ehdr;
   typedef typename ELFT::Phdr Elf_Phdr;
 
+  void initSymbolTable(const llvm::object::ELFFile<ELFT> &ElfFile,
+                       SymbolTableSection *SymTab);
   SecPtr makeSection(const llvm::object::ELFFile<ELFT> &ElfFile,
                      const Elf_Shdr &Shdr);
   void readProgramHeaders(const llvm::object::ELFFile<ELFT> &ElfFile);
@@ -129,6 +164,7 @@ private:
 
 protected:
   StringTableSection *SectionNames;
+  SymbolTableSection *SymbolTable;
   std::vector<SecPtr> Sections;
   std::vector<SegPtr> Segments;