[Runtime] Use dyld SPI for getting macho sections when available.

When available, use objc_addLoadImageFunc2 and _dyld_register_func_for_add_image to look up mach-o sections instead of using getsectiondata. When not available, we fall back to objc_addLoadImageFunc or _dyld_register_func_for_add_image and getsectiondata as before.

rdar://51760462

Author: mikeash

stdlib/public/runtime/ImageInspectionMachO.cpp

@@ -30,6 +30,31 @@
 #include <assert.h>
 #include <dlfcn.h>
 
+#if __has_include(<objc/objc-internal.h>) && __has_include(<mach-o/dyld_priv.h>)
+#include <mach-o/dyld_priv.h>
+#include <objc/objc-internal.h>
+#else
+
+// Bring our own definition of enum _dyld_section_location_kind and some of its
+// values. They'll be unused when we can't include dyld_priv.h, but the code
+// needs them in order to compile.
+enum _dyld_section_location_kind {
+  _dyld_section_location_text_swift5_protos,
+  _dyld_section_location_text_swift5_proto,
+  _dyld_section_location_text_swift5_types,
+  _dyld_section_location_text_swift5_replace,
+  _dyld_section_location_text_swift5_replace2,
+  _dyld_section_location_text_swift5_ac_funcs,
+};
+
+#endif
+
+#if !OBJC_ADDLOADIMAGEFUNC2_DEFINED
+// If we don't have objc_addLoadImageFunc2, fall back to objc_addLoadImageFunc.
+// Use a #define so we don't have to conditionalize the calling code below.
+#define objc_addLoadImageFunc2 objc_addLoadImageFunc
+#endif
+
 using namespace swift;
 
 namespace {
@@ -53,35 +78,62 @@ using mach_header_platform = mach_header_64;
 using mach_header_platform = mach_header;
 #endif
 
+// Callback for objc_addLoadImageFunc that just takes a mach_header.
 template <const char *SEGMENT_NAME, const char *SECTION_NAME,
-         void CONSUME_BLOCK(const void *baseAddress,
-                            const void *start, uintptr_t size)>
+          enum _dyld_section_location_kind SECTION_KIND,
+          void CONSUME_BLOCK(const void *baseAddress, const void *start,
+                             uintptr_t size)>
 void addImageCallback(const mach_header *mh) {
 #if __POINTER_WIDTH__ == 64
   assert(mh->magic == MH_MAGIC_64 && "loaded non-64-bit image?!");
 #endif
-  
-  // Look for a __swift5_proto section.
+
   unsigned long size;
   const uint8_t *section =
   getsectiondata(reinterpret_cast<const mach_header_platform *>(mh),
                  SEGMENT_NAME, SECTION_NAME,
                  &size);
-  
+
   if (!section)
     return;
-  
+
   CONSUME_BLOCK(mh, section, size);
 }
+
+// Callback for objc_addLoadImageFunc2 that takes a mach_header and dyld info.
+#if OBJC_ADDLOADIMAGEFUNC2_DEFINED
+template <const char *SEGMENT_NAME, const char *SECTION_NAME,
+          enum _dyld_section_location_kind SECTION_KIND,
+          void CONSUME_BLOCK(const void *baseAddress, const void *start,
+                             uintptr_t size)>
+void addImageCallback(const mach_header *mh,
+                      struct _dyld_section_location_info_s *dyldInfo) {
+#if __POINTER_WIDTH__ == 64
+  assert(mh->magic == MH_MAGIC_64 && "loaded non-64-bit image?!");
+#endif
+
+  auto result = _dyld_lookup_section_info(mh, dyldInfo, SECTION_KIND);
+  if (result.buffer)
+    CONSUME_BLOCK(mh, result.buffer, result.bufferSize);
+}
+#endif
+
+// Callback for _dyld_register_func_for_add_image that takes a mach_header and a
+// slide.
 template <const char *SEGMENT_NAME, const char *SECTION_NAME,
-         void CONSUME_BLOCK(const void *baseAddress,
-                            const void *start, uintptr_t size)>
+          enum _dyld_section_location_kind SECTION_KIND,
+          void CONSUME_BLOCK(const void *baseAddress, const void *start,
+                             uintptr_t size)>
 void addImageCallback(const mach_header *mh, intptr_t vmaddr_slide) {
-  addImageCallback<SEGMENT_NAME, SECTION_NAME, CONSUME_BLOCK>(mh);
+  addImageCallback<SEGMENT_NAME, SECTION_NAME, SECTION_KIND, CONSUME_BLOCK>(mh);
 }
 
+// Callback for objc_addLoadImageFunc that just takes a mach_header.
+
 template <const char *SEGMENT_NAME, const char *SECTION_NAME,
           const char *SEGMENT_NAME2, const char *SECTION_NAME2,
+          enum _dyld_section_location_kind SECTION_KIND,
+          enum _dyld_section_location_kind SECTION_KIND2,
           void CONSUME_BLOCK(const void *baseAddress,
                              const void *start, uintptr_t size,
                              const void *start2, uintptr_t size2)>
@@ -111,25 +163,63 @@ void addImageCallback2Sections(const mach_header *mh) {
 
   CONSUME_BLOCK(mh, section, size, section2, size2);
 }
+
+// Callback for objc_addLoadImageFunc2 that takes a mach_header and dyld info.
+#if OBJC_ADDLOADIMAGEFUNC2_DEFINED
 template <const char *SEGMENT_NAME, const char *SECTION_NAME,
           const char *SEGMENT_NAME2, const char *SECTION_NAME2,
+          enum _dyld_section_location_kind SECTION_KIND,
+          enum _dyld_section_location_kind SECTION_KIND2,
           void CONSUME_BLOCK(const void *baseAddress,
                              const void *start, uintptr_t size,
                              const void *start2, uintptr_t size2)>
+void addImageCallback2Sections(const mach_header *mh,
+                               struct _dyld_section_location_info_s *dyldInfo) {
+#if __POINTER_WIDTH__ == 64
+  assert(mh->magic == MH_MAGIC_64 && "loaded non-64-bit image?!");
+#endif
+
+  // Look for a section.
+  auto result = _dyld_lookup_section_info(mh, dyldInfo, SECTION_KIND);
+  if (!result.buffer)
+    return;
+
+  auto result2 = _dyld_lookup_section_info(mh, dyldInfo, SECTION_KIND2);
+  // No NULL check here, we allow this one not to be present. dyld gives us
+  // a NULL pointer and 0 size when the section isn't in the dylib so we don't
+  // need to zero anything out.
+
+  CONSUME_BLOCK(mh, result.buffer, result.bufferSize, result2.buffer,
+                result2.bufferSize);
+}
+#endif
+
+// Callback for _dyld_register_func_for_add_image that takes a mach_header and a
+// slide.
+template <const char *SEGMENT_NAME, const char *SECTION_NAME,
+          const char *SEGMENT_NAME2, const char *SECTION_NAME2,
+          enum _dyld_section_location_kind SECTION_KIND,
+          enum _dyld_section_location_kind SECTION_KIND2,
+          void CONSUME_BLOCK(const void *baseAddress, const void *start,
+                             uintptr_t size, const void *start2,
+                             uintptr_t size2)>
 void addImageCallback2Sections(const mach_header *mh, intptr_t vmaddr_slide) {
-  addImageCallback2Sections<SEGMENT_NAME, SECTION_NAME,
-                            SEGMENT_NAME2, SECTION_NAME2,
+  addImageCallback2Sections<SEGMENT_NAME, SECTION_NAME, SEGMENT_NAME2,
+                            SECTION_NAME2, SECTION_KIND, SECTION_KIND2,
                             CONSUME_BLOCK>(mh);
 }
 
 } // end anonymous namespace
 
 #if OBJC_ADDLOADIMAGEFUNC_DEFINED && SWIFT_OBJC_INTEROP
-#define REGISTER_FUNC(...)                                               \
-  if (__builtin_available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)) { \
-    objc_addLoadImageFunc(__VA_ARGS__);                                  \
-  } else {                                                               \
-    _dyld_register_func_for_add_image(__VA_ARGS__);                      \
+#define REGISTER_FUNC(...)                                                     \
+  if (SWIFT_RUNTIME_WEAK_CHECK(objc_addLoadImageFunc2)) {                      \
+    SWIFT_RUNTIME_WEAK_USE(objc_addLoadImageFunc2(__VA_ARGS__));               \
+  } else if (__builtin_available(macOS 10.15, iOS 13, tvOS 13,                 \
+                                 watchOS 6, *)) {                              \
+    objc_addLoadImageFunc(__VA_ARGS__);                                        \
+  } else {                                                                     \
+    _dyld_register_func_for_add_image(__VA_ARGS__);                            \
   }
 #else
 #define REGISTER_FUNC(...) _dyld_register_func_for_add_image(__VA_ARGS__)
@@ -144,30 +234,36 @@ void addImageCallback2Sections(const mach_header *mh, intptr_t vmaddr_slide) {
 
 void swift::initializeProtocolLookup() {
   REGISTER_FUNC(addImageCallback<TextSegment, ProtocolsSection,
+                                 _dyld_section_location_text_swift5_protos,
                                  addImageProtocolsBlockCallbackUnsafe>);
 }
 
 void swift::initializeProtocolConformanceLookup() {
   REGISTER_FUNC(
       addImageCallback<TextSegment, ProtocolConformancesSection,
+                       _dyld_section_location_text_swift5_proto,
                        addImageProtocolConformanceBlockCallbackUnsafe>);
 }
 void swift::initializeTypeMetadataRecordLookup() {
   REGISTER_FUNC(
       addImageCallback<TextSegment, TypeMetadataRecordSection,
+                       _dyld_section_location_text_swift5_types,
                        addImageTypeMetadataRecordBlockCallbackUnsafe>);
 }
 
 void swift::initializeDynamicReplacementLookup() {
   REGISTER_FUNC(
       addImageCallback2Sections<TextSegment, DynamicReplacementSection,
                                 TextSegment, DynamicReplacementSomeSection,
+                                _dyld_section_location_text_swift5_replace,
+                                _dyld_section_location_text_swift5_replace2,
                                 addImageDynamicReplacementBlockCallback>);
 }
 
 void swift::initializeAccessibleFunctionsLookup() {
   REGISTER_FUNC(
       addImageCallback<TextSegment, AccessibleFunctionsSection,
+                       _dyld_section_location_text_swift5_ac_funcs,
                        addImageAccessibleFunctionsBlockCallbackUnsafe>);
 }