[llvm] a6f3fed - [objcopy] Refactor CommonConfig to add posibility to specify added/updated sections as MemoryBuffer.

Alexey Lapshin via llvm-commits llvm-commits at lists.llvm.org
Tue Mar 1 03:50:41 PST 2022


Author: Alexey Lapshin
Date: 2022-03-01T14:49:41+03:00
New Revision: a6f3fedc3f12a444924b4d30a4dcad1b6eaad965

URL: https://github.com/llvm/llvm-project/commit/a6f3fedc3f12a444924b4d30a4dcad1b6eaad965
DIFF: https://github.com/llvm/llvm-project/commit/a6f3fedc3f12a444924b4d30a4dcad1b6eaad965.diff

LOG: [objcopy] Refactor CommonConfig to add posibility to specify added/updated sections as MemoryBuffer.

Current objcopy implementation has a possibility to add or update sections.
The incoming section is specified as a pair: section name and name of the file
containing section data. The interface does not allow to specify incoming
section as a memory buffer. This patch adds possibility to specify incoming
section as a memory buffer.

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

Added: 
    

Modified: 
    llvm/include/llvm/ObjCopy/CommonConfig.h
    llvm/lib/ObjCopy/COFF/COFFObjcopy.cpp
    llvm/lib/ObjCopy/ELF/ELFObjcopy.cpp
    llvm/lib/ObjCopy/MachO/MachOObjcopy.cpp
    llvm/lib/ObjCopy/wasm/WasmObjcopy.cpp
    llvm/test/tools/llvm-objcopy/COFF/add-section.test
    llvm/test/tools/llvm-objcopy/ELF/add-section.test
    llvm/test/tools/llvm-objcopy/MachO/add-section-error.test
    llvm/test/tools/llvm-objcopy/MachO/update-section.test
    llvm/tools/llvm-objcopy/ObjcopyOptions.cpp
    llvm/unittests/ObjCopy/ObjCopyTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ObjCopy/CommonConfig.h b/llvm/include/llvm/ObjCopy/CommonConfig.h
index ecb169a4e8ec2..24503caed3427 100644
--- a/llvm/include/llvm/ObjCopy/CommonConfig.h
+++ b/llvm/include/llvm/ObjCopy/CommonConfig.h
@@ -18,6 +18,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Object/ELFTypes.h"
 #include "llvm/Support/GlobPattern.h"
+#include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Regex.h"
 // Necessary for llvm::DebugCompressionType::None
 #include "llvm/Target/TargetOptions.h"
@@ -186,6 +187,16 @@ struct NewSymbolInfo {
   std::vector<StringRef> BeforeSyms;
 };
 
+// Specify section name and section body for newly added or updated section.
+struct NewSectionInfo {
+  NewSectionInfo() = default;
+  NewSectionInfo(StringRef Name, std::unique_ptr<MemoryBuffer> &&Buffer)
+      : SectionName(Name), SectionData(std::move(Buffer)) {}
+
+  StringRef SectionName;
+  std::shared_ptr<MemoryBuffer> SectionData;
+};
+
 // Configuration for copying/stripping a single file.
 struct CommonConfig {
   // Main input/output options
@@ -208,9 +219,9 @@ struct CommonConfig {
   DiscardType DiscardMode = DiscardType::None;
 
   // Repeated options
-  std::vector<StringRef> AddSection;
+  std::vector<NewSectionInfo> AddSection;
   std::vector<StringRef> DumpSection;
-  std::vector<StringRef> UpdateSection;
+  std::vector<NewSectionInfo> UpdateSection;
 
   // Section matchers
   NameMatcher KeepSection;

diff  --git a/llvm/lib/ObjCopy/COFF/COFFObjcopy.cpp b/llvm/lib/ObjCopy/COFF/COFFObjcopy.cpp
index 16b57703391bd..cda93ce0fb3c2 100644
--- a/llvm/lib/ObjCopy/COFF/COFFObjcopy.cpp
+++ b/llvm/lib/ObjCopy/COFF/COFFObjcopy.cpp
@@ -230,56 +230,41 @@ static Error handleArgs(const CommonConfig &Config,
             It->second.NewFlags, Sec.Header.Characteristics);
     }
 
-  for (const auto &Flag : Config.AddSection) {
-    StringRef SecName, FileName;
-    std::tie(SecName, FileName) = Flag.split("=");
-
-    auto BufOrErr = MemoryBuffer::getFile(FileName);
-    if (!BufOrErr)
-      return createFileError(FileName, errorCodeToError(BufOrErr.getError()));
-    auto Buf = std::move(*BufOrErr);
-
+  for (const NewSectionInfo &NewSection : Config.AddSection) {
     uint32_t Characteristics;
-    const auto It = Config.SetSectionFlags.find(SecName);
+    const auto It = Config.SetSectionFlags.find(NewSection.SectionName);
     if (It != Config.SetSectionFlags.end())
       Characteristics = flagsToCharacteristics(It->second.NewFlags, 0);
     else
       Characteristics = IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_ALIGN_1BYTES;
 
-    addSection(
-        Obj, SecName,
-        makeArrayRef(reinterpret_cast<const uint8_t *>(Buf->getBufferStart()),
-                     Buf->getBufferSize()),
-        Characteristics);
+    addSection(Obj, NewSection.SectionName,
+               makeArrayRef(reinterpret_cast<const uint8_t *>(
+                                NewSection.SectionData->getBufferStart()),
+                            NewSection.SectionData->getBufferSize()),
+               Characteristics);
   }
 
-  for (StringRef Flag : Config.UpdateSection) {
-    StringRef SecName, FileName;
-    std::tie(SecName, FileName) = Flag.split('=');
-
-    auto BufOrErr = MemoryBuffer::getFile(FileName);
-    if (!BufOrErr)
-      return createFileError(FileName, errorCodeToError(BufOrErr.getError()));
-    auto Buf = std::move(*BufOrErr);
-
-    auto It = llvm::find_if(Obj.getMutableSections(), [SecName](auto &Sec) {
-      return Sec.Name == SecName;
+  for (const NewSectionInfo &NewSection : Config.UpdateSection) {
+    auto It = llvm::find_if(Obj.getMutableSections(), [&](auto &Sec) {
+      return Sec.Name == NewSection.SectionName;
     });
     if (It == Obj.getMutableSections().end())
       return createStringError(errc::invalid_argument,
                                "could not find section with name '%s'",
-                               SecName.str().c_str());
+                               NewSection.SectionName.str().c_str());
     size_t ContentSize = It->getContents().size();
     if (!ContentSize)
       return createStringError(
           errc::invalid_argument,
           "section '%s' cannot be updated because it does not have contents",
-          SecName.str().c_str());
-    if (ContentSize < Buf->getBufferSize())
+          NewSection.SectionName.str().c_str());
+    if (ContentSize < NewSection.SectionData->getBufferSize())
       return createStringError(
           errc::invalid_argument,
           "new section cannot be larger than previous section");
-    It->setOwnedContents({Buf->getBufferStart(), Buf->getBufferEnd()});
+    It->setOwnedContents({NewSection.SectionData->getBufferStart(),
+                          NewSection.SectionData->getBufferEnd()});
   }
 
   if (!Config.AddGnuDebugLink.empty())

diff  --git a/llvm/lib/ObjCopy/ELF/ELFObjcopy.cpp b/llvm/lib/ObjCopy/ELF/ELFObjcopy.cpp
index 9059e4a17afd4..756f6ccbd288e 100644
--- a/llvm/lib/ObjCopy/ELF/ELFObjcopy.cpp
+++ b/llvm/lib/ObjCopy/ELF/ELFObjcopy.cpp
@@ -586,19 +586,12 @@ static void addSymbol(Object &Obj, const NewSymbolInfo &SymInfo,
 }
 
 static Error
-handleUserSection(StringRef Flag,
+handleUserSection(const NewSectionInfo &NewSection,
                   function_ref<Error(StringRef, ArrayRef<uint8_t>)> F) {
-  std::pair<StringRef, StringRef> SecPair = Flag.split("=");
-  StringRef SecName = SecPair.first;
-  StringRef File = SecPair.second;
-  ErrorOr<std::unique_ptr<MemoryBuffer>> BufOrErr = MemoryBuffer::getFile(File);
-  if (!BufOrErr)
-    return createFileError(File, errorCodeToError(BufOrErr.getError()));
-  std::unique_ptr<MemoryBuffer> Buf = std::move(*BufOrErr);
-  ArrayRef<uint8_t> Data(
-      reinterpret_cast<const uint8_t *>(Buf->getBufferStart()),
-      Buf->getBufferSize());
-  return F(SecName, Data);
+  ArrayRef<uint8_t> Data(reinterpret_cast<const uint8_t *>(
+                             NewSection.SectionData->getBufferStart()),
+                         NewSection.SectionData->getBufferSize());
+  return F(NewSection.SectionName, Data);
 }
 
 // This function handles the high level operations of GNU objcopy including
@@ -717,7 +710,7 @@ static Error handleArgs(const CommonConfig &Config, const ELFConfig &ELFConfig,
       if (Sec.Flags & SHF_ALLOC && Sec.Type != SHT_NOTE)
         Sec.Type = SHT_NOBITS;
 
-  for (const auto &Flag : Config.AddSection) {
+  for (const NewSectionInfo &AddedSection : Config.AddSection) {
     auto AddSection = [&](StringRef Name, ArrayRef<uint8_t> Data) {
       OwnedDataSection &NewSection =
           Obj.addSection<OwnedDataSection>(Name, Data);
@@ -725,15 +718,15 @@ static Error handleArgs(const CommonConfig &Config, const ELFConfig &ELFConfig,
         NewSection.Type = SHT_NOTE;
       return Error::success();
     };
-    if (Error E = handleUserSection(Flag, AddSection))
+    if (Error E = handleUserSection(AddedSection, AddSection))
       return E;
   }
 
-  for (StringRef Flag : Config.UpdateSection) {
+  for (const NewSectionInfo &NewSection : Config.UpdateSection) {
     auto UpdateSection = [&](StringRef Name, ArrayRef<uint8_t> Data) {
       return Obj.updateSection(Name, Data);
     };
-    if (Error E = handleUserSection(Flag, UpdateSection))
+    if (Error E = handleUserSection(NewSection, UpdateSection))
       return E;
   }
 

diff  --git a/llvm/lib/ObjCopy/MachO/MachOObjcopy.cpp b/llvm/lib/ObjCopy/MachO/MachOObjcopy.cpp
index b8e21222aa3e1..8eebdd6f8c659 100644
--- a/llvm/lib/ObjCopy/MachO/MachOObjcopy.cpp
+++ b/llvm/lib/ObjCopy/MachO/MachOObjcopy.cpp
@@ -99,7 +99,7 @@ static void updateAndRemoveSymbols(const CommonConfig &Config,
       Sym.Name = std::string(I->getValue());
   }
 
-  auto RemovePred = [Config, MachOConfig,
+  auto RemovePred = [&Config, &MachOConfig,
                      &Obj](const std::unique_ptr<SymbolEntry> &N) {
     if (N->Referenced)
       return false;
@@ -283,17 +283,12 @@ static Error dumpSectionToFile(StringRef SecName, StringRef Filename,
                            SecName.str().c_str());
 }
 
-static Error addSection(StringRef SecName, StringRef Filename, Object &Obj) {
-  ErrorOr<std::unique_ptr<MemoryBuffer>> BufOrErr =
-      MemoryBuffer::getFile(Filename);
-  if (!BufOrErr)
-    return createFileError(Filename, errorCodeToError(BufOrErr.getError()));
-  std::unique_ptr<MemoryBuffer> Buf = std::move(*BufOrErr);
-
-  std::pair<StringRef, StringRef> Pair = SecName.split(',');
+static Error addSection(const NewSectionInfo &NewSection, Object &Obj) {
+  std::pair<StringRef, StringRef> Pair = NewSection.SectionName.split(',');
   StringRef TargetSegName = Pair.first;
   Section Sec(TargetSegName, Pair.second);
-  Sec.Content = Obj.NewSectionsContents.save(Buf->getBuffer());
+  Sec.Content =
+      Obj.NewSectionsContents.save(NewSection.SectionData->getBuffer());
   Sec.Size = Sec.Content.size();
 
   // Add the a section into an existing segment.
@@ -342,24 +337,18 @@ static Expected<Section &> findSection(StringRef SecName, Object &O) {
   return *FoundSec->get();
 }
 
-static Error updateSection(StringRef SecName, StringRef Filename, Object &O) {
-  Expected<Section &> SecToUpdateOrErr = findSection(SecName, O);
+static Error updateSection(const NewSectionInfo &NewSection, Object &O) {
+  Expected<Section &> SecToUpdateOrErr = findSection(NewSection.SectionName, O);
 
   if (!SecToUpdateOrErr)
     return SecToUpdateOrErr.takeError();
   Section &Sec = *SecToUpdateOrErr;
 
-  ErrorOr<std::unique_ptr<MemoryBuffer>> BufOrErr =
-      MemoryBuffer::getFile(Filename);
-  if (!BufOrErr)
-    return createFileError(Filename, errorCodeToError(BufOrErr.getError()));
-  std::unique_ptr<MemoryBuffer> Buf = std::move(*BufOrErr);
-
-  if (Buf->getBufferSize() > Sec.Size)
+  if (NewSection.SectionData->getBufferSize() > Sec.Size)
     return createStringError(
         errc::invalid_argument,
         "new section cannot be larger than previous section");
-  Sec.Content = O.NewSectionsContents.save(Buf->getBuffer());
+  Sec.Content = O.NewSectionsContents.save(NewSection.SectionData->getBuffer());
   Sec.Size = Sec.Content.size();
   return Error::success();
 }
@@ -411,23 +400,17 @@ static Error handleArgs(const CommonConfig &Config,
       for (std::unique_ptr<Section> &Sec : LC.Sections)
         Sec->Relocations.clear();
 
-  for (const auto &Flag : Config.AddSection) {
-    std::pair<StringRef, StringRef> SecPair = Flag.split("=");
-    StringRef SecName = SecPair.first;
-    StringRef File = SecPair.second;
-    if (Error E = isValidMachOCannonicalName(SecName))
+  for (const NewSectionInfo &NewSection : Config.AddSection) {
+    if (Error E = isValidMachOCannonicalName(NewSection.SectionName))
       return E;
-    if (Error E = addSection(SecName, File, Obj))
+    if (Error E = addSection(NewSection, Obj))
       return E;
   }
 
-  for (const auto &Flag : Config.UpdateSection) {
-    StringRef SectionName;
-    StringRef FileName;
-    std::tie(SectionName, FileName) = Flag.split('=');
-    if (Error E = isValidMachOCannonicalName(SectionName))
+  for (const NewSectionInfo &NewSection : Config.UpdateSection) {
+    if (Error E = isValidMachOCannonicalName(NewSection.SectionName))
       return E;
-    if (Error E = updateSection(SectionName, FileName, Obj))
+    if (Error E = updateSection(NewSection, Obj))
       return E;
   }
 

diff  --git a/llvm/lib/ObjCopy/wasm/WasmObjcopy.cpp b/llvm/lib/ObjCopy/wasm/WasmObjcopy.cpp
index ff56df82eaf74..6877cd68bee47 100644
--- a/llvm/lib/ObjCopy/wasm/WasmObjcopy.cpp
+++ b/llvm/lib/ObjCopy/wasm/WasmObjcopy.cpp
@@ -121,21 +121,19 @@ static Error handleArgs(const CommonConfig &Config, Object &Obj) {
 
   removeSections(Config, Obj);
 
-  for (StringRef Flag : Config.AddSection) {
-    StringRef SecName, FileName;
-    std::tie(SecName, FileName) = Flag.split("=");
-    ErrorOr<std::unique_ptr<MemoryBuffer>> BufOrErr =
-        MemoryBuffer::getFile(FileName);
-    if (!BufOrErr)
-      return createFileError(FileName, errorCodeToError(BufOrErr.getError()));
+  for (const NewSectionInfo &NewSection : Config.AddSection) {
     Section Sec;
     Sec.SectionType = llvm::wasm::WASM_SEC_CUSTOM;
-    Sec.Name = SecName;
-    std::unique_ptr<MemoryBuffer> Buf = std::move(*BufOrErr);
+    Sec.Name = NewSection.SectionName;
+
+    std::unique_ptr<MemoryBuffer> BufferCopy = MemoryBuffer::getMemBufferCopy(
+        NewSection.SectionData->getBufferStart(),
+        NewSection.SectionData->getBufferIdentifier());
     Sec.Contents = makeArrayRef<uint8_t>(
-        reinterpret_cast<const uint8_t *>(Buf->getBufferStart()),
-        Buf->getBufferSize());
-    Obj.addSectionWithOwnedContents(Sec, std::move(Buf));
+        reinterpret_cast<const uint8_t *>(BufferCopy->getBufferStart()),
+        BufferCopy->getBufferSize());
+
+    Obj.addSectionWithOwnedContents(Sec, std::move(BufferCopy));
   }
 
   return Error::success();

diff  --git a/llvm/test/tools/llvm-objcopy/COFF/add-section.test b/llvm/test/tools/llvm-objcopy/COFF/add-section.test
index 61434c242e605..db241a3eef061 100644
--- a/llvm/test/tools/llvm-objcopy/COFF/add-section.test
+++ b/llvm/test/tools/llvm-objcopy/COFF/add-section.test
@@ -59,9 +59,9 @@
 
 ## Test that llvm-objcopy produces an error if the file with section contents
 ## to be added does not exist.
-# RUN: not llvm-objcopy --add-section=.another.section=%t2 %t %t3 2>&1 | FileCheck -DFILE1=%t -DFILE2=%t2 -DMSG=%errc_ENOENT %s --check-prefixes=ERR1
+# RUN: not llvm-objcopy --add-section=.another.section=%t2 %t %t3 2>&1 | FileCheck -DFILE=%t2 -DMSG=%errc_ENOENT %s --check-prefixes=ERR1
 
-# ERR1: error: '[[FILE1]]': '[[FILE2]]': [[MSG]]
+# ERR1: error: '[[FILE]]': [[MSG]]
 
 ## Another negative test for invalid --add-sections command line argument.
 # RUN: not llvm-objcopy --add-section=.another.section %t %t3 2>&1 | FileCheck %s --check-prefixes=ERR2

diff  --git a/llvm/test/tools/llvm-objcopy/ELF/add-section.test b/llvm/test/tools/llvm-objcopy/ELF/add-section.test
index 2331157fb5ac8..15a2dc302b851 100644
--- a/llvm/test/tools/llvm-objcopy/ELF/add-section.test
+++ b/llvm/test/tools/llvm-objcopy/ELF/add-section.test
@@ -49,9 +49,9 @@ Sections:
 
 ## Test that llvm-objcopy produces an error if the file with section contents
 ## to be added does not exist.
-# RUN: not llvm-objcopy --add-section=.section.name=%t.missing %t %t.out 2>&1 | FileCheck -DFILE1=%t -DFILE2=%t.missing -DMSG=%errc_ENOENT %s --check-prefixes=ERR1
+# RUN: not llvm-objcopy --add-section=.section.name=%t.missing %t %t.out 2>&1 | FileCheck -DFILE=%t.missing -DMSG=%errc_ENOENT %s --check-prefixes=ERR1
 
-# ERR1: error: '[[FILE1]]': '[[FILE2]]': [[MSG]]
+# ERR1: error: '[[FILE]]': [[MSG]]
 
 ## Negative test for invalid --add-sections argument - missing '='.
 # RUN: not llvm-objcopy --add-section=.section.name %t %t.out 2>&1 | FileCheck %s --check-prefixes=ERR2

diff  --git a/llvm/test/tools/llvm-objcopy/MachO/add-section-error.test b/llvm/test/tools/llvm-objcopy/MachO/add-section-error.test
index 41f25cae04fc5..4e5ada940fac2 100644
--- a/llvm/test/tools/llvm-objcopy/MachO/add-section-error.test
+++ b/llvm/test/tools/llvm-objcopy/MachO/add-section-error.test
@@ -6,7 +6,7 @@
 ## Error case 1: Nonexistent input file is specified by --add-section.
 # RUN: not llvm-objcopy --add-section __TEXT,__text=%t.missing %t %t.nonexistent-file 2>&1 \
 # RUN:   | FileCheck %s -DINPUT=%t -DSECTION_DATA_FILE=%t.missing -DMSG=%errc_ENOENT --check-prefix=NONEXSITENT-FILE
-# NONEXSITENT-FILE: error: '[[INPUT]]': '[[SECTION_DATA_FILE]]': [[MSG]]
+# NONEXSITENT-FILE: error: '[[SECTION_DATA_FILE]]': [[MSG]]
 
 ## Error case 2: Too long segment name.
 # RUN: not llvm-objcopy --add-section __TOOOOOOOOO_LONG,__text=%t.data %t %t.too-long-seg-name 2>&1 \

diff  --git a/llvm/test/tools/llvm-objcopy/MachO/update-section.test b/llvm/test/tools/llvm-objcopy/MachO/update-section.test
index a4fa5423ee52b..d9fb23d41474e 100644
--- a/llvm/test/tools/llvm-objcopy/MachO/update-section.test
+++ b/llvm/test/tools/llvm-objcopy/MachO/update-section.test
@@ -26,7 +26,7 @@
 # RUN: llvm-objcopy --update-section __TEXT,__text=%t.
diff  %t - | obj2yaml | FileCheck %s --check-prefix=FULL-SECNAME
 # FULL-SECNAME: content: '41414142'
 
-# RUN: not llvm-objcopy --update-section __text=%t.dff %t /dev/null 2>&1 | FileCheck %s --check-prefix=NON-CANONICAL-SECNAME
+# RUN: not llvm-objcopy --update-section __text=%t.
diff  %t /dev/null 2>&1 | FileCheck %s --check-prefix=NON-CANONICAL-SECNAME
 # NON-CANONICAL-SECNAME: error: {{.*}}invalid section name '__text' (should be formatted as '<segment name>,<section name>')
 
 --- !mach-o

diff  --git a/llvm/tools/llvm-objcopy/ObjcopyOptions.cpp b/llvm/tools/llvm-objcopy/ObjcopyOptions.cpp
index 65bbd033d3c28..db3823d1f9c89 100644
--- a/llvm/tools/llvm-objcopy/ObjcopyOptions.cpp
+++ b/llvm/tools/llvm-objcopy/ObjcopyOptions.cpp
@@ -558,6 +558,32 @@ static Expected<NewSymbolInfo> parseNewSymbolInfo(StringRef FlagValue) {
   return SI;
 }
 
+// Parse input option \p ArgValue and load section data. This function
+// extracts section name and name of the file keeping section data from
+// ArgValue, loads data from the file, and stores section name and data
+// into the vector of new sections \p NewSections.
+static Error loadNewSectionData(StringRef ArgValue, StringRef OptionName,
+                                std::vector<NewSectionInfo> &NewSections) {
+  if (!ArgValue.contains('='))
+    return createStringError(errc::invalid_argument,
+                             "bad format for " + OptionName + ": missing '='");
+
+  std::pair<StringRef, StringRef> SecPair = ArgValue.split("=");
+  if (SecPair.second.empty())
+    return createStringError(errc::invalid_argument, "bad format for " +
+                                                         OptionName +
+                                                         ": missing file name");
+
+  ErrorOr<std::unique_ptr<MemoryBuffer>> BufOrErr =
+      MemoryBuffer::getFile(SecPair.second);
+  if (!BufOrErr)
+    return createFileError(SecPair.second,
+                           errorCodeToError(BufOrErr.getError()));
+
+  NewSections.push_back({SecPair.first, std::move(*BufOrErr)});
+  return Error::success();
+}
+
 // ParseObjcopyOptions returns the config and sets the input arguments. If a
 // help flag is set then ParseObjcopyOptions will print the help messege and
 // exit.
@@ -848,26 +874,14 @@ objcopy::parseObjcopyOptions(ArrayRef<const char *> RawArgsArr,
             Arg->getValue(), SectionMatchStyle, ErrorCallback)))
       return std::move(E);
   for (auto Arg : InputArgs.filtered(OBJCOPY_add_section)) {
-    StringRef ArgValue(Arg->getValue());
-    if (!ArgValue.contains('='))
-      return createStringError(errc::invalid_argument,
-                               "bad format for --add-section: missing '='");
-    if (ArgValue.split("=").second.empty())
-      return createStringError(
-          errc::invalid_argument,
-          "bad format for --add-section: missing file name");
-    Config.AddSection.push_back(ArgValue);
+    if (Error Err = loadNewSectionData(Arg->getValue(), "--add-section",
+                                       Config.AddSection))
+      return std::move(Err);
   }
   for (auto Arg : InputArgs.filtered(OBJCOPY_update_section)) {
-    StringRef ArgValue(Arg->getValue());
-    if (!ArgValue.contains('='))
-      return createStringError(errc::invalid_argument,
-                               "bad format for --update-section: missing '='");
-    if (ArgValue.split("=").second.empty())
-      return createStringError(
-          errc::invalid_argument,
-          "bad format for --update-section: missing file name");
-    Config.UpdateSection.push_back(ArgValue);
+    if (Error Err = loadNewSectionData(Arg->getValue(), "--update-section",
+                                       Config.UpdateSection))
+      return std::move(Err);
   }
   for (auto *Arg : InputArgs.filtered(OBJCOPY_dump_section)) {
     StringRef Value(Arg->getValue());

diff  --git a/llvm/unittests/ObjCopy/ObjCopyTest.cpp b/llvm/unittests/ObjCopy/ObjCopyTest.cpp
index 8d208be988a71..bc16e954ca560 100644
--- a/llvm/unittests/ObjCopy/ObjCopyTest.cpp
+++ b/llvm/unittests/ObjCopy/ObjCopyTest.cpp
@@ -116,3 +116,252 @@ TEST(CopySimpleInMemoryFile, Wasm) {
 )",
       [](const Binary &File) { return File.isWasm(); });
 }
+
+enum Action : uint8_t { AddSection, UpdateSection };
+
+void addOrUpdateSectionToFileImpl(
+    const char *YamlCreationString,
+    std::function<bool(const Binary &File)> IsValidFormat,
+    StringRef NewSectionName, StringRef NewSectionData, Action SectionAction) {
+  auto ErrHandler = [&](const Twine &Msg) { FAIL() << "Error: " << Msg; };
+
+  // Create Object file from YAML description.
+  SmallVector<char> Storage;
+  std::unique_ptr<ObjectFile> Obj =
+      yaml2ObjectFile(Storage, YamlCreationString, ErrHandler);
+  ASSERT_TRUE(Obj);
+  ASSERT_TRUE(IsValidFormat(*Obj));
+
+  std::unique_ptr<MemoryBuffer> NewSectionBuffer =
+      MemoryBuffer::getMemBuffer(NewSectionData, NewSectionName, false);
+  std::string Name;
+  if (Obj->isMachO())
+    Name = "__TEXT," + NewSectionName.str();
+  else
+    Name = NewSectionName.str();
+
+  ConfigManager Config;
+  Config.Common.OutputFilename = "a.out";
+  if (SectionAction == AddSection)
+    Config.Common.AddSection.push_back({Name, std::move(NewSectionBuffer)});
+  else
+    Config.Common.UpdateSection.push_back({Name, std::move(NewSectionBuffer)});
+
+  // Call executeObjcopyOnBinary()
+  SmallVector<char> DataVector;
+  raw_svector_ostream OutStream(DataVector);
+  Error Err = objcopy::executeObjcopyOnBinary(Config, *Obj.get(), OutStream);
+  ASSERT_FALSE(std::move(Err));
+
+  MemoryBufferRef Buffer(StringRef(DataVector.data(), DataVector.size()),
+                         Config.Common.OutputFilename);
+
+  // Check copied file.
+  Expected<std::unique_ptr<Binary>> Result = createBinary(Buffer);
+  ASSERT_THAT_EXPECTED(Result, Succeeded());
+  ASSERT_TRUE(IsValidFormat(**Result));
+  ASSERT_TRUE((*Result)->isObject());
+
+  // Check that copied file has the new section.
+  bool HasNewSection = false;
+  for (const object::SectionRef &Sect :
+       static_cast<ObjectFile *>((*Result).get())->sections()) {
+    Expected<StringRef> SectNameOrErr = Sect.getName();
+    ASSERT_THAT_EXPECTED(SectNameOrErr, Succeeded());
+
+    if (*SectNameOrErr == NewSectionName) {
+      HasNewSection = true;
+      Expected<StringRef> SectionData = Sect.getContents();
+      ASSERT_THAT_EXPECTED(SectionData, Succeeded());
+      EXPECT_TRUE(Sect.getSize() == NewSectionData.size());
+      EXPECT_TRUE(memcmp(SectionData->data(), NewSectionData.data(),
+                         NewSectionData.size()) == 0);
+      break;
+    }
+  }
+  EXPECT_TRUE(HasNewSection);
+}
+
+TEST(AddSection, COFF) {
+  SCOPED_TRACE("addSectionToFileCOFF");
+
+  addOrUpdateSectionToFileImpl(
+      R"(
+--- !COFF
+header:
+  Machine:         IMAGE_FILE_MACHINE_AMD64
+  Characteristics: [  ]
+sections:
+  - Name:            .text
+    Characteristics: [  ]
+    Alignment:       4
+    SectionData:     E800000000C3C3C3
+symbols:
+...
+)",
+      [](const Binary &File) { return File.isCOFF(); }, ".foo", "1234",
+      AddSection);
+}
+
+TEST(AddSection, ELF) {
+  SCOPED_TRACE("addSectionToFileELF");
+
+  addOrUpdateSectionToFileImpl(
+      R"(
+--- !ELF
+FileHeader:
+   Class:    ELFCLASS64
+   Data:     ELFDATA2LSB
+   Type:     ET_REL)",
+      [](const Binary &File) { return File.isELF(); }, ".foo", "1234",
+      AddSection);
+}
+
+TEST(AddSection, MachO) {
+  SCOPED_TRACE("addSectionToFileMachO");
+
+  addOrUpdateSectionToFileImpl(
+      R"(
+--- !mach-o
+FileHeader:
+  magic:           0xFEEDFACF
+  cputype:         0x01000007
+  cpusubtype:      0x80000003
+  filetype:        0x00000001
+  ncmds:           1
+  sizeofcmds:      152
+  flags:           0x00002000
+  reserved:        0x00000000
+LoadCommands:
+  - cmd:             LC_SEGMENT_64
+    cmdsize:         152
+    segname:         __TEXT
+    vmaddr:          0
+    vmsize:          4
+    fileoff:         184
+    filesize:        4
+    maxprot:         7
+    initprot:        7
+    nsects:          1
+    flags:           0
+    Sections:
+      - sectname:        __text
+        segname:         __TEXT
+        addr:            0x0000000000000000
+        content:         'AABBCCDD'
+        size:            4
+        offset:          184
+        align:           0
+        reloff:          0x00000000
+        nreloc:          0
+        flags:           0x80000400
+        reserved1:       0x00000000
+        reserved2:       0x00000000
+        reserved3:       0x00000000
+...
+)",
+      [](const Binary &File) { return File.isMachO(); }, "__foo", "1234",
+      AddSection);
+}
+
+TEST(AddSection, Wasm) {
+  SCOPED_TRACE("addSectionToFileWasm");
+
+  addOrUpdateSectionToFileImpl(
+      R"(
+--- !WASM
+FileHeader:
+  Version:         0x00000001
+...
+)",
+      [](const Binary &File) { return File.isWasm(); }, ".foo", "1234",
+      AddSection);
+}
+
+TEST(UpdateSection, COFF) {
+  SCOPED_TRACE("updateSectionToFileCOFF");
+
+  addOrUpdateSectionToFileImpl(
+      R"(
+--- !COFF
+header:
+  Machine:         IMAGE_FILE_MACHINE_AMD64
+  Characteristics: [  ]
+sections:
+  - Name:            .foo
+    Characteristics: [  ]
+    Alignment:       4
+    SectionData:     E800000000C3C3C3
+symbols:
+...
+)",
+      [](const Binary &File) { return File.isCOFF(); }, ".foo", "1234",
+      UpdateSection);
+}
+
+TEST(UpdateSection, ELF) {
+  SCOPED_TRACE("updateSectionToFileELF");
+
+  addOrUpdateSectionToFileImpl(
+      R"(
+--- !ELF
+FileHeader:
+  Class:    ELFCLASS64
+  Data:     ELFDATA2LSB
+  Type:     ET_REL
+Sections:
+  - Name:            .foo
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC ]
+    Content:        "12345678"
+)",
+      [](const Binary &File) { return File.isELF(); }, ".foo", "1234",
+      UpdateSection);
+}
+
+TEST(UpdateSection, MachO) {
+  SCOPED_TRACE("updateSectionToFileMachO");
+
+  addOrUpdateSectionToFileImpl(
+      R"(
+--- !mach-o
+FileHeader:
+  magic:           0xFEEDFACF
+  cputype:         0x01000007
+  cpusubtype:      0x80000003
+  filetype:        0x00000001
+  ncmds:           1
+  sizeofcmds:      152
+  flags:           0x00002000
+  reserved:        0x00000000
+LoadCommands:
+  - cmd:             LC_SEGMENT_64
+    cmdsize:         152
+    segname:         __TEXT
+    vmaddr:          0
+    vmsize:          4
+    fileoff:         184
+    filesize:        4
+    maxprot:         7
+    initprot:        7
+    nsects:          1
+    flags:           0
+    Sections:
+      - sectname:        __foo
+        segname:         __TEXT
+        addr:            0x0000000000000000
+        content:         'AABBCCDD'
+        size:            4
+        offset:          184
+        align:           0
+        reloff:          0x00000000
+        nreloc:          0
+        flags:           0x80000400
+        reserved1:       0x00000000
+        reserved2:       0x00000000
+        reserved3:       0x00000000
+...
+)",
+      [](const Binary &File) { return File.isMachO(); }, "__foo", "1234",
+      UpdateSection);
+}


        


More information about the llvm-commits mailing list