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

[PRESIDENT810]

Summary

This fixes https://github.com/llvm/llvm-project/issues/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.

[int3]

Thanks for working on this!

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.

I'm not aware of anyone asking for this right now, but it does seem like something we'll want eventually :) feel free to work on that in a followup if you're interested

lld/MachO/Arch/ARM.cpp

177–178

let's include a comment about how this only implements -no_dtrace_dof, and that this deviates from ld64's default behavior

184

I don't think we should make direct references to ld64 code like this

204–206

code convention is to omit braces for one-liners

214

nit: can you set your editor to include newlines at EOF?

lld/MachO/Arch/ARM64Common.cpp

117–121

we should handle this generically in validateRelocationInfo (if we aren't doing that already) rather than here

lld/MachO/Dtrace.cpp

1–2 ↗

(On Diff #442020)

nit: fix formatting

17 ↗

(On Diff #442020)

18 ↗

(On Diff #442020)

could use startswith()

24–40 ↗

(On Diff #442020)

not sure this function is necessary. Couldn't we hoist the error() into TargetInfo::handleDtraceReloc and call that directly?

lld/test/MachO/arm-dtrace.ll

35–61 ↗

(On Diff #442020)

can we minimize the test case? probably most of the instructions aside from the calls to dtrace functions can be removed. The !srcloc references too

[BertalanD]

lld/MachO/Arch/X86_64.cpp

222

lld/MachO/Dtrace.cpp

17–22 ↗

(On Diff #442020)

It looks like this function might be quite hot. Do you think it would make sense for it to be an inline function in the header file?

lld/test/MachO/arm-dtrace.ll

1 ↗

(On Diff #442020)

llvm-as can only be run for this test case if the ARM backend is enabled.

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

1 ↗

(On Diff #442020)

lld/test/MachO/arm64-dtrace.ll

1 ↗

(On Diff #442020)

lld/test/MachO/x86_64-dtrace.ll

1 ↗

(On Diff #442020)

[BertalanD]

lld/MachO/Arch/X86_64.cpp

222

or even loc[-1] would work, I guess

[int3]

Actually, is there any reason the test cases need to be in bitcode rather than assembly? In general we should prefer assembly if we are not testing things that require LTO

[int3]

lld/MachO/Dtrace.cpp

17–22 ↗

(On Diff #442020)

Agreed. In fact we could take it one step further though and eliminate the function entirely... the call to this can just be replaced by name.startswith("___dtrace_").

[PRESIDENT810]

Remove some unnecessary code. Add comment for handleDtraceReloc function. Replace IR tests with assembly tests.

[oontvoo]

lld/MachO/Arch/ARM.cpp

44

Let's not copy the reloc.

177

179

why do we need this assert if the default in the following switch already handles the checking?

187

these magic numbers are a bit unfortunate ... i wonder if it's better to use constants?

209

consider using toString(Symbol*) instead. (That way, we get the demangled vs non-demangled handling for the symbol name, depending on whether -demangle is set)

lld/MachO/Arch/ARM64Common.h

33

lld/MachO/Arch/X86_64.cpp

42

[PRESIDENT810]

Pass reloc by reference instead of by value. Remove the unnecessary assertion. Use toString(Symbol *) instead of getName().

I'm not sure about the magic number part. ld64 just uses these numbers directly to overwrite the original instructions, and they are not actually "magic numbers", but rather assembly instructions corresponding to different architectures, I think? Also, the comments above could give illustrations about what assembly instructions these numbers represent. So I'm not very certain if it's necessary to use constants instead.

[int3]

Just some nits but otherwise lgtm :)

lld/MachO/Arch/ARM.cpp

182–183

could we hoist this to the top of the function? seems like we'd just never want to modify any dtrace refs if we are outputting an object file

lld/MachO/Arch/X86_64.cpp

210–213

I kinda preferred the previous version where you used write32le here. I think @BertalanD's suggestion of writing to loc[-1] was to make it clear that we are only modifying a single byte. But if we're emitting a 4-byte instruction, write32le() makes sense

i.e. I'm suggesting

loc[-1] = 0x90;
write32le(loc, 0x00401F0F);

lld/test/MachO/arm-dtrace.s

5

are the -demangle -dynamic flags actually necessary? Also, %lld by default specifies a platform version of 11.0. Any reason we can't use that default?

30–31

are these branching instructions necessary? doesn't seem like they really test anything

lld/test/MachO/x86_64-dtrace.s

10–16

nit: can we column-align the right hand side?

[PRESIDENT810]

Update according to int3's advice.

[int3]

lgtm, thanks! Do you have commit permissions or do you need me to land this for you?

[PRESIDENT810]

In D129062#3636389, @int3 wrote:

lgtm, thanks! Do you have commit permissions or do you need me to land this for you?

I actually don't commit access so please help me land this. Please use "Kaining Zhong <[e-mail]>" if any author info is needed. Thanks!

lld/MachO/Arch/ARM.cpp

struct ARM : TargetInfo {
   void writeStub(uint8_t *buf, const Symbol &) const override;
   void writeStubHelperHeader(uint8_t *buf) const override;
   void writeStubHelperEntry(uint8_t *buf, const Symbol &,
                             uint64_t entryAddr) const override;
 
   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
 
 const RelocAttrs &ARM::getRelocAttrs(uint8_t type) const {
   static const std::array<RelocAttrs, 10> relocAttrsArray{{
 #define B(x) RelocAttrBits::x
       {"VANILLA", /* FIXME populate this */ B(_0)},
ARM::ARM(uint32_t cpuSubtype) : TargetInfo(ILP32()) {
   stubHelperHeaderSize = 0 /* FIXME */;
   stubHelperEntrySize = 0 /* FIXME */;
 }
 
 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 {
+  switch (r.type) {
+  case ARM_RELOC_BR24:
+    if (sym->getName().startswith("___dtrace_probe")) {
+      if (config->outputType == MH_OBJECT)
+        return;
+      // change call site to a NOP
+      write32le(loc, 0xE1A00000);
+    } else if (sym->getName().startswith("___dtrace_isenabled")) {
+      if (config->outputType == MH_OBJECT)
+        return;
+      // change call site to 'eor r0, r0, r0'
+      write32le(loc, 0xE0200000);
+    } else {
+      error("Unrecognized dtrace symbol prefix: " + sym->getName());
+    }
+    break;
+  case ARM_THUMB_RELOC_BR22:
+    if (sym->getName().startswith("___dtrace_probe")) {
+      if (config->outputType == MH_OBJECT)
+        return;
+      // change 32-bit blx call site to two thumb NOPs
+      write32le(loc, 0x46C046C0);
+    } else if (sym->getName().startswith("___dtrace_isenabled")) {
+      if (config->outputType == MH_OBJECT)
+        return;
+      // 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.h

struct ARM64Common : TargetInfo {
 
   int64_t getEmbeddedAddend(MemoryBufferRef, uint64_t offset,
                             const llvm::MachO::relocation_info) const override;
   void relocateOne(uint8_t *loc, const Reloc &, uint64_t va,
                    uint64_t pc) const override;
 
   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) {
   return ((value >> right) & ((1 << width) - 1)) << left;
 }
 
 //              25                                                0
 // +-----------+---------------------------------------------------+

lld/MachO/Arch/ARM64Common.cpp

if ((instruction & 0xbfc00000) != 0xb9400000)
     error(getRelocAttrs(type).name + " reloc requires LDR instruction");
   assert(((instruction >> 10) & 0xfff) == 0 &&
          "non-zero embedded LDR immediate");
   // C6.2.4 ADD (immediate)
   // ADD <Xd|SP>, <Xn|SP>, #<imm>{, <shift>}
   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 (sym->getName().startswith("___dtrace_probe")) {
+    if (config->outputType == MH_OBJECT)
+      return;
+    // change call site to a NOP
+    write32le(loc, 0xD503201F);
+  } else if (sym->getName().startswith("___dtrace_isenabled")) {
+    if (config->outputType == MH_OBJECT)
+      return;
+    // change call site to 'MOVZ X0,0'
+    write32le(loc, 0xD2800000);
+  } else {
+    error("Unrecognized dtrace symbol prefix: " + toString(*sym));
+  }
+}

lld/MachO/Arch/X86_64.cpp

struct X86_64 : TargetInfo {
   void writeStub(uint8_t *buf, const Symbol &) const override;
   void writeStubHelperHeader(uint8_t *buf) const override;
   void writeStubHelperEntry(uint8_t *buf, const Symbol &,
                             uint64_t entryAddr) const override;
 
   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
 
 const RelocAttrs &X86_64::getRelocAttrs(uint8_t type) const {
   static const std::array<RelocAttrs, 10> relocAttrsArray{{
 #define B(x) RelocAttrBits::x
       {"UNSIGNED",
X86_64::X86_64() : TargetInfo(LP64()) {
   stubHelperHeaderSize = sizeof(stubHelperHeader);
   stubHelperEntrySize = sizeof(stubHelperEntry);
 }
 
 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 (sym->getName().startswith("___dtrace_probe")) {
+    if (config->outputType == MH_OBJECT)
+      return;
+    // change call site to a NOP
+    loc[-1] = 0x90; // 1-byte nop
+    loc[0] = 0x0F;  // 4-byte nop
+    loc[1] = 0x1F;
+    loc[2] = 0x40;
+    loc[3] = 0x00;
+  } else if (sym->getName().startswith("___dtrace_isenabled")) {
+    if (config->outputType == MH_OBJECT)
+      return;
+    // change call site to a clear eax
+    loc[-1] = 0x33; // xorl eax,eax
+    loc[0] = 0xC0;
+    loc[1] = 0x90; // 1-byte nop
+    loc[2] = 0x90; // 1-byte nop
+    loc[3] = 0x90; // 1-byte nop
+  } else {
+    error("Unrecognized dtrace symbol prefix: " + toString(*sym));
+  }
+}

lld/MachO/InputSection.cpp

if (target->hasAttr(r.type, RelocAttrBits::SUBTRAHEND)) {
         assert(!::shouldOmitFromOutput(referentIsec));
         minuendVA = referentIsec->getVA(minuend.addend);
       }
       referentVA = minuendVA - fromSym->getVA();
     } else if (auto *referentSym = r.referent.dyn_cast<Symbol *>()) {
       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())) {
         // References from thread-local variable sections are treated as offsets
         // relative to the start of the thread-local data memory area, which
         // is initialized via copying all the TLV data sections (which are all
         // contiguous).
         if (isa<Defined>(referentSym))

lld/MachO/SymbolTable.cpp

if (name.consume_front("section$start$"))
     return handleSectionBoundarySymbol(sym, name, Boundary::Start);
   if (name.consume_front("section$end$"))
     return handleSectionBoundarySymbol(sym, name, Boundary::End);
   if (name.consume_front("segment$start$"))
     return handleSegmentBoundarySymbol(sym, name, Boundary::Start);
   if (name.consume_front("segment$end$"))
     return handleSegmentBoundarySymbol(sym, name, Boundary::End);
 
+  // Leave dtrace symbols, since we will handle them when we do the relocation
+  if (name.startswith("___dtrace_"))
+    return;
+
   // Handle -U.
   if (config->explicitDynamicLookups.count(sym.getName())) {
     symtab->addDynamicLookup(sym.getName());
     return;
   }
 
   // Handle -undefined.
   auto message = [source, &sym]() {

lld/MachO/Target.h

public:
   }
 
   bool hasAttr(uint8_t type, RelocAttrBits bit) const {
     return getRelocAttrs(type).hasAttr(bit);
   }
 
   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");
+  }
+
   uint32_t magic;
   llvm::MachO::CPUType cpuType;
   uint32_t cpuSubtype;
 
   uint64_t pageZeroSize;
   size_t headerSize;
   size_t stubSize;
   size_t stubHelperHeaderSize;

lld/test/MachO/arm-dtrace.s

+# REQUIRES: arm
+# RUN: rm -rf %t; split-file %s %t
+
+# RUN: llvm-mc -filetype=obj -triple=armv4t-apple-macosx12.0.0 %t/armv4t-dtrace.s -o %t/armv4t-dtrace.o
+# RUN: %lld -demangle -dynamic -arch armv4t -platform_version macos 12.0.0 12.3 -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-macosx12.0.0 %t/armv7-dtrace.s -o %t/armv7-dtrace.o
+# RUN: %lld -demangle -dynamic -arch armv7 -platform_version macos 12.0.0 12.3 -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
+  beq LBB0_4
+  b LBB0_1
+LBB0_1:
+  b LBB0_2
+LBB0_2:
+  .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
+  b LBB0_3
+LBB0_3:
+  b LBB0_4
+LBB0_4:
+  bx  lr
+
+.subsections_via_symbols
+
+;--- armv7-dtrace.s
+  .globl  _main
+  .thumb_func _main
+_main:
+  bl  ___dtrace_isenabled$Foo$added$v1
+  cbz r0, LBB0_4
+  b LBB0_2
+LBB0_2:
+  .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
+@ %bb.3:
+  b LBB0_4
+LBB0_4:
+
+.subsections_via_symbols

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

+# REQUIRES: aarch64
+# RUN: rm -rf %t; split-file %s %t
+
+# RUN: llvm-mc -filetype=obj -triple=arm64_32-apple-macosx12.0.0 %t/arm64-32-dtrace.s -o %t/arm64-32-dtrace.o
+# RUN: %lld -demangle -dynamic -arch arm64_32 -platform_version macos 12.0.0 12.3 -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
+  cbz w8, LBB0_4
+  b LBB0_1
+LBB0_1:
+  b LBB0_2
+LBB0_2:
+  .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
+  b LBB0_3
+LBB0_3:
+  b LBB0_4
+LBB0_4:
+  ret
+
+.subsections_via_symbols

lld/test/MachO/arm64-dtrace.s

+# REQUIRES: aarch64
+# RUN: rm -rf %t; split-file %s %t
+
+# RUN: llvm-mc -filetype=obj -triple=arm64-apple-macosx12.0.0 %t/arm64-dtrace.s -o %t/arm64-dtrace.o
+# RUN: %lld -demangle -dynamic -arch arm64 -platform_version macos 12.0.0 12.3 -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
+  cbz w8, LBB0_4
+  b LBB0_1
+LBB0_1:
+  b LBB0_2
+LBB0_2:
+  .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
+  b LBB0_3
+LBB0_3:
+  b LBB0_4
+LBB0_4:
+  ret
+
+.subsections_via_symbols

lld/test/MachO/x86_64-dtrace.s

+# 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 -demangle -dynamic -arch x86_64 -platform_version macos 12.0.0 12.3 -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
+  je  LBB0_4
+  jmp LBB0_2
+LBB0_2:
+  .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
+  jmp LBB0_4
+LBB0_4:
+  retq
+
+.subsections_via_symbols