[Runtime] Fix a race with dlopen and libobjc.

stdlib/public/runtime/ImageInspectionMachO.cpp

@@ -23,6 +23,7 @@
 #include "ImageInspection.h"
 #include <mach-o/dyld.h>
 #include <mach-o/getsect.h>
+#include <objc/runtime.h>
 #include <assert.h>
 #include <dlfcn.h>
 
@@ -55,7 +56,7 @@ extern "C" void *_NSGetMachExecuteHeader();
 
 template <const char *SEGMENT_NAME, const char *SECTION_NAME,
          void CONSUME_BLOCK(const void *start, uintptr_t size)>
-void addImageCallback(const mach_header *mh, intptr_t vmaddr_slide) {
+void addImageCallback(const mach_header *mh) {
 #if __POINTER_WIDTH__ == 64
   assert(mh->magic == MH_MAGIC_64 && "loaded non-64-bit image?!");
 #endif
@@ -72,12 +73,17 @@ void addImageCallback(const mach_header *mh, intptr_t vmaddr_slide) {
 
   CONSUME_BLOCK(section, size);
 }
+template <const char *SEGMENT_NAME, const char *SECTION_NAME,
+         void CONSUME_BLOCK(const void *start, uintptr_t size)>
+void addImageCallback(const mach_header *mh, intptr_t vmaddr_slide) {
+  addImageCallback<SEGMENT_NAME, SECTION_NAME, CONSUME_BLOCK>(mh);
+}
 
 template <const char *SEGMENT_NAME, const char *SECTION_NAME,
           const char *SEGMENT_NAME2, const char *SECTION_NAME2,
           void CONSUME_BLOCK(const void *start, uintptr_t size,
                              const void *start2, uintptr_t size2)>
-void addImageCallback2Sections(const mach_header *mh, intptr_t vmaddr_slide) {
+void addImageCallback2Sections(const mach_header *mh) {
 #if __POINTER_WIDTH__ == 64
   assert(mh->magic == MH_MAGIC_64 && "loaded non-64-bit image?!");
 #endif
@@ -103,27 +109,48 @@ void addImageCallback2Sections(const mach_header *mh, intptr_t vmaddr_slide) {
 
   CONSUME_BLOCK(section, size, section2, size2);
 }
+template <const char *SEGMENT_NAME, const char *SECTION_NAME,
+          const char *SEGMENT_NAME2, const char *SECTION_NAME2,
+          void CONSUME_BLOCK(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,
+                            CONSUME_BLOCK>(mh);
+}
+
 } // end anonymous namespace
 
+#if OBJC_ADDLOADIMAGEFUNC_DEFINED
+#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__);                      \
+  }
+#else
+#define REGISTER_FUNC(...) _dyld_register_func_for_add_image(__VA_ARGS__)
+#endif
+
 void swift::initializeProtocolLookup() {
-  _dyld_register_func_for_add_image(
+  REGISTER_FUNC(
     addImageCallback<TextSegment, ProtocolsSection,
                      addImageProtocolsBlockCallback>);
 }
 
 void swift::initializeProtocolConformanceLookup() {
-  _dyld_register_func_for_add_image(
+  REGISTER_FUNC(
     addImageCallback<TextSegment, ProtocolConformancesSection,
                      addImageProtocolConformanceBlockCallback>);
 }
 void swift::initializeTypeMetadataRecordLookup() {
-  _dyld_register_func_for_add_image(
+  REGISTER_FUNC(
     addImageCallback<TextSegment, TypeMetadataRecordSection,
                      addImageTypeMetadataRecordBlockCallback>);
 }
 
 void swift::initializeDynamicReplacementLookup() {
-  _dyld_register_func_for_add_image(
+  REGISTER_FUNC(
       addImageCallback2Sections<TextSegment, DynamicReplacementSection,
                                 TextSegment, DynamicReplacementSomeSection,
                                 addImageDynamicReplacementBlockCallback>);

test/stdlib/Inputs/dlopen_race_dylib.swift

@@ -0,0 +1,6 @@
+
+class C: CustomStringConvertible {
+  var description: String {
+    return "description"
+  }
+}

test/stdlib/dlopen_race.swift

@@ -0,0 +1,56 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift -emit-library -o %t/dlopen_race.dylib %S/Inputs/dlopen_race_dylib.swift
+// RUN: %target-build-swift -o %t/dlopen_race %s
+// RUN: %target-run %t/dlopen_race
+// REQUIRES: executable_test
+// REQUIRES: objc_interop
+
+import StdlibUnittest
+
+import Darwin
+
+var DlopenRaceTests = TestSuite("DlopenRace")
+
+typealias add_image_callback = @convention(c) (UnsafeRawPointer?, Int) -> Void
+func register_func_for_add_image(_ f: add_image_callback) {
+  let RTLD_DEFAULT = UnsafeMutableRawPointer(bitPattern: -2)
+  let _dyld_register_func_for_add_image_ptr =
+    dlsym(RTLD_DEFAULT, "_dyld_register_func_for_add_image")
+  expectNotNil(_dyld_register_func_for_add_image_ptr)
+
+  typealias _dyld_register_func_for_add_image_func =
+    @convention(c) (add_image_callback) -> Void
+  let _dyld_register_func_for_add_image = unsafeBitCast(
+    _dyld_register_func_for_add_image_ptr,
+    to: _dyld_register_func_for_add_image_func.self)
+  _dyld_register_func_for_add_image(f)
+}
+
+// Make sure Swift doesn't register newly opened images before ObjC is ready
+// for them to be used. rdar://problem/49742015
+var add_image_count = 0
+DlopenRaceTests.test("race") {
+  // This test is expected to fail unless the ObjC notification is supported.
+  let RTLD_DEFAULT = UnsafeMutableRawPointer(bitPattern: -2)
+  let objc_addLoadImageFunc = dlsym(RTLD_DEFAULT, "objc_addLoadImageFunc");
+  if objc_addLoadImageFunc == nil { return }
+
+  register_func_for_add_image({ header, slide in
+    // The protocol conformance check in the print call is enough to trigger
+    // ObjC class initialization in the newly opened image if Swift has
+    // registered the conformance records in that image. While we would be
+    // unlikely to make this sort of call directly in the callback in a real
+    // program, it could happen at this time in another thread.
+    print(header, slide)
+    add_image_count += 1
+  })
+
+  let dylibPath = CommandLine.arguments[0] + ".dylib"
+
+  let beforeCount = add_image_count
+  let handle = dlopen(dylibPath, RTLD_LAZY)
+  expectNotNil(handle, String(cString: dlerror()))
+  expectEqual(add_image_count, beforeCount + 1)
+}
+
+runAllTests()