[llvm-objcopy] --add-symbol: fix crash if SHT_SYMTAB does not exist

Exposed by D69041. If SHT_SYMTAB does not exist, ELFObjcopy.cpp:handleArgs will crash due
to a null pointer dereference.

  for (const NewSymbolInfo &SI : Config.ELF->SymbolsToAdd) {
    ...
    Obj.SymbolTable->addSymbol(

Fix this by creating .symtab and .strtab on demand in ELFBuilder<ELFT>::readSections,
if --add-symbol is specified.

Reviewed By: grimar

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

llvm-svn: 375105

3 files changed, 41 insertions, 20 deletions

diff --git a/llvm/tools/llvm-objcopy/ELF/ELFObjcopy.cpp b/llvm/tools/llvm-objcopy/ELF/ELFObjcopy.cpp
index dd6a7d7e14b..8bf7e0f8801 100644
--- a/llvm/tools/llvm-objcopy/ELF/ELFObjcopy.cpp
+++ b/llvm/tools/llvm-objcopy/ELF/ELFObjcopy.cpp
@@ -263,7 +263,7 @@ static Error linkToBuildIdDir(const CopyConfig &Config, StringRef ToLink,
 
 static Error splitDWOToFile(const CopyConfig &Config, const Reader &Reader,
                             StringRef File, ElfType OutputElfType) {
-  auto DWOFile = Reader.create();
+  auto DWOFile = Reader.create(false);
   auto OnlyKeepDWOPred = [&DWOFile](const SectionBase &Sec) {
     return onlyKeepDWOPred(*DWOFile, Sec);
   };
@@ -743,7 +743,7 @@ static Error writeOutput(const CopyConfig &Config, Object &Obj, Buffer &Out,
 Error executeObjcopyOnIHex(const CopyConfig &Config, MemoryBuffer &In,
                            Buffer &Out) {
   IHexReader Reader(&In);
-  std::unique_ptr<Object> Obj = Reader.create();
+  std::unique_ptr<Object> Obj = Reader.create(true);
   const ElfType OutputElfType =
     getOutputElfType(Config.OutputArch.getValueOr(MachineInfo()));
   if (Error E = handleArgs(Config, *Obj, Reader, OutputElfType))
@@ -756,7 +756,7 @@ Error executeObjcopyOnRawBinary(const CopyConfig &Config, MemoryBuffer &In,
   uint8_t NewSymbolVisibility =
       Config.ELF->NewSymbolVisibility.getValueOr((uint8_t)ELF::STV_DEFAULT);
   BinaryReader Reader(&In, NewSymbolVisibility);
-  std::unique_ptr<Object> Obj = Reader.create();
+  std::unique_ptr<Object> Obj = Reader.create(true);
 
   // Prefer OutputArch (-O<format>) if set, otherwise fallback to BinaryArch
   // (-B<arch>).
@@ -770,7 +770,7 @@ Error executeObjcopyOnRawBinary(const CopyConfig &Config, MemoryBuffer &In,
 Error executeObjcopyOnBinary(const CopyConfig &Config,
                              object::ELFObjectFileBase &In, Buffer &Out) {
   ELFReader Reader(&In, Config.ExtractPartition);
-  std::unique_ptr<Object> Obj = Reader.create();
+  std::unique_ptr<Object> Obj = Reader.create(!Config.SymbolsToAdd.empty());
   // Prefer OutputArch (-O<format>) if set, otherwise infer it from the input.
   const ElfType OutputElfType =
       Config.OutputArch ? getOutputElfType(Config.OutputArch.getValue())
diff --git a/llvm/tools/llvm-objcopy/ELF/Object.cpp b/llvm/tools/llvm-objcopy/ELF/Object.cpp
index 0220b2b5a92..74145dad6e6 100644
--- a/llvm/tools/llvm-objcopy/ELF/Object.cpp
+++ b/llvm/tools/llvm-objcopy/ELF/Object.cpp
@@ -1527,7 +1527,7 @@ template <class ELFT> void ELFBuilder<ELFT>::readSectionHeaders() {
   }
 }
 
-template <class ELFT> void ELFBuilder<ELFT>::readSections() {
+template <class ELFT> void ELFBuilder<ELFT>::readSections(bool EnsureSymtab) {
   // If a section index table exists we'll need to initialize it before we
   // initialize the symbol table because the symbol table might need to
   // reference it.
@@ -1540,6 +1540,27 @@ template <class ELFT> void ELFBuilder<ELFT>::readSections() {
   if (Obj.SymbolTable) {
     Obj.SymbolTable->initialize(Obj.sections());
     initSymbolTable(Obj.SymbolTable);
+  } else if (EnsureSymtab) {
+    // Reuse the existing SHT_STRTAB section if exists.
+    StringTableSection *StrTab = nullptr;
+    for (auto &Sec : Obj.sections()) {
+      if (Sec.Type == ELF::SHT_STRTAB && !(Sec.Flags & SHF_ALLOC)) {
+        StrTab = static_cast<StringTableSection *>(&Sec);
+
+        // Prefer .strtab to .shstrtab.
+        if (Obj.SectionNames != &Sec)
+          break;
+      }
+    }
+    if (!StrTab)
+      StrTab = &Obj.addSection<StringTableSection>();
+
+    SymbolTableSection &SymTab = Obj.addSection<SymbolTableSection>();
+    SymTab.Name = ".symtab";
+    SymTab.Link = StrTab->Index;
+    SymTab.initialize(Obj.sections());
+    SymTab.addSymbol("", 0, 0, nullptr, 0, 0, 0, 0);
+    Obj.SymbolTable = &SymTab;
   }
 
   // Now that all sections and symbols have been added we can add
@@ -1578,7 +1599,7 @@ template <class ELFT> void ELFBuilder<ELFT>::readSections() {
                 " is not a string table");
 }
 
-template <class ELFT> void ELFBuilder<ELFT>::build() {
+template <class ELFT> void ELFBuilder<ELFT>::build(bool EnsureSymtab) {
   readSectionHeaders();
   findEhdrOffset();
 
@@ -1597,7 +1618,7 @@ template <class ELFT> void ELFBuilder<ELFT>::build() {
   Obj.Entry = Ehdr.e_entry;
   Obj.Flags = Ehdr.e_flags;
 
-  readSections();
+  readSections(EnsureSymtab);
   readProgramHeaders(HeadersFile);
 }
 
@@ -1605,7 +1626,7 @@ Writer::~Writer() {}
 
 Reader::~Reader() {}
 
-std::unique_ptr<Object> BinaryReader::create() const {
+std::unique_ptr<Object> BinaryReader::create(bool /*EnsureSymtab*/) const {
   return BinaryELFBuilder(MemBuf, NewSymbolVisibility).build();
 }
 
@@ -1635,28 +1656,28 @@ Expected<std::vector<IHexRecord>> IHexReader::parse() const {
   return std::move(Records);
 }
 
-std::unique_ptr<Object> IHexReader::create() const {
+std::unique_ptr<Object> IHexReader::create(bool /*EnsureSymtab*/) const {
   std::vector<IHexRecord> Records = unwrapOrError(parse());
   return IHexELFBuilder(Records).build();
 }
 
-std::unique_ptr<Object> ELFReader::create() const {
+std::unique_ptr<Object> ELFReader::create(bool EnsureSymtab) const {
   auto Obj = std::make_unique<Object>();
   if (auto *O = dyn_cast<ELFObjectFile<ELF32LE>>(Bin)) {
     ELFBuilder<ELF32LE> Builder(*O, *Obj, ExtractPartition);
-    Builder.build();
+    Builder.build(EnsureSymtab);
     return Obj;
   } else if (auto *O = dyn_cast<ELFObjectFile<ELF64LE>>(Bin)) {
     ELFBuilder<ELF64LE> Builder(*O, *Obj, ExtractPartition);
-    Builder.build();
+    Builder.build(EnsureSymtab);
     return Obj;
   } else if (auto *O = dyn_cast<ELFObjectFile<ELF32BE>>(Bin)) {
     ELFBuilder<ELF32BE> Builder(*O, *Obj, ExtractPartition);
-    Builder.build();
+    Builder.build(EnsureSymtab);
     return Obj;
   } else if (auto *O = dyn_cast<ELFObjectFile<ELF64BE>>(Bin)) {
     ELFBuilder<ELF64BE> Builder(*O, *Obj, ExtractPartition);
-    Builder.build();
+    Builder.build(EnsureSymtab);
     return Obj;
   }
   error("invalid file type");
diff --git a/llvm/tools/llvm-objcopy/ELF/Object.h b/llvm/tools/llvm-objcopy/ELF/Object.h
index d74b5f410af..eeacb014e4d 100644
--- a/llvm/tools/llvm-objcopy/ELF/Object.h
+++ b/llvm/tools/llvm-objcopy/ELF/Object.h
@@ -863,7 +863,7 @@ public:
 class Reader {
 public:
   virtual ~Reader();
-  virtual std::unique_ptr<Object> create() const = 0;
+  virtual std::unique_ptr<Object> create(bool EnsureSymtab) const = 0;
 };
 
 using object::Binary;
@@ -926,7 +926,7 @@ private:
   void initGroupSection(GroupSection *GroupSec);
   void initSymbolTable(SymbolTableSection *SymTab);
   void readSectionHeaders();
-  void readSections();
+  void readSections(bool EnsureSymtab);
   void findEhdrOffset();
   SectionBase &makeSection(const Elf_Shdr &Shdr);
 
@@ -936,7 +936,7 @@ public:
       : ElfFile(*ElfObj.getELFFile()), Obj(Obj),
         ExtractPartition(ExtractPartition) {}
 
-  void build();
+  void build(bool EnsureSymtab);
 };
 
 class BinaryReader : public Reader {
@@ -946,7 +946,7 @@ class BinaryReader : public Reader {
 public:
   BinaryReader(MemoryBuffer *MB, const uint8_t NewSymbolVisibility)
       : MemBuf(MB), NewSymbolVisibility(NewSymbolVisibility) {}
-  std::unique_ptr<Object> create() const override;
+  std::unique_ptr<Object> create(bool EnsureSymtab) const override;
 };
 
 class IHexReader : public Reader {
@@ -968,7 +968,7 @@ class IHexReader : public Reader {
 public:
   IHexReader(MemoryBuffer *MB) : MemBuf(MB) {}
 
-  std::unique_ptr<Object> create() const override;
+  std::unique_ptr<Object> create(bool EnsureSymtab) const override;
 };
 
 class ELFReader : public Reader {
@@ -976,7 +976,7 @@ class ELFReader : public Reader {
   Optional<StringRef> ExtractPartition;
 
 public:
-  std::unique_ptr<Object> create() const override;
+  std::unique_ptr<Object> create(bool EnsureSymtab) const override;
   explicit ELFReader(Binary *B, Optional<StringRef> ExtractPartition)
       : Bin(B), ExtractPartition(ExtractPartition) {}
 };