[lld-macho] add code signature for native arm64 macOS

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

lld/MachO/OutputSegment.h

@@ -45,6 +45,7 @@ class OutputSegment {
   size_t numNonHiddenSections() const;
 
   uint64_t fileOff = 0;
+  uint64_t fileSize = 0;
   StringRef name;
   uint32_t maxProt = 0;
   uint32_t initProt = 0;

lld/MachO/SyntheticSections.cpp

@@ -24,6 +24,11 @@
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/LEB128.h"
 #include "llvm/Support/Path.h"
+#include "llvm/Support/SHA256.h"
+
+#if defined(__APPLE__)
+#include <sys/mman.h>
+#endif
 
 using namespace llvm;
 using namespace llvm::support;
@@ -887,3 +892,93 @@ void StringTableSection::writeTo(uint8_t *buf) const {
     off += str.size() + 1; // account for null terminator
   }
 }
+
+CodeSignatureSection::CodeSignatureSection()
+    : LinkEditSection(segment_names::linkEdit, section_names::codeSignature) {
+  align = 16; // required by libstuff
+  fileName = config->outputFile;
+  size_t slashIndex = fileName.rfind("/");
+  if (slashIndex != std::string::npos)
+    fileName = fileName.drop_front(slashIndex + 1);
+  allHeadersSize = alignTo<16>(fixedHeadersSize + fileName.size() + 1);
+  fileNamePad = allHeadersSize - fixedHeadersSize - fileName.size();
+}
+
+uint32_t CodeSignatureSection::getBlockCount() const {
+  return (fileOff + blockSize - 1) / blockSize;
+}
+
+uint64_t CodeSignatureSection::getRawSize() const {
+  return allHeadersSize + getBlockCount() * hashSize;
+}
+
+void CodeSignatureSection::writeHashes(uint8_t *buf) const {
+  uint8_t *code = buf;
+  uint8_t *codeEnd = buf + fileOff;
+  uint8_t *hashes = codeEnd + allHeadersSize;
+  while (code < codeEnd) {
+    StringRef block(reinterpret_cast<char *>(code),
+                    std::min(codeEnd - code, static_cast<ssize_t>(blockSize)));
+    SHA256 hasher;
+    hasher.update(block);
+    StringRef hash = hasher.final();
+    assert(hash.size() == hashSize);
+    memcpy(hashes, hash.data(), hashSize);
+    code += blockSize;
+    hashes += hashSize;
+  }
+#if defined(__APPLE__)
+  // This is macOS-specific work-around and makes no sense for any
+  // other host OS. See https://openradar.appspot.com/FB8914231
+  //
+  // The macOS kernel maintains a signature-verification cache to
+  // quickly validate applications at time of execve(2).  The trouble
+  // is that for the kernel creates the cache entry at the time of the
+  // mmap(2) call, before we have a chance to write either the code to
+  // sign or the signature header+hashes.  The fix is to invalidate
+  // all cached data associated with the output file, thus discarding
+  // the bogus prematurely-cached signature.
+  msync(buf, fileOff + getSize(), MS_INVALIDATE);
+#endif
+}
+
+void CodeSignatureSection::writeTo(uint8_t *buf) const {
+  using namespace llvm::MachO;
+  uint32_t signatureSize = static_cast<uint32_t>(getSize());
+  auto *superBlob = reinterpret_cast<CS_SuperBlob *>(buf);
+  write32be(&superBlob->magic, CSMAGIC_EMBEDDED_SIGNATURE);
+  write32be(&superBlob->length, signatureSize);
+  write32be(&superBlob->count, 1);
+  auto *blobIndex = reinterpret_cast<CS_BlobIndex *>(&superBlob[1]);
+  write32be(&blobIndex->type, CSSLOT_CODEDIRECTORY);
+  write32be(&blobIndex->offset, blobHeadersSize);
+  auto *codeDirectory =
+      reinterpret_cast<CS_CodeDirectory *>(buf + blobHeadersSize);
+  write32be(&codeDirectory->magic, CSMAGIC_CODEDIRECTORY);
+  write32be(&codeDirectory->length, signatureSize - blobHeadersSize);
+  write32be(&codeDirectory->version, CS_SUPPORTSEXECSEG);
+  write32be(&codeDirectory->flags, CS_ADHOC | CS_LINKER_SIGNED);
+  write32be(&codeDirectory->hashOffset,
+            sizeof(CS_CodeDirectory) + fileName.size() + fileNamePad);
+  write32be(&codeDirectory->identOffset, sizeof(CS_CodeDirectory));
+  codeDirectory->nSpecialSlots = 0;
+  write32be(&codeDirectory->nCodeSlots, getBlockCount());
+  write32be(&codeDirectory->codeLimit, fileOff);
+  codeDirectory->hashSize = static_cast<uint8_t>(hashSize);
+  codeDirectory->hashType = kSecCodeSignatureHashSHA256;
+  codeDirectory->platform = 0;
+  codeDirectory->pageSize = blockSizeShift;
+  codeDirectory->spare2 = 0;
+  codeDirectory->scatterOffset = 0;
+  codeDirectory->teamOffset = 0;
+  codeDirectory->spare3 = 0;
+  codeDirectory->codeLimit64 = 0;
+  OutputSegment *textSeg = getOrCreateOutputSegment(segment_names::text);
+  write64be(&codeDirectory->execSegBase, textSeg->fileOff);
+  write64be(&codeDirectory->execSegLimit, textSeg->fileSize);
+  write64be(&codeDirectory->execSegFlags,
+            config->outputType == MH_EXECUTE ? CS_EXECSEG_MAIN_BINARY : 0);
+  auto *id = reinterpret_cast<char *>(&codeDirectory[1]);
+  memcpy(id, fileName.begin(), fileName.size());
+  memset(id + fileName.size(), 0, fileNamePad);
+}

lld/MachO/SyntheticSections.h

@@ -18,6 +18,7 @@
 
 #include "llvm/ADT/PointerUnion.h"
 #include "llvm/ADT/SetVector.h"
+#include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace llvm {
@@ -40,6 +41,7 @@ constexpr const char export_[] = "__export";
 constexpr const char symbolTable[] = "__symbol_table";
 constexpr const char indirectSymbolTable[] = "__ind_sym_tab";
 constexpr const char stringTable[] = "__string_table";
+constexpr const char codeSignature[] = "__code_signature";
 constexpr const char got[] = "__got";
 constexpr const char threadPtrs[] = "__thread_ptrs";
 constexpr const char unwindInfo[] = "__unwind_info";
@@ -94,7 +96,7 @@ class LinkEditSection : public SyntheticSection {
   // NOTE: This assumes that the extra bytes required for alignment can be
   // zero-valued bytes.
   uint64_t getSize() const override final {
-    return llvm::alignTo(getRawSize(), WordSize);
+    return llvm::alignTo(getRawSize(), align);
   }
 };
 
@@ -482,6 +484,32 @@ class IndirectSymtabSection : public LinkEditSection {
   void writeTo(uint8_t *buf) const override;
 };
 
+// The code signature comes at the very end of the linked output file.
+class CodeSignatureSection : public LinkEditSection {
+public:
+  static constexpr uint8_t blockSizeShift = 12;
+  static constexpr size_t blockSize = (1 << blockSizeShift); // 4 KiB
+  static constexpr size_t hashSize = 256 / 8;
+  static constexpr size_t blobHeadersSize = llvm::alignTo<8>(
+      sizeof(llvm::MachO::CS_SuperBlob) + sizeof(llvm::MachO::CS_BlobIndex));
+  static constexpr uint32_t fixedHeadersSize =
+      blobHeadersSize + sizeof(llvm::MachO::CS_CodeDirectory);
+
+  uint32_t fileNamePad = 0;
+  uint32_t allHeadersSize = 0;
+  StringRef fileName;
+
+  CodeSignatureSection();
+  uint64_t getRawSize() const override;
+  bool isNeeded() const override { return true; }
+  void writeTo(uint8_t *buf) const override;
+  uint32_t getBlockCount() const;
+  void writeHashes(uint8_t *buf) const;
+};
+
+static_assert((CodeSignatureSection::blobHeadersSize % 8) == 0, "");
+static_assert((CodeSignatureSection::fixedHeadersSize % 8) == 0, "");
+
 struct InStruct {
   MachHeaderSection *header = nullptr;
   RebaseSection *rebase = nullptr;

lld/MachO/Writer.cpp

@@ -52,6 +52,7 @@ class Writer {
   void openFile();
   void writeSections();
   void writeUuid();
+  void writeCodeSignature();
 
   void run();
 
@@ -62,6 +63,7 @@ class Writer {
   StringTableSection *stringTableSection = nullptr;
   SymtabSection *symtabSection = nullptr;
   IndirectSymtabSection *indirectSymtabSection = nullptr;
+  CodeSignatureSection *codeSignatureSection = nullptr;
   UnwindInfoSection *unwindInfoSection = nullptr;
   LCUuid *uuidCommand = nullptr;
 };
@@ -400,6 +402,23 @@ class LCUuid : public LoadCommand {
   mutable uint8_t *uuidBuf;
 };
 
+class LCCodeSignature : public LoadCommand {
+public:
+  LCCodeSignature(CodeSignatureSection *section) : section(section) {}
+
+  uint32_t getSize() const override { return sizeof(linkedit_data_command); }
+
+  void writeTo(uint8_t *buf) const override {
+    auto *c = reinterpret_cast<linkedit_data_command *>(buf);
+    c->cmd = LC_CODE_SIGNATURE;
+    c->cmdsize = getSize();
+    c->dataoff = static_cast<uint32_t>(section->fileOff);
+    c->datasize = section->getSize();
+  }
+
+  CodeSignatureSection *section;
+};
+
 } // namespace
 
 static void prepareSymbolRelocation(lld::macho::Symbol *sym,
@@ -521,6 +540,9 @@ void Writer::createLoadCommands() {
     }
   }
 
+  if (codeSignatureSection)
+    in.header->addLoadCommand(make<LCCodeSignature>(codeSignatureSection));
+
   const uint32_t MACOS_MAXPATHLEN = 1024;
   config->headerPad = std::max(
       config->headerPad, (config->headerPadMaxInstallNames
@@ -624,6 +646,7 @@ static int sectionOrder(OutputSection *osec) {
         .Case(section_names::symbolTable, -3)
         .Case(section_names::indirectSymbolTable, -2)
         .Case(section_names::stringTable, -1)
+        .Case(section_names::codeSignature, std::numeric_limits<int>::max())
         .Default(0);
   }
   // ZeroFill sections must always be the at the end of their segments,
@@ -678,6 +701,9 @@ void Writer::createOutputSections() {
   unwindInfoSection = make<UnwindInfoSection>(); // TODO(gkm): only when no -r
   symtabSection = make<SymtabSection>(*stringTableSection);
   indirectSymtabSection = make<IndirectSymtabSection>();
+  if (config->outputType == MH_EXECUTE &&
+      (config->arch == AK_arm64 || config->arch == AK_arm64e))
+    codeSignatureSection = make<CodeSignatureSection>();
 
   switch (config->outputType) {
   case MH_EXECUTE:
@@ -743,6 +769,7 @@ void Writer::assignAddresses(OutputSegment *seg) {
     addr += osec->getSize();
     fileOff += osec->getFileSize();
   }
+  seg->fileSize = fileOff - seg->fileOff;
 }
 
 void Writer::openFile() {
@@ -770,6 +797,11 @@ void Writer::writeUuid() {
   uuidCommand->writeUuid(digest);
 }
 
+void Writer::writeCodeSignature() {
+  if (codeSignatureSection)
+    codeSignatureSection->writeHashes(buffer->getBufferStart());
+}
+
 void Writer::run() {
   // dyld requires __LINKEDIT segment to always exist (even if empty).
   OutputSegment *linkEditSegment =
@@ -817,6 +849,7 @@ void Writer::run() {
 
   writeSections();
   writeUuid();
+  writeCodeSignature();
 
   if (auto e = buffer->commit())
     error("failed to write to the output file: " + toString(std::move(e)));