[llvm-objcopy][MachO] Implement --update-sectionbase_offset

Implements `--update-section` which is currently supported for ELF for Mach-O as well

Reviewed By: alexander-shaposhnikov

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

2 files changed, 171 insertions, 0 deletions

diff --git a/llvm/test/tools/llvm-objcopy/MachO/update-section.test b/llvm/test/tools/llvm-objcopy/MachO/update-section.test
new file mode 100644
index 000000000000..a4fa5423ee52
--- /dev/null
+++ b/llvm/test/tools/llvm-objcopy/MachO/update-section.test
@@ -0,0 +1,115 @@
+# RUN: echo -n AAAB > %t.diff
+# RUN: echo -n AAA > %t.smaller
+# RUN: echo -n AAAAAAAAA > %t.larger
+
+# RUN: yaml2obj --docnum=1 %s -o %t
+
+# RUN: llvm-objcopy --update-section __TEXT,__text=%t.diff %t - | obj2yaml | FileCheck %s
+# CHECK: content: '41414142'
+
+# RUN: llvm-objcopy --update-section __TEXT,__text=%t.smaller %t - | obj2yaml | FileCheck %s --check-prefix=SMALLER
+# SMALLER: content: '414141'
+
+# RUN: not llvm-objcopy --update-section __TEXT,__text=%t.larger %t /dev/null 2>&1 | FileCheck %s --check-prefix=TOO-LARGE
+# TOO-LARGE: error: {{.*}}new section cannot be larger than previous section
+
+# RUN: not llvm-objcopy --update-section __TEXT,__text=%t.noexist %t /dev/null
+
+# RUN: not llvm-objcopy --update-section __NOEXIST,__text=%t.diff %t /dev/null 2>&1 | FileCheck %s --check-prefix=NO-SEGMENT
+# NO-SEGMENT: error: {{.*}}could not find segment with name '__NOEXIST'
+
+# RUN: not llvm-objcopy --update-section __TEXT,__noexist=%t.diff %t /dev/null 2>&1 | FileCheck %s --check-prefix=NO-SECTION
+# NO-SECTION: error: {{.*}}could not find section with name '__noexist'
+
+# RUN: yaml2obj --docnum=2 %s -o %t
+
+# 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
+# NON-CANONICAL-SECNAME: error: {{.*}}invalid section name '__text' (should be formatted as '<segment name>,<section name>')
+
+--- !mach-o
+FileHeader:
+  magic:           0xFEEDFACF
+  cputype:         0x01000007
+  cpusubtype:      0x00000003
+  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:         '41414141'
+        size:            4
+        offset:          184
+        align:           0
+        reloff:          0x00000000
+        nreloc:          0
+        flags:           0x80000400
+        reserved1:       0x00000000
+        reserved2:       0x00000000
+        reserved3:       0x00000000
+
+--- !mach-o
+FileHeader:
+  magic:           0xFEEDFACF
+  cputype:         0x01000007
+  cpusubtype:      0x00000003
+  filetype:        0x00000001
+  ncmds:           1
+  sizeofcmds:      312
+  flags:           0x00002000
+  reserved:        0x00000000
+LoadCommands:
+  - cmd:             LC_SEGMENT_64
+    cmdsize:         312
+    segname:         '__TEXT'
+    vmaddr:          0
+    vmsize:          12
+    fileoff:         344
+    filesize:        12
+    maxprot:         7
+    initprot:        7
+    nsects:          3
+    flags:           0
+    Sections:
+      - sectname:        __text
+        segname:         __TEXT
+        addr:            0x0000000000000000
+        content:         'AABBCCDD'
+        size:            4
+        offset:          344
+        align:           0
+        reloff:          0x00000000
+        nreloc:          0
+        flags:           0x80000400
+        reserved1:       0x00000000
+        reserved2:       0x00000000
+      - sectname:        __text
+        segname:         __TEXT2
+        addr:            0x0000000000000004
+        content:         ''
+        size:            0
+        offset:          348
+        align:           0
+        reloff:          0x00000000
+        nreloc:          0
+        flags:           0x00000000
+        reserved1:       0x00000000
+        reserved2:       0x00000000
diff --git a/llvm/tools/llvm-objcopy/MachO/MachOObjcopy.cpp b/llvm/tools/llvm-objcopy/MachO/MachOObjcopy.cpp
index 915394b65b12..0f92ca516bef 100644
--- a/llvm/tools/llvm-objcopy/MachO/MachOObjcopy.cpp
+++ b/llvm/tools/llvm-objcopy/MachO/MachOObjcopy.cpp
@@ -317,6 +317,52 @@ static Error addSection(StringRef SecName, StringRef Filename, Object &Obj) {
   return Error::success();
 }
 
+static Expected<Section &> findSection(StringRef SecName, Object &O) {
+  StringRef SegName;
+  std::tie(SegName, SecName) = SecName.split(",");
+  auto FoundSeg =
+      llvm::find_if(O.LoadCommands, [SegName](const LoadCommand &LC) {
+        return LC.getSegmentName() == SegName;
+      });
+  if (FoundSeg == O.LoadCommands.end())
+    return createStringError(errc::invalid_argument,
+                             "could not find segment with name '%s'",
+                             SegName.str().c_str());
+  auto FoundSec = llvm::find_if(FoundSeg->Sections,
+                                [SecName](const std::unique_ptr<Section> &Sec) {
+                                  return Sec->Sectname == SecName;
+                                });
+  if (FoundSec == FoundSeg->Sections.end())
+    return createStringError(errc::invalid_argument,
+                             "could not find section with name '%s'",
+                             SecName.str().c_str());
+
+  assert(FoundSec->get()->CanonicalName == (SegName + "," + SecName).str());
+  return *FoundSec->get();
+}
+
+static Error updateSection(StringRef SecName, StringRef Filename, Object &O) {
+  Expected<Section &> SecToUpdateOrErr = findSection(SecName, 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)
+    return createStringError(
+        errc::invalid_argument,
+        "new section cannot be larger than previous section");
+  Sec.Content = O.NewSectionsContents.save(Buf->getBuffer());
+  Sec.Size = Sec.Content.size();
+  return Error::success();
+}
+
 // isValidMachOCannonicalName returns success if Name is a MachO cannonical name
 // ("<segment>,<section>") and lengths of both segment and section names are
 // valid.
@@ -374,6 +420,16 @@ static Error handleArgs(const CommonConfig &Config,
       return E;
   }
 
+  for (const auto &Flag : Config.UpdateSection) {
+    StringRef SectionName;
+    StringRef FileName;
+    std::tie(SectionName, FileName) = Flag.split('=');
+    if (Error E = isValidMachOCannonicalName(SectionName))
+      return E;
+    if (Error E = updateSection(SectionName, FileName, Obj))
+      return E;
+  }
+
   if (Error E = processLoadCommands(MachOConfig, Obj))
     return E;