[llvm-objcopy] Implement --set-section-flags.

Summary:
--set-section-flags is used to change the section flags (e.g. SHF_ALLOC) for given sections. The flags allowed are the same from the existing --rename-section=.old=.new[,flags] feature.

Additionally, make sure that --set-section-flag cannot be used with --rename-section (either the source or destination), since --rename-section accepts flags. This avoids ambiguity for something like "--rename-section=.foo=.bar,alloc --set-section-flag=.bar,code".

Reviewers: jhenderson, jakehehrlich, alexshap, espindola

Reviewed By: jhenderson, jakehehrlich

Subscribers: llvm-commits, emaste, arichardson

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

llvm-svn: 352505
diff --git a/llvm/tools/llvm-objcopy/COFF/COFFObjcopy.cpp b/llvm/tools/llvm-objcopy/COFF/COFFObjcopy.cpp
index b7b3d3c..cf7e2ae 100644
--- a/llvm/tools/llvm-objcopy/COFF/COFFObjcopy.cpp
+++ b/llvm/tools/llvm-objcopy/COFF/COFFObjcopy.cpp
@@ -177,10 +177,10 @@
       !Config.SymbolsToGlobalize.empty() || !Config.SymbolsToKeep.empty() ||
       !Config.SymbolsToLocalize.empty() || !Config.SymbolsToWeaken.empty() ||
       !Config.SymbolsToKeepGlobal.empty() || !Config.SectionsToRename.empty() ||
-      !Config.SymbolsToRename.empty() || Config.ExtractDWO ||
-      Config.KeepFileSymbols || Config.LocalizeHidden || Config.PreserveDates ||
-      Config.StripDWO || Config.StripNonAlloc || Config.StripSections ||
-      Config.Weaken || Config.DecompressDebugSections) {
+      !Config.SetSectionFlags.empty() || !Config.SymbolsToRename.empty() ||
+      Config.ExtractDWO || Config.KeepFileSymbols || Config.LocalizeHidden ||
+      Config.PreserveDates || Config.StripDWO || Config.StripNonAlloc ||
+      Config.StripSections || Config.Weaken || Config.DecompressDebugSections) {
     return createStringError(llvm::errc::invalid_argument,
                              "Option not supported by llvm-objcopy for COFF");
   }
diff --git a/llvm/tools/llvm-objcopy/CopyConfig.cpp b/llvm/tools/llvm-objcopy/CopyConfig.cpp
index 38c9de3..9702eb6 100644
--- a/llvm/tools/llvm-objcopy/CopyConfig.cpp
+++ b/llvm/tools/llvm-objcopy/CopyConfig.cpp
@@ -128,6 +128,32 @@
       .Default(SectionFlag::SecNone);
 }
 
+static uint64_t parseSectionFlagSet(ArrayRef<StringRef> SectionFlags) {
+  SectionFlag ParsedFlags = SectionFlag::SecNone;
+  for (StringRef Flag : SectionFlags) {
+    SectionFlag ParsedFlag = parseSectionRenameFlag(Flag);
+    if (ParsedFlag == SectionFlag::SecNone)
+      error("Unrecognized section flag '" + Flag +
+            "'. Flags supported for GNU compatibility: alloc, load, noload, "
+            "readonly, debug, code, data, rom, share, contents, merge, "
+            "strings.");
+    ParsedFlags |= ParsedFlag;
+  }
+
+  uint64_t NewFlags = 0;
+  if (ParsedFlags & SectionFlag::SecAlloc)
+    NewFlags |= ELF::SHF_ALLOC;
+  if (!(ParsedFlags & SectionFlag::SecReadonly))
+    NewFlags |= ELF::SHF_WRITE;
+  if (ParsedFlags & SectionFlag::SecCode)
+    NewFlags |= ELF::SHF_EXECINSTR;
+  if (ParsedFlags & SectionFlag::SecMerge)
+    NewFlags |= ELF::SHF_MERGE;
+  if (ParsedFlags & SectionFlag::SecStrings)
+    NewFlags |= ELF::SHF_STRINGS;
+  return NewFlags;
+}
+
 static SectionRename parseRenameSectionValue(StringRef FlagValue) {
   if (!FlagValue.contains('='))
     error("Bad format for --rename-section: missing '='");
@@ -142,34 +168,29 @@
   Old2New.second.split(NameAndFlags, ',');
   SR.NewName = NameAndFlags[0];
 
-  if (NameAndFlags.size() > 1) {
-    SectionFlag Flags = SectionFlag::SecNone;
-    for (size_t I = 1, Size = NameAndFlags.size(); I < Size; ++I) {
-      SectionFlag Flag = parseSectionRenameFlag(NameAndFlags[I]);
-      if (Flag == SectionFlag::SecNone)
-        error("Unrecognized section flag '" + NameAndFlags[I] +
-              "'. Flags supported for GNU compatibility: alloc, load, noload, "
-              "readonly, debug, code, data, rom, share, contents, merge, "
-              "strings.");
-      Flags |= Flag;
-    }
-
-    SR.NewFlags = 0;
-    if (Flags & SectionFlag::SecAlloc)
-      *SR.NewFlags |= ELF::SHF_ALLOC;
-    if (!(Flags & SectionFlag::SecReadonly))
-      *SR.NewFlags |= ELF::SHF_WRITE;
-    if (Flags & SectionFlag::SecCode)
-      *SR.NewFlags |= ELF::SHF_EXECINSTR;
-    if (Flags & SectionFlag::SecMerge)
-      *SR.NewFlags |= ELF::SHF_MERGE;
-    if (Flags & SectionFlag::SecStrings)
-      *SR.NewFlags |= ELF::SHF_STRINGS;
-  }
+  if (NameAndFlags.size() > 1)
+    SR.NewFlags = parseSectionFlagSet(makeArrayRef(NameAndFlags).drop_front());
 
   return SR;
 }
 
+static SectionFlagsUpdate parseSetSectionFlagValue(StringRef FlagValue) {
+  if (!StringRef(FlagValue).contains('='))
+    error("Bad format for --set-section-flags: missing '='");
+
+  // Initial split: ".foo" = "f1,f2,..."
+  auto Section2Flags = StringRef(FlagValue).split('=');
+  SectionFlagsUpdate SFU;
+  SFU.Name = Section2Flags.first;
+
+  // Flags split: "f1" "f2" ...
+  SmallVector<StringRef, 6> SectionFlags;
+  Section2Flags.second.split(SectionFlags, ',');
+  SFU.NewFlags = parseSectionFlagSet(SectionFlags);
+
+  return SFU;
+}
+
 static const StringMap<MachineInfo> ArchMap{
     // Name, {EMachine, 64bit, LittleEndian}
     {"aarch64", {ELF::EM_AARCH64, true, true}},
@@ -327,6 +348,24 @@
     if (!Config.SectionsToRename.try_emplace(SR.OriginalName, SR).second)
       error("Multiple renames of section " + SR.OriginalName);
   }
+  for (auto Arg : InputArgs.filtered(OBJCOPY_set_section_flags)) {
+    SectionFlagsUpdate SFU = parseSetSectionFlagValue(Arg->getValue());
+    if (!Config.SetSectionFlags.try_emplace(SFU.Name, SFU).second)
+      error("--set-section-flags set multiple times for section " + SFU.Name);
+  }
+  // Prohibit combinations of --set-section-flags when the section name is used
+  // by --rename-section, either as a source or a destination.
+  for (const auto &E : Config.SectionsToRename) {
+    const SectionRename &SR = E.second;
+    if (Config.SetSectionFlags.count(SR.OriginalName))
+      error("--set-section-flags=" + SR.OriginalName +
+            " conflicts with --rename-section=" + SR.OriginalName + "=" +
+            SR.NewName);
+    if (Config.SetSectionFlags.count(SR.NewName))
+      error("--set-section-flags=" + SR.NewName +
+            " conflicts with --rename-section=" + SR.OriginalName + "=" +
+            SR.NewName);
+  }
 
   for (auto Arg : InputArgs.filtered(OBJCOPY_remove_section))
     Config.ToRemove.push_back(Arg->getValue());
diff --git a/llvm/tools/llvm-objcopy/CopyConfig.h b/llvm/tools/llvm-objcopy/CopyConfig.h
index 103e12a..c9af373 100644
--- a/llvm/tools/llvm-objcopy/CopyConfig.h
+++ b/llvm/tools/llvm-objcopy/CopyConfig.h
@@ -37,6 +37,11 @@
   Optional<uint64_t> NewFlags;
 };
 
+struct SectionFlagsUpdate {
+  StringRef Name;
+  uint64_t NewFlags;
+};
+
 // Configuration for copying/stripping a single file.
 struct CopyConfig {
   // Main input/output options
@@ -73,6 +78,7 @@
 
   // Map options
   StringMap<SectionRename> SectionsToRename;
+  StringMap<SectionFlagsUpdate> SetSectionFlags;
   StringMap<StringRef> SymbolsToRename;
 
   // Boolean options
diff --git a/llvm/tools/llvm-objcopy/ELF/ELFObjcopy.cpp b/llvm/tools/llvm-objcopy/ELF/ELFObjcopy.cpp
index 2a52f1f..9259996 100644
--- a/llvm/tools/llvm-objcopy/ELF/ELFObjcopy.cpp
+++ b/llvm/tools/llvm-objcopy/ELF/ELFObjcopy.cpp
@@ -70,6 +70,17 @@
   return !isDWOSection(Sec);
 }
 
+static uint64_t setSectionFlagsPreserveMask(uint64_t OldFlags,
+                                            uint64_t NewFlags) {
+  // Preserve some flags which should not be dropped when setting flags.
+  // Also, preserve anything OS/processor dependant.
+  const uint64_t PreserveMask = ELF::SHF_COMPRESSED | ELF::SHF_EXCLUDE |
+                                ELF::SHF_GROUP | ELF::SHF_LINK_ORDER |
+                                ELF::SHF_MASKOS | ELF::SHF_MASKPROC |
+                                ELF::SHF_TLS | ELF::SHF_INFO_LINK;
+  return (OldFlags & PreserveMask) | (NewFlags & ~PreserveMask);
+}
+
 static ElfType getOutputElfType(const Binary &Bin) {
   // Infer output ELF type from the input ELF object
   if (isa<ELFObjectFile<ELF32LE>>(Bin))
@@ -484,16 +495,19 @@
       if (Iter != Config.SectionsToRename.end()) {
         const SectionRename &SR = Iter->second;
         Sec.Name = SR.NewName;
-        if (SR.NewFlags.hasValue()) {
-          // Preserve some flags which should not be dropped when setting flags.
-          // Also, preserve anything OS/processor dependant.
-          const uint64_t PreserveMask = ELF::SHF_COMPRESSED | ELF::SHF_EXCLUDE |
-                                        ELF::SHF_GROUP | ELF::SHF_LINK_ORDER |
-                                        ELF::SHF_MASKOS | ELF::SHF_MASKPROC |
-                                        ELF::SHF_TLS | ELF::SHF_INFO_LINK;
-          Sec.Flags = (Sec.Flags & PreserveMask) |
-                      (SR.NewFlags.getValue() & ~PreserveMask);
-        }
+        if (SR.NewFlags.hasValue())
+          Sec.Flags =
+              setSectionFlagsPreserveMask(Sec.Flags, SR.NewFlags.getValue());
+      }
+    }
+  }
+
+  if (!Config.SetSectionFlags.empty()) {
+    for (auto &Sec : Obj.sections()) {
+      const auto Iter = Config.SetSectionFlags.find(Sec.Name);
+      if (Iter != Config.SetSectionFlags.end()) {
+        const SectionFlagsUpdate &SFU = Iter->second;
+        Sec.Flags = setSectionFlagsPreserveMask(Sec.Flags, SFU.NewFlags);
       }
     }
   }
diff --git a/llvm/tools/llvm-objcopy/ObjcopyOpts.td b/llvm/tools/llvm-objcopy/ObjcopyOpts.td
index 57d8bf1..9c54ad4 100644
--- a/llvm/tools/llvm-objcopy/ObjcopyOpts.td
+++ b/llvm/tools/llvm-objcopy/ObjcopyOpts.td
@@ -86,6 +86,13 @@
          "Make a section named <section> with the contents of <file>.">,
       MetaVarName<"section=file">;
 
+defm set_section_flags
+    : Eq<"set-section-flags",
+         "Set section flags for a given section. Flags supported for GNU "
+         "compatibility: alloc, load, noload, readonly, debug, code, data, "
+         "rom, share, contents, merge, strings.">,
+      MetaVarName<"section=flag1[,flag2,...]">;
+
 def strip_all
     : Flag<["-", "--"], "strip-all">,
       HelpText<