[lld-macho] Support LC_ENCRYPTION_INFO

This load command records a range spanning from the end of the load
commands to the end of the `__TEXT` segment. Presumably the kernel will encrypt
all this data.

Reviewed By: #lld-macho, thakis

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

lld/MachO/Config.h

@@ -90,6 +90,7 @@ struct Configuration {
   bool adhocCodesign = false;
   bool emitFunctionStarts = false;
   bool emitBitcodeBundle = false;
+  bool emitEncryptionInfo = false;
   bool timeTraceEnabled = false;
   uint32_t headerPad;
   uint32_t dylibCompatibilityVersion = 0;

lld/MachO/Driver.cpp

@@ -963,6 +963,12 @@ bool macho::link(ArrayRef<const char *> argsArr, bool canExitEarly,
   config->emitFunctionStarts = !args.hasArg(OPT_no_function_starts);
   config->emitBitcodeBundle = args.hasArg(OPT_bitcode_bundle);
 
+  std::array<PlatformKind, 3> encryptablePlatforms{
+      PlatformKind::iOS, PlatformKind::watchOS, PlatformKind::tvOS};
+  config->emitEncryptionInfo = args.hasFlag(
+      OPT_encryptable, OPT_no_encryption,
+      is_contained(encryptablePlatforms, config->platformInfo.target.Platform));
+
 #ifndef HAVE_LIBXAR
   if (config->emitBitcodeBundle)
     error("-bitcode_bundle unsupported because LLD wasn't built with libxar");

lld/MachO/Options.td

@@ -1158,8 +1158,10 @@ def dyld_env : Flag<["-"], "dyld_env">,
     Flags<[HelpHidden]>,
     Group<grp_undocumented>;
 def encryptable : Flag<["-"], "encryptable">,
-    HelpText<"This option is undocumented in ld64">,
-    Flags<[HelpHidden]>,
+    HelpText<"Generate the LC_ENCRYPTION_INFO load command">,
+    Group<grp_undocumented>;
+def no_encryption : Flag<["-"], "no_encryption">,
+    HelpText<"Do not generate the LC_ENCRYPTION_INFO load command">,
     Group<grp_undocumented>;
 def executable_path : Flag<["-"], "executable_path">,
     HelpText<"This option is undocumented in ld64">,
@@ -1253,10 +1255,6 @@ def no_dtrace_dof : Flag<["-"], "no_dtrace_dof">,
     HelpText<"This option is undocumented in ld64">,
     Flags<[HelpHidden]>,
     Group<grp_undocumented>;
-def no_encryption : Flag<["-"], "no_encryption">,
-    HelpText<"This option is undocumented in ld64">,
-    Flags<[HelpHidden]>,
-    Group<grp_undocumented>;
 def no_new_main : Flag<["-"], "no_new_main">,
     HelpText<"This option is undocumented in ld64">,
     Flags<[HelpHidden]>,

lld/MachO/SyntheticSections.cpp

@@ -85,7 +85,13 @@ template <class LP> MachHeaderSection *macho::makeMachHeaderSection() {
 }
 
 template <class LP> uint64_t MachHeaderSectionImpl<LP>::getSize() const {
-  return sizeof(typename LP::mach_header) + sizeOfCmds + config->headerPad;
+  uint64_t size =
+      sizeof(typename LP::mach_header) + sizeOfCmds + config->headerPad;
+  // If we are emitting an encryptable binary, our load commands must have a
+  // separate (non-encrypted) page to themselves.
+  if (config->emitEncryptionInfo)
+    size = alignTo(size, target->getPageSize());
+  return size;
 }
 
 static uint32_t cpuSubtype() {

lld/MachO/Target.h

@@ -86,9 +86,12 @@ struct LP64 {
   using nlist = structs::nlist_64;
   using segment_command = llvm::MachO::segment_command_64;
   using section = llvm::MachO::section_64;
+  using encryption_info_command = llvm::MachO::encryption_info_command_64;
 
   static constexpr uint32_t magic = llvm::MachO::MH_MAGIC_64;
   static constexpr uint32_t segmentLCType = llvm::MachO::LC_SEGMENT_64;
+  static constexpr uint32_t encryptionInfoLCType =
+      llvm::MachO::LC_ENCRYPTION_INFO_64;
 
   static constexpr uint64_t pageZeroSize = 1ull << 32;
   static constexpr size_t wordSize = 8;
@@ -99,9 +102,12 @@ struct ILP32 {
   using nlist = structs::nlist;
   using segment_command = llvm::MachO::segment_command;
   using section = llvm::MachO::section;
+  using encryption_info_command = llvm::MachO::encryption_info_command;
 
   static constexpr uint32_t magic = llvm::MachO::MH_MAGIC;
   static constexpr uint32_t segmentLCType = llvm::MachO::LC_SEGMENT;
+  static constexpr uint32_t encryptionInfoLCType =
+      llvm::MachO::LC_ENCRYPTION_INFO;
 
   static constexpr uint64_t pageZeroSize = 1ull << 12;
   static constexpr size_t wordSize = 4;

lld/MachO/Writer.cpp

@@ -471,6 +471,29 @@ public:
   mutable uint8_t *uuidBuf;
 };
 
+template <class LP> class LCEncryptionInfo : public LoadCommand {
+public:
+  LCEncryptionInfo() {}
+
+  uint32_t getSize() const override {
+    return sizeof(typename LP::encryption_info_command);
+  }
+
+  void writeTo(uint8_t *buf) const override {
+    using EncryptionInfo = typename LP::encryption_info_command;
+    auto *c = reinterpret_cast<EncryptionInfo *>(buf);
+    buf += sizeof(EncryptionInfo);
+    c->cmd = LP::encryptionInfoLCType;
+    c->cmdsize = getSize();
+    c->cryptoff = in.header->getSize();
+    auto it = find_if(outputSegments, [](const OutputSegment *seg) {
+      return seg->name == segment_names::text;
+    });
+    assert(it != outputSegments.end());
+    c->cryptsize = (*it)->fileSize - c->cryptoff;
+  }
+};
+
 class LCCodeSignature : public LoadCommand {
 public:
   LCCodeSignature(CodeSignatureSection *section) : section(section) {}
@@ -621,6 +644,8 @@ template <class LP> void Writer::createLoadCommands() {
       make<LCDysymtab>(symtabSection, indirectSymtabSection));
   if (functionStartsSection)
     in.header->addLoadCommand(make<LCFunctionStarts>(functionStartsSection));
+  if (config->emitEncryptionInfo)
+    in.header->addLoadCommand(make<LCEncryptionInfo<LP>>());
   for (StringRef path : config->runtimePaths)
     in.header->addLoadCommand(make<LCRPath>(path));
 

lld/test/MachO/encryption-info.s

@@ -0,0 +1,35 @@
+# REQUIRES: x86
+# RUN: rm -rf %t; mkdir -p %t
+# RUN: llvm-mc -filetype=obj -triple=x86_64-apple-darwin %s -o %t/test.o
+# RUN: llvm-mc -filetype=obj -triple=arm64_32-apple-watchos %s -o %t/watchos-test.o
+
+# RUN: %lld -lSystem -o %t/test %t/test.o
+# RUN: llvm-objdump --macho --all-headers %t/test | FileCheck %s --check-prefix=NO-ENCRYPTION -DSUFFIX=_64
+
+# RUN: %lld -lSystem -encryptable -o %t/test %t/test.o
+# RUN: llvm-objdump --macho --all-headers %t/test | FileCheck %s --check-prefix=ENCRYPTION -DSUFFIX=_64 -D#PAGE_SIZE=4096
+
+# RUN: %lld-watchos -lSystem -o %t/watchos-test %t/watchos-test.o
+# RUN: llvm-objdump --macho --all-headers %t/watchos-test | FileCheck %s --check-prefix=ENCRYPTION -DSUFFIX= -D#PAGE_SIZE=16384
+
+# RUN: %lld-watchos -lSystem -no_encryption -o %t/watchos-test %t/watchos-test.o
+# RUN: llvm-objdump --macho --all-headers %t/watchos-test | FileCheck %s --check-prefix=NO-ENCRYPTION -DSUFFIX=
+
+# ENCRYPTION:      segname __TEXT
+# ENCRYPTION-NEXT: vmaddr
+# ENCRYPTION-NEXT: vmsize
+# ENCRYPTION-NEXT: fileoff 0
+# ENCRYPTION-NEXT: filesize [[#TEXT_SIZE:]]
+
+# ENCRYPTION:      cmd LC_ENCRYPTION_INFO[[SUFFIX]]{{$}}
+# ENCRYPTION-NEXT: cmdsize
+# ENCRYPTION-NEXT: cryptoff [[#PAGE_SIZE]]
+# ENCRYPTION-NEXT: cryptsize [[#TEXT_SIZE - PAGE_SIZE]]
+# ENCRYPTION-NEXT: cryptid 0
+
+# NO-ENCRYPTION-NOT: LC_ENCRYPTION_INFO[[SUFFIX]]{{$}}
+
+.globl _main
+.p2align 2
+_main:
+  ret