[lld-macho] Implement -segprot

Addresses llvm.org/PR49405.

Reviewed By: #lld-macho, oontvoo

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

GitOrigin-RevId: a43f588e0128bfb5ba10f6dbe7069e55bfa5ec51
diff --git a/MachO/Config.h b/MachO/Config.h
index 810ae7b..c629c84 100644
--- a/MachO/Config.h
+++ b/MachO/Config.h
@@ -50,6 +50,12 @@
   dynamic_lookup,
 };
 
+struct SegmentProtection {
+  llvm::StringRef name;
+  uint32_t maxProt;
+  uint32_t initProt;
+};
+
 class SymbolPatterns {
 public:
   // GlobPattern can also match literals,
@@ -104,6 +110,9 @@
   std::vector<llvm::StringRef> frameworkSearchPaths;
   std::vector<llvm::StringRef> runtimePaths;
   std::vector<Symbol *> explicitUndefineds;
+  // There are typically very few custom segmentProtections, so use a vector
+  // instead of a map.
+  std::vector<SegmentProtection> segmentProtections;
 
   llvm::DenseMap<llvm::StringRef, SymbolPriorityEntry> priorities;
   SectionRenameMap sectionRenameMap;
diff --git a/MachO/Driver.cpp b/MachO/Driver.cpp
index d262ad5..96f8195 100644
--- a/MachO/Driver.cpp
+++ b/MachO/Driver.cpp
@@ -724,6 +724,29 @@
   return version.rawValue();
 }
 
+static uint32_t parseProtection(StringRef protStr) {
+  uint32_t prot = 0;
+  for (char c : protStr) {
+    switch (c) {
+    case 'r':
+      prot |= VM_PROT_READ;
+      break;
+    case 'w':
+      prot |= VM_PROT_WRITE;
+      break;
+    case 'x':
+      prot |= VM_PROT_EXECUTE;
+      break;
+    case '-':
+      break;
+    default:
+      error("unknown -segprot letter '" + Twine(c) + "' in " + protStr);
+      return 0;
+    }
+  }
+  return prot;
+}
+
 void SymbolPatterns::clear() {
   literals.clear();
   globs.clear();
@@ -966,6 +989,18 @@
         validName(arg->getValue(1));
   }
 
+  for (const Arg *arg : args.filtered(OPT_segprot)) {
+    StringRef segName = arg->getValue(0);
+    uint32_t maxProt = parseProtection(arg->getValue(1));
+    uint32_t initProt = parseProtection(arg->getValue(2));
+    if (maxProt != initProt && config->target.Arch != AK_i386)
+      error("invalid argument '" + arg->getAsString(args) +
+            "': max and init must be the same for non-i386 archs");
+    if (segName == segment_names::linkEdit)
+      error("-segprot cannot be used to change __LINKEDIT's protections");
+    config->segmentProtections.push_back({segName, maxProt, initProt});
+  }
+
   handleSymbolPatterns(args, config->exportedSymbols, OPT_exported_symbol,
                        OPT_exported_symbols_list);
   handleSymbolPatterns(args, config->unexportedSymbols, OPT_unexported_symbol,
diff --git a/MachO/Options.td b/MachO/Options.td
index da2c53e..6a963ff 100644
--- a/MachO/Options.td
+++ b/MachO/Options.td
@@ -772,8 +772,7 @@
     Group<grp_rare>;
 def segprot : MultiArg<["-"], "segprot", 3>,
     MetaVarName<"<segment> <max> <init>">,
-    HelpText<"Specifies the <max> and <init> virtual memory protection of <segment> as r/w/x/-seg_addr_table path Specify hex base addresses and dylib install names on successive lines in <path>. This option is obsolete">,
-    Flags<[HelpHidden]>,
+    HelpText<"Specifies the <max> and <init> virtual memory protection of <segment> as r/w/x/-seg_addr_table path">,
     Group<grp_rare>;
 def segs_read_write_addr : Separate<["-"], "segs_read_write_addr">,
     MetaVarName<"<address>">,
diff --git a/MachO/OutputSegment.cpp b/MachO/OutputSegment.cpp
index edfc474..71da2ef 100644
--- a/MachO/OutputSegment.cpp
+++ b/MachO/OutputSegment.cpp
@@ -21,6 +21,12 @@
 using namespace lld::macho;
 
 static uint32_t initProt(StringRef name) {
+  auto it = find_if(
+      config->segmentProtections,
+      [&](const SegmentProtection &segprot) { return segprot.name == name; });
+  if (it != config->segmentProtections.end())
+    return it->initProt;
+
   if (name == segment_names::text)
     return VM_PROT_READ | VM_PROT_EXECUTE;
   if (name == segment_names::pageZero)
diff --git a/test/MachO/segprot.s b/test/MachO/segprot.s
new file mode 100644
index 0000000..a4e91d1
--- /dev/null
+++ b/test/MachO/segprot.s
@@ -0,0 +1,48 @@
+# REQUIRES: x86
+# RUN: llvm-mc -filetype=obj -triple=x86_64-apple-darwin %s -o %t.o
+
+## Make sure the option parser doesn't think --x and -w are flags.
+# RUN: %lld -dylib -o %t %t.o -segprot FOO rwx xwr -segprot BAR --x --x -segprot BAZ -w -w
+# RUN: llvm-readobj --macho-segment %t | FileCheck %s
+
+# CHECK:        Name: FOO
+# CHECK-NEXT:   Size:
+# CHECK-NEXT:   vmaddr:
+# CHECK-NEXT:   vmsize:
+# CHECK-NEXT:   fileoff:
+# CHECK-NEXT:   filesize:
+# CHECK-NEXT:   maxprot: rwx
+# CHECK-NEXT:   initprot: rwx
+
+# CHECK:        Name: BAR
+# CHECK-NEXT:   Size:
+# CHECK-NEXT:   vmaddr:
+# CHECK-NEXT:   vmsize:
+# CHECK-NEXT:   fileoff:
+# CHECK-NEXT:   filesize:
+# CHECK-NEXT:   maxprot: --x
+# CHECK-NEXT:   initprot: --x
+
+# CHECK:        Name: BAZ
+# CHECK-NEXT:   Size:
+# CHECK-NEXT:   vmaddr:
+# CHECK-NEXT:   vmsize:
+# CHECK-NEXT:   fileoff:
+# CHECK-NEXT:   filesize:
+# CHECK-NEXT:   maxprot: -w-
+# CHECK-NEXT:   initprot: -w-
+
+# RUN: not %lld -dylib -o /dev/null %t.o -segprot FOO rwx rw 2>&1 | FileCheck %s --check-prefix=MISMATCH
+# RUN: not %lld -dylib -o /dev/null %t.o -segprot __LINKEDIT rwx rwx 2>&1 | FileCheck %s --check-prefix=NO-LINKEDIT
+# RUN: not %lld -dylib -o /dev/null %t.o -segprot FOO uhh wat 2>&1 | FileCheck %s --check-prefix=MISPARSE
+# RUN: not %lld -dylib -o /dev/null %t.o -segprot FOO rwx 2>&1 | FileCheck %s --check-prefix=MISSING
+
+# MISMATCH:    error: invalid argument '-segprot FOO rwx rw': max and init must be the same for non-i386 archs
+# NO-LINKEDIT: error: -segprot cannot be used to change __LINKEDIT's protections
+# MISPARSE:    error: unknown -segprot letter 'u' in uhh
+# MISPARSE:    error: unknown -segprot letter 'a' in wat
+# MISSING:     error: -segprot: missing argument
+
+.section FOO,foo
+.section BAR,bar
+.section BAZ,baz