[llvm-objcopy] Fix crash when removing symbol table at same time as adding a symbol

This patch resolves crash that occurs when user wanted to remove all
symbols and add a brand new one using:

```
llvm-objcopy -R .symtab --add-symbol foo=1234 in.o out.o
```

Before these changes the symbol table internally being null when adding
new symbols. For now we will regenerate symtab in this case.

This fixes: https://bugs.llvm.org/show_bug.cgi?id=43930

Reviewed By: jhenderson

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

1 files changed, 28 insertions, 21 deletions

diff --git a/llvm/tools/llvm-objcopy/ELF/Object.cpp b/llvm/tools/llvm-objcopy/ELF/Object.cpp
index 944bfdf..e15fb24 100644
--- a/llvm/tools/llvm-objcopy/ELF/Object.cpp
+++ b/llvm/tools/llvm-objcopy/ELF/Object.cpp
@@ -1588,27 +1588,7 @@ template <class ELFT> void ELFBuilder<ELFT>::readSections(bool EnsureSymtab) {
     Obj.SymbolTable->initialize(Obj.sections());
     initSymbolTable(Obj.SymbolTable);
   } else if (EnsureSymtab) {
-    // Reuse an existing SHT_STRTAB section if it 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 a string table that is not the section header string table, if
-        // such a table exists.
-        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;
+    Obj.addNewSymbolTable();
   }
 
   // Now that all sections and symbols have been added we can add
@@ -1900,6 +1880,33 @@ Error Object::removeSymbols(function_ref<bool(const Symbol &)> ToRemove) {
   return Error::success();
 }
 
+void Object::addNewSymbolTable() {
+  assert(!SymbolTable && "Object must not has a SymbolTable.");
+
+  // Reuse an existing SHT_STRTAB section if it exists.
+  StringTableSection *StrTab = nullptr;
+  for (SectionBase &Sec : sections()) {
+    if (Sec.Type == ELF::SHT_STRTAB && !(Sec.Flags & SHF_ALLOC)) {
+      StrTab = static_cast<StringTableSection *>(&Sec);
+
+      // Prefer a string table that is not the section header string table, if
+      // such a table exists.
+      if (SectionNames != &Sec)
+        break;
+    }
+  }
+  if (!StrTab)
+    StrTab = &addSection<StringTableSection>();
+
+  SymbolTableSection &SymTab = addSection<SymbolTableSection>();
+  SymTab.Name = ".symtab";
+  SymTab.Link = StrTab->Index;
+  SymTab.initialize(sections());
+  SymTab.addSymbol("", 0, 0, nullptr, 0, 0, 0, 0);
+
+  SymbolTable = &SymTab;
+}
+
 void Object::sortSections() {
   // Use stable_sort to maintain the original ordering as closely as possible.
   llvm::stable_sort(Sections, [](const SecPtr &A, const SecPtr &B) {