[lld-macho] Handle user-provided dtrace symbols to avoid linking failure

This fixes #56238. ld64.lld currently does not generate __dof section in Mach-O, and -no_dtrace_dof option is on by default. However when there are user-defined dtrace symbols, ld64.lld will treat them as undefined symbols, which causes the linking to fail because lld cannot find their definitions. This patch allows ld64.lld to rewrite the instructions calling dtrace symbols to instructions like nop as what ld64 does; therefore, when encountered with user-provided dtrace probes, the linking can still succeed.

I'm not sure whether support for dtrace is expected in lld, so for now I didn't add codes to make lld emit __dof section like ld64, and only made it possible to link with dtrace symbols provided. If this feature is needed, I can add that part in Dtrace.cpp & Dtrace.h.

Reviewed By: int3, #lld-macho

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

lld/MachO/Arch/ARM.cpp

@@ -40,6 +40,9 @@ struct ARM : TargetInfo {
   void relaxGotLoad(uint8_t *loc, uint8_t type) const override;
   const RelocAttrs &getRelocAttrs(uint8_t type) const override;
   uint64_t getPageSize() const override { return 4 * 1024; }
+
+  void handleDtraceReloc(const Symbol *sym, const Reloc &r,
+                         uint8_t *loc) const override;
 };
 
 } // namespace
@@ -170,3 +173,36 @@ TargetInfo *macho::createARMTargetInfo(uint32_t cpuSubtype) {
   static ARM t(cpuSubtype);
   return &t;
 }
+
+void ARM::handleDtraceReloc(const Symbol *sym, const Reloc &r,
+                            uint8_t *loc) const {
+  if (config->outputType == MH_OBJECT)
+    return;
+
+  switch (r.type) {
+  case ARM_RELOC_BR24:
+    if (sym->getName().startswith("___dtrace_probe")) {
+      // change call site to a NOP
+      write32le(loc, 0xE1A00000);
+    } else if (sym->getName().startswith("___dtrace_isenabled")) {
+      // change call site to 'eor r0, r0, r0'
+      write32le(loc, 0xE0200000);
+    } else {
+      error("Unrecognized dtrace symbol prefix: " + toString(*sym));
+    }
+    break;
+  case ARM_THUMB_RELOC_BR22:
+    if (sym->getName().startswith("___dtrace_probe")) {
+      // change 32-bit blx call site to two thumb NOPs
+      write32le(loc, 0x46C046C0);
+    } else if (sym->getName().startswith("___dtrace_isenabled")) {
+      // change 32-bit blx call site to 'nop', 'eor r0, r0'
+      write32le(loc, 0x46C04040);
+    } else {
+      error("Unrecognized dtrace symbol prefix: " + toString(*sym));
+    }
+    break;
+  default:
+    llvm_unreachable("Unsupported dtrace relocation type for ARM");
+  }
+}

lld/MachO/Arch/ARM64Common.cpp

@@ -109,3 +109,21 @@ void ARM64Common::relaxGotLoad(uint8_t *loc, uint8_t type) const {
   instruction = ((instruction & 0x001fffff) | 0x91000000);
   write32le(loc, instruction);
 }
+
+void ARM64Common::handleDtraceReloc(const Symbol *sym, const Reloc &r,
+                                    uint8_t *loc) const {
+  assert(r.type == ARM64_RELOC_BRANCH26);
+
+  if (config->outputType == MH_OBJECT)
+    return;
+
+  if (sym->getName().startswith("___dtrace_probe")) {
+    // change call site to a NOP
+    write32le(loc, 0xD503201F);
+  } else if (sym->getName().startswith("___dtrace_isenabled")) {
+    // change call site to 'MOVZ X0,0'
+    write32le(loc, 0xD2800000);
+  } else {
+    error("Unrecognized dtrace symbol prefix: " + toString(*sym));
+  }
+}

lld/MachO/Arch/ARM64Common.h

@@ -29,6 +29,9 @@ struct ARM64Common : TargetInfo {
 
   void relaxGotLoad(uint8_t *loc, uint8_t type) const override;
   uint64_t getPageSize() const override { return 16 * 1024; }
+
+  void handleDtraceReloc(const Symbol *sym, const Reloc &r,
+                         uint8_t *loc) const override;
 };
 
 inline uint64_t bitField(uint64_t value, int right, int width, int left) {

lld/MachO/Arch/X86_64.cpp

@@ -39,6 +39,9 @@ struct X86_64 : TargetInfo {
   void relaxGotLoad(uint8_t *loc, uint8_t type) const override;
   const RelocAttrs &getRelocAttrs(uint8_t type) const override;
   uint64_t getPageSize() const override { return 4 * 1024; }
+
+  void handleDtraceReloc(const Symbol *sym, const Reloc &r,
+                         uint8_t *loc) const override;
 };
 
 } // namespace
@@ -199,3 +202,23 @@ TargetInfo *macho::createX86_64TargetInfo() {
   static X86_64 t;
   return &t;
 }
+
+void X86_64::handleDtraceReloc(const Symbol *sym, const Reloc &r,
+                               uint8_t *loc) const {
+  assert(r.type == X86_64_RELOC_BRANCH);
+
+  if (config->outputType == MH_OBJECT)
+    return;
+
+  if (sym->getName().startswith("___dtrace_probe")) {
+    // change call site to a NOP
+    loc[-1] = 0x90;
+    write32le(loc, 0x00401F0F);
+  } else if (sym->getName().startswith("___dtrace_isenabled")) {
+    // change call site to a clear eax
+    loc[-1] = 0x33;
+    write32le(loc, 0x909090C0);
+  } else {
+    error("Unrecognized dtrace symbol prefix: " + toString(*sym));
+  }
+}

lld/MachO/InputSection.cpp

@@ -201,6 +201,12 @@ void ConcatInputSection::writeTo(uint8_t *buf) {
       if (target->hasAttr(r.type, RelocAttrBits::LOAD) &&
           !referentSym->isInGot())
         target->relaxGotLoad(loc, r.type);
+      // For dtrace symbols, do not handle them as normal undefined symbols
+      if (referentSym->getName().startswith("___dtrace_")) {
+        // Change dtrace call site to pre-defined instructions
+        target->handleDtraceReloc(referentSym, r, loc);
+        continue;
+      }
       referentVA = resolveSymbolVA(referentSym, r.type) + r.addend;
 
       if (isThreadLocalVariables(getFlags())) {

lld/MachO/SymbolTable.cpp

@@ -332,6 +332,10 @@ static bool recoverFromUndefinedSymbol(const Undefined &sym) {
     return true;
   }
 
+  // Leave dtrace symbols, since we will handle them when we do the relocation
+  if (name.startswith("___dtrace_"))
+    return true;
+
   // Handle -U.
   if (config->explicitDynamicLookups.count(sym.getName())) {
     symtab->addDynamicLookup(sym.getName());

lld/MachO/Target.h

@@ -79,6 +79,15 @@ class TargetInfo {
 
   bool usesThunks() const { return thunkSize > 0; }
 
+  // For now, handleDtraceReloc only implements -no_dtrace_dof, and ensures
+  // that the linking would not fail even when there are user-provided dtrace
+  // symbols. However, unlike ld64, lld currently does not emit __dof sections.
+  virtual void handleDtraceReloc(const Symbol *sym, const Reloc &r,
+                                 uint8_t *loc) const {
+    llvm_unreachable("Unsupported architecture for dtrace symbols");
+  }
+
+
   virtual void applyOptimizationHints(uint8_t *buf, const ConcatInputSection *,
                                       llvm::ArrayRef<uint64_t>) const {};
 

lld/test/MachO/arm-dtrace.s

@@ -0,0 +1,45 @@
+# REQUIRES: arm
+# RUN: rm -rf %t; split-file %s %t
+
+# RUN: llvm-mc -filetype=obj -triple=armv4t-apple-darwin %t/armv4t-dtrace.s -o %t/armv4t-dtrace.o
+# RUN: %lld -arch armv4t -o %t/armv4t-dtrace %t/armv4t-dtrace.o
+
+## If references of dtrace symbols are handled by lld, their relocation should be replaced with the following instructions
+# RUN: llvm-objdump --macho -D %t/armv4t-dtrace | FileCheck %s --check-prefix=CHECK-armv4t
+
+# CHECK-armv4t: 00 00 20 e0  eor     r0, r0, r0
+
+# CHECK-armv4t: 00 00 a0 e1  mov     r0, r0
+
+# RUN: llvm-mc -filetype=obj -triple=thumbv7-apple-darwin %t/armv7-dtrace.s -o %t/armv7-dtrace.o
+# RUN: %lld -arch armv7 -o %t/armv7-dtrace %t/armv7-dtrace.o
+
+## If references of dtrace symbols are handled by lld, their relocation should be replaced with the following instructions
+# RUN: llvm-objdump --macho -D %t/armv7-dtrace | FileCheck %s --check-prefix=CHECK-armv7
+
+# CHECK-armv7:      40 40   eors    r0, r0
+# CHECK-armv7-NEXT: c0 46   mov     r8, r8
+
+# CHECK-armv7:      c0 46   mov     r8, r8
+# CHECK-armv7-NEXT: c0 46   mov     r8, r8
+
+;--- armv4t-dtrace.s
+	.globl	_main
+_main:
+	bl	___dtrace_isenabled$Foo$added$v1
+	.reference	___dtrace_typedefs$Foo$v2
+	bl	___dtrace_probe$Foo$added$v1$696e74
+	.reference	___dtrace_stability$Foo$v1$1_1_0_1_1_0_1_1_0_1_1_0_1_1_0
+
+.subsections_via_symbols
+
+;--- armv7-dtrace.s
+	.globl	_main
+	.thumb_func	_main
+_main:
+	bl	___dtrace_isenabled$Foo$added$v1
+	.reference	___dtrace_typedefs$Foo$v2
+	bl	___dtrace_probe$Foo$added$v1$696e74
+	.reference	___dtrace_stability$Foo$v1$1_1_0_1_1_0_1_1_0_1_1_0_1_1_0
+
+.subsections_via_symbols

lld/test/MachO/arm64-32-dtrace.s

@@ -0,0 +1,23 @@
+# REQUIRES: aarch64
+# RUN: rm -rf %t; split-file %s %t
+
+# RUN: llvm-mc -filetype=obj -triple=arm64_32-apple-darwin %t/arm64-32-dtrace.s -o %t/arm64-32-dtrace.o
+# RUN: %lld -arch arm64_32 -o %t/arm64-32-dtrace %t/arm64-32-dtrace.o
+
+## If references of dtrace symbols are handled by lld, their relocation should be replaced with the following instructions
+# RUN: llvm-objdump --macho -D %t/arm64-32-dtrace | FileCheck %s --check-prefix=CHECK
+
+# CHECK: 00 00 80 d2  mov     x0, #0
+
+# CHECK: 1f 20 03 d5  nop
+
+#--- arm64-32-dtrace.s
+	.globl	_main
+_main:
+	bl	___dtrace_isenabled$Foo$added$v1
+	.reference	___dtrace_typedefs$Foo$v2
+	bl	___dtrace_probe$Foo$added$v1$696e74
+	.reference	___dtrace_stability$Foo$v1$1_1_0_1_1_0_1_1_0_1_1_0_1_1_0
+	ret
+
+.subsections_via_symbols

lld/test/MachO/arm64-dtrace.s

@@ -0,0 +1,23 @@
+# REQUIRES: aarch64
+# RUN: rm -rf %t; split-file %s %t
+
+# RUN: llvm-mc -filetype=obj -triple=arm64-apple-darwin %t/arm64-dtrace.s -o %t/arm64-dtrace.o
+# RUN: %lld -arch arm64 -o %t/arm64-dtrace %t/arm64-dtrace.o
+
+## If references of dtrace symbols are handled by lld, their relocation should be replaced with the following instructions
+# RUN: llvm-objdump --macho -D %t/arm64-dtrace | FileCheck %s --check-prefix=CHECK
+
+# CHECK: 00 00 80 d2  mov     x0, #0
+
+# CHECK: 1f 20 03 d5  nop
+
+#--- arm64-dtrace.s
+	.globl	_main
+_main:
+	bl	___dtrace_isenabled$Foo$added$v1
+	.reference	___dtrace_typedefs$Foo$v2
+	bl	___dtrace_probe$Foo$added$v1$696e74
+	.reference	___dtrace_stability$Foo$v1$1_1_0_1_1_0_1_1_0_1_1_0_1_1_0
+	ret
+
+.subsections_via_symbols

lld/test/MachO/x86_64-dtrace.s

@@ -0,0 +1,27 @@
+# REQUIRES: x86
+# RUN: rm -rf %t; split-file %s %t
+
+# RUN: llvm-mc -filetype=obj -triple=x86_64-apple-darwin %t/x86_64-dtrace.s -o %t/x86_64-dtrace.o
+# RUN: %lld -arch x86_64 -o %t/x86_64-dtrace %t/x86_64-dtrace.o
+
+## If references of dtrace symbols are handled by lld, their relocation should be replaced with the following instructions
+# RUN: llvm-objdump --macho -D %t/x86_64-dtrace | FileCheck %s --check-prefix=CHECK
+
+# CHECK:      33 c0                   xorl    %eax, %eax
+# CHECK-NEXT: 90                      nop
+# CHECK-NEXT: 90                      nop
+# CHECK-NEXT: 90                      nop
+
+# CHECK:      90                      nop
+# CHECK-NEXT: 0f 1f 40 00             nopl    (%rax)
+
+#--- x86_64-dtrace.s
+	.globl	_main
+_main:
+	callq	___dtrace_isenabled$Foo$added$v1
+	.reference	___dtrace_typedefs$Foo$v2
+	callq	___dtrace_probe$Foo$added$v1$696e74
+	.reference	___dtrace_stability$Foo$v1$1_1_0_1_1_0_1_1_0_1_1_0_1_1_0
+	retq
+
+.subsections_via_symbols