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

[abrachet]

Summary

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

[abrachet]

From what I can gather about Mach-O, it seems like maybe updateSection could just be removeSection + addSection, but I wasn't sure. Maybe someone with more Mach-O experience could say for sure.

[jhenderson]

Yeah, I'm not really qualified to do this review, as it's fairly format-specific heavy, so I'll defer to others to do it.

[alexander-shaposhnikov]

thanks for the patch! I'll review it over the weekend

[abrachet]

ping :)

[alexander-shaposhnikov]

llvm/tools/llvm-objcopy/MachO/MachOObjcopy.cpp

343

we probably don't need 2 cases here (specified / unspecified segment name),
for consistency it would be better to always require the full name, e.g. TEXT,text
(otool -d also understands this notation, llvm-objcopy --add-section uses it as well)

424

nit:

std::tie(SectionName, FileName) = Flag.split('=');

[abrachet]

Remove ability to update section without prepending segment name

[alexander-shaposhnikov]

llvm/tools/llvm-objcopy/MachO/MachOObjcopy.cpp

344

immediately called lambda is unnecessary here,
let's split it out into a helper function

426

we need to validate the name (isValidMachOCannonicalName can be of some help here) (+ test case)

[alexander-shaposhnikov]

small nit, otherwise LGTM

llvm/tools/llvm-objcopy/MachO/MachOObjcopy.cpp

320

nit: const Object &O

llvm/test/tools/llvm-objcopy/MachO/update-section.test

+# 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

llvm/tools/llvm-objcopy/MachO/MachOObjcopy.cpp

static Error addSection(StringRef SecName, StringRef Filename, Object &Obj) {
   // Insert a new section into it.
   LoadCommand &NewSegment =
       Obj.addSegment(TargetSegName, alignTo(Sec.Size, 16384));
   NewSegment.Sections.push_back(std::make_unique<Section>(Sec));
   NewSegment.Sections.back()->Addr = *NewSegment.getSegmentVMAddr();
   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.
 static Error isValidMachOCannonicalName(StringRef Name) {
   if (Name.count(',') != 1)
     return createStringError(errc::invalid_argument,
                              "invalid section name '%s' (should be formatted "
                              "as '<segment name>,<section name>')",
for (const auto &Flag : Config.AddSection) {
     StringRef SecName = SecPair.first;
     StringRef File = SecPair.second;
     if (Error E = isValidMachOCannonicalName(SecName))
       return E;
     if (Error E = addSection(SecName, File, Obj))
       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;
 
   return Error::success();
 }
 
 Error objcopy::macho::executeObjcopyOnBinary(const CommonConfig &Config,
                                              const MachOConfig &MachOConfig,