Merge pull request #18060 from jrose-apple/and-you-get-a-class_getImageName

[runtime] Backwards-deployment support for class_getImageName

https://bugs.swift.org/browse/SR-1917
rdar://problem/41535552
diff --git a/lib/Driver/DarwinToolChains.cpp b/lib/Driver/DarwinToolChains.cpp
index b665db6..cbb63ff 100644
--- a/lib/Driver/DarwinToolChains.cpp
+++ b/lib/Driver/DarwinToolChains.cpp
@@ -208,13 +208,14 @@
   // When updating the versions listed here, please record the most recent
   // feature being depended on and when it was introduced:
   //
-  // - The hook to override class_getImageName (macOS 10.14 and equivalent)
+  // - Make assigning 'nil' to an NSMutableDictionary subscript delete the
+  //   entry, like it does for Swift.Dictionary, rather than trap.
   if (triple.isiOS())
-    return triple.isOSVersionLT(12);
+    return triple.isOSVersionLT(9);
   if (triple.isMacOSX())
-    return triple.isMacOSXVersionLT(10, 14);
+    return triple.isMacOSXVersionLT(10, 11);
   if (triple.isWatchOS())
-    return triple.isOSVersionLT(5);
+    return false;
   llvm_unreachable("unknown Darwin OS");
 }
 
diff --git a/stdlib/public/runtime/CMakeLists.txt b/stdlib/public/runtime/CMakeLists.txt
index d53c725..50b4481 100644
--- a/stdlib/public/runtime/CMakeLists.txt
+++ b/stdlib/public/runtime/CMakeLists.txt
@@ -31,6 +31,7 @@
     SwiftObject.mm
     SwiftValue.mm
     ReflectionMirror.mm
+    ObjCRuntimeGetImageNameFromClass.cpp
     "${SWIFT_SOURCE_DIR}/lib/Demangling/OldRemangler.cpp"
     "${SWIFT_SOURCE_DIR}/lib/Demangling/Remangler.cpp"
     "${SWIFT_SOURCE_DIR}/lib/Demangling/TypeDecoder.cpp"
diff --git a/stdlib/public/runtime/Metadata.cpp b/stdlib/public/runtime/Metadata.cpp
index f00eb12..05d43d8 100644
--- a/stdlib/public/runtime/Metadata.cpp
+++ b/stdlib/public/runtime/Metadata.cpp
@@ -55,8 +55,7 @@
 #endif
 
 #if SWIFT_OBJC_INTEROP
-#include <dlfcn.h>
-#include <objc/runtime.h>
+#include "ObjCRuntimeGetImageNameFromClass.h"
 #endif
 
 #include <cstdio>
@@ -1976,45 +1975,6 @@
   return self;
 }
 
-#if SWIFT_OBJC_INTEROP
-
-// FIXME: This is from a later version of <objc/runtime.h>. Once the declaration
-// is available in SDKs, we can remove this typedef.
-typedef BOOL (*objc_hook_getImageName)(
-    Class _Nonnull cls, const char * _Nullable * _Nonnull outImageName);
-
-/// \see customGetImageNameFromClass
-static objc_hook_getImageName defaultGetImageNameFromClass = nullptr;
-
-/// A custom implementation of Objective-C's class_getImageName for Swift
-/// classes, which knows how to handle dynamically-initialized class metadata.
-///
-/// Per the documentation for objc_setHook_getImageName, any non-Swift classes
-/// will still go through the normal implementation of class_getImageName,
-/// which is stored in defaultGetImageNameFromClass.
-static BOOL
-customGetImageNameFromClass(Class _Nonnull objcClass,
-                            const char * _Nullable * _Nonnull outImageName) {
-  auto *classAsMetadata = reinterpret_cast<const ClassMetadata *>(objcClass);
-
-  // Is this a Swift class?
-  if (classAsMetadata->isTypeMetadata() &&
-      !classAsMetadata->isArtificialSubclass()) {
-    const void *descriptor = classAsMetadata->getDescription();
-    assert(descriptor &&
-           "all non-artificial Swift classes should have a descriptor");
-    Dl_info imageInfo = {};
-    if (!dladdr(descriptor, &imageInfo))
-      return NO;
-    *outImageName = imageInfo.dli_fname;
-    return imageInfo.dli_fname != nullptr;
-  }
-
-  // If not, fall back to the default implementation.
-  return defaultGetImageNameFromClass(objcClass, outImageName);
-}
-#endif
-
 /// Initialize the field offset vector for a dependent-layout class, using the
 /// "Universal" layout strategy.
 void
@@ -2028,15 +1988,7 @@
   static swift_once_t onceToken;
   swift_once(&onceToken, [](void *unused) {
     (void)unused;
-    // FIXME: This is from a later version of <objc/runtime.h>. Once the
-    // declaration is available in SDKs, we can access this directly instead of
-    // using dlsym.
-    if (void *setHookPtr = dlsym(RTLD_DEFAULT, "objc_setHook_getImageName")) {
-      auto setHook = reinterpret_cast<
-          void(*)(objc_hook_getImageName _Nonnull,
-                  objc_hook_getImageName _Nullable * _Nonnull)>(setHookPtr);
-      setHook(customGetImageNameFromClass, &defaultGetImageNameFromClass);
-    }
+    setUpObjCRuntimeGetImageNameFromClass();
   }, nullptr);
 #endif
 
diff --git a/stdlib/public/runtime/ObjCRuntimeGetImageNameFromClass.cpp b/stdlib/public/runtime/ObjCRuntimeGetImageNameFromClass.cpp
new file mode 100644
index 0000000..c25beaf
--- /dev/null
+++ b/stdlib/public/runtime/ObjCRuntimeGetImageNameFromClass.cpp
@@ -0,0 +1,297 @@
+//===--- ObjCRuntimeGetImageNameFromClass.cpp - ObjC hook setup -----------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2018 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+// Setup for the Objective-C runtime function class_getImageName, making it
+// understand Swift classes. This is tricky because before Apple's 2018 OSs,
+// this function was not designed to be hooked.
+//
+//===----------------------------------------------------------------------===//
+
+#include "ObjCRuntimeGetImageNameFromClass.h"
+#include "swift/Runtime/Config.h"
+
+#if SWIFT_OBJC_INTEROP
+
+#include "swift/Runtime/Metadata.h"
+
+#include <dlfcn.h>
+#include <objc/runtime.h>
+
+// Note: There are more #includes below under "Function patching machinery".
+// Those are only relevant to the function patching machinery.
+
+using namespace swift;
+
+
+// FIXME: This is from a later version of <objc/runtime.h>. Once the declaration
+// is available in SDKs, we can remove this typedef.
+typedef BOOL (*objc_hook_getImageName)(
+    Class _Nonnull cls, const char * _Nullable * _Nonnull outImageName);
+
+/// \see customGetImageNameFromClass
+static objc_hook_getImageName defaultGetImageNameFromClass = nullptr;
+
+/// A custom implementation of Objective-C's class_getImageName for Swift
+/// classes, which knows how to handle dynamically-initialized class metadata.
+///
+/// Per the documentation for objc_setHook_getImageName, any non-Swift classes
+/// will still go through the normal implementation of class_getImageName,
+/// which is stored in defaultGetImageNameFromClass.
+static BOOL
+getImageNameFromSwiftClass(Class _Nonnull objcClass,
+                           const char * _Nullable * _Nonnull outImageName) {
+  auto *classAsMetadata = reinterpret_cast<const ClassMetadata *>(objcClass);
+
+  // Is this a Swift class?
+  if (classAsMetadata->isTypeMetadata() &&
+      !classAsMetadata->isArtificialSubclass()) {
+    const void *descriptor = classAsMetadata->getDescription();
+    assert(descriptor &&
+           "all non-artificial Swift classes should have a descriptor");
+    Dl_info imageInfo = {};
+    if (!dladdr(descriptor, &imageInfo))
+      return NO;
+    *outImageName = imageInfo.dli_fname;
+    return imageInfo.dli_fname != nullptr;
+  }
+
+  // If not, fall back to the default implementation.
+  return defaultGetImageNameFromClass(objcClass, outImageName);
+}
+
+/***************************************************************************/
+/* Function patching machinery *********************************************/
+/***************************************************************************/
+
+#include "llvm/ADT/ArrayRef.h"
+
+#include <mach-o/dyld.h>
+#include <mach-o/loader.h>
+#include <mach-o/nlist.h>
+
+#include <cstring>
+
+using llvm::ArrayRef;
+
+namespace {
+
+#if __LP64__
+# define LC_SEGMENT_COMMAND         LC_SEGMENT_64
+# define LC_ROUTINES_COMMAND        LC_ROUTINES_64
+  typedef struct mach_header_64     macho_header;
+  typedef struct section_64         macho_section;
+  typedef struct nlist_64           macho_nlist;
+  typedef struct segment_command_64 macho_segment_command;
+#else
+# define LC_SEGMENT_COMMAND         LC_SEGMENT
+# define LC_ROUTINES_COMMAND        LC_ROUTINES
+  typedef struct mach_header        macho_header;
+  typedef struct section            macho_section;
+  typedef struct nlist              macho_nlist;
+  typedef struct segment_command    macho_segment_command;
+#endif
+
+  struct patch_t {
+    const char *name;
+    const void *fn;
+
+    template<typename T>
+    patch_t(const char *newName, const T *newFn)
+      : name(newName), fn((const void*)newFn) { 
+    }
+  };
+} // end anonymous namespace
+
+/// Overwrite a cross-image symbol reference by directly editing symbol tables
+/// in a Mach-O image.
+///
+/// This technique only works for certain versions of Apple's dynamic linker;
+/// fortunately we only even attempt to invoke it when running on the OSs where
+/// it works. Newer OSs already have the hook we need; older ones don't support
+/// Swift at all.
+///
+/// Also, if the symbol being patched has references within the image where it
+/// was originaly defined, those references will \e not be patched.
+static void patchLazyPointers(const mach_header *mh, patch_t patch) {
+  // Get linkEditBase
+  const uint32_t cmd_count = mh->ncmds;
+  const load_command * const cmds =
+    (const load_command *)((const char *)mh + sizeof(macho_header));
+  const load_command *cmd;
+
+  const uint8_t *linkEditBase = nullptr;
+  intptr_t slide = 0;
+
+  cmd = cmds;
+  for (uint32_t i = 0; i < cmd_count; ++i) {
+    if (cmd->cmd == LC_SEGMENT_COMMAND) {
+      const macho_segment_command *seg = (const macho_segment_command *)cmd;
+      if (strcmp(seg->segname, "__TEXT") == 0) 
+        slide = (uintptr_t)mh - seg->vmaddr;
+      else if (strcmp(seg->segname,"__LINKEDIT") == 0) 
+        linkEditBase = (const uint8_t *)(seg->vmaddr + slide - seg->fileoff);
+    }
+    cmd = (const load_command *)(((const char *)cmd)+cmd->cmdsize);
+  }
+  if (linkEditBase == nullptr)
+    return;
+
+  // Gather symbol table info
+  const macho_nlist *symbolTable = nullptr;
+  const char *stringTable = nullptr;
+  uint32_t stringTableBytes = 0;
+  const uint32_t *indirectSymbolTable = nullptr;
+
+  cmd = cmds;
+  for (uint32_t i = 0; i < cmd_count; ++i) {
+    switch (cmd->cmd) {
+    case LC_SYMTAB: {
+      const symtab_command *symtab = (const symtab_command *)cmd;
+      stringTable = (const char *)&linkEditBase[symtab->stroff];
+      stringTableBytes = symtab->strsize;
+      symbolTable = (const macho_nlist *)(&linkEditBase[symtab->symoff]);
+      break;
+    }
+    case LC_DYSYMTAB: {
+      const dysymtab_command *dsymtab = (const dysymtab_command *)cmd;
+      indirectSymbolTable =
+          (const uint32_t *)(&linkEditBase[dsymtab->indirectsymoff]);
+      break;
+    }
+    default:
+      break;
+    }
+    cmd = (const load_command *)(((const char *)cmd)+cmd->cmdsize);
+  }
+  if (symbolTable == nullptr || stringTable == nullptr ||
+      indirectSymbolTable == nullptr) {
+    return;
+  }
+
+  // Find lazy pointer section
+  cmd = cmds;
+  for (uint32_t i = 0; i < cmd_count; ++i) {
+    if (cmd->cmd == LC_SEGMENT_COMMAND) {
+      const macho_segment_command *seg = (const macho_segment_command *)cmd;
+      const macho_section * const sectionsStart = 
+          (const macho_section *)(seg + 1);
+      ArrayRef<macho_section> sections(sectionsStart, seg->nsects);
+
+      for (const macho_section &sect : sections) {
+        const uint8_t type = sect.flags & SECTION_TYPE;
+        if (type != S_LAZY_SYMBOL_POINTERS)
+          continue;
+        
+        const size_t pointerCount = sect.size / sizeof(uintptr_t);
+        uintptr_t * const symbolPointers = (uintptr_t *)(sect.addr + slide);
+        const uint32_t indirectTableOffset = sect.reserved1;
+        for (uint32_t lazyIndex = 0; lazyIndex < pointerCount; ++lazyIndex) {
+          uint32_t symbolIndex =
+              indirectSymbolTable[indirectTableOffset + lazyIndex];
+          if (symbolIndex >= stringTableBytes) {
+            // Presumably INDIRECT_SYMBOL_LOCAL or some other special value.
+            continue;
+          }
+
+          // Found symbol for this lazy pointer, now lookup address.
+          const char *lazyTargetName = 
+              &stringTable[symbolTable[symbolIndex].n_un.n_strx];
+          if (strcmp(patch.name, lazyTargetName) == 0) {
+            // Can't use the value currently stored here because it may 
+            // be a dyld stub binder that will undo our patch if called.
+            symbolPointers[lazyIndex] = (uintptr_t)patch.fn;
+          }
+        }
+      }
+    }
+    cmd = (const load_command *)(((const char *)cmd)+cmd->cmdsize);
+  }
+}
+
+/// \see callUnpatchedGetImageNameFromClass
+static decltype(&class_getImageName) unpatchedGetImageNameFromClass = nullptr;
+
+/// A fallback implementation of class_getImageName that just calls
+/// unpatchedGetImageNameFromClass, with the signature of
+/// objc_hook_getImageName.
+///
+/// This is used on older OSs where objc_setHook_getImageName isn't supported.
+/// In this case, we invoke the Swift implementation above
+/// (customGetImageNameFromClass), but set it up to fall back to this one.
+/// Then this one can call the system's original version, which should be stored
+/// in unpatchedGetImageNameFromClass.
+static BOOL callUnpatchedGetImageNameFromClass(
+    Class _Nonnull objcClass, const char * _Nullable * _Nonnull outImageName) {
+  *outImageName = unpatchedGetImageNameFromClass(objcClass);
+  return outImageName != nullptr;
+}
+
+/// A patched version of class_getImageName that always uses the Swift 
+/// implementation.
+///
+/// The Swift implementation is always set up to chain to another
+/// implementation, so on older OSs we just have to make sure that that chained
+/// implementation is the original system version. See
+/// callUnpatchedGetImageNameFromClass.
+static const char *patchedGetImageNameFromClassForOldOSs(Class _Nullable cls) {
+  if (!cls)
+    return nullptr;
+  const char *result;
+  if (getImageNameFromSwiftClass(cls, &result))
+    return result;
+  return nullptr;
+}
+
+/// A hook for _dyld_register_func_for_add_image that overwrites any references
+/// to class_getImageName with our custom implementation.
+static void patchGetImageNameInImage(const struct mach_header *mh,
+                                     intptr_t vmaddr_slide) {
+  (void)vmaddr_slide;
+  patchLazyPointers(mh, patch_t("_class_getImageName",
+                                &patchedGetImageNameFromClassForOldOSs));
+}
+
+/***************************************************************************/
+/* Installing the hook *****************************************************/
+/***************************************************************************/
+
+void swift::setUpObjCRuntimeGetImageNameFromClass() {
+  assert(defaultGetImageNameFromClass == nullptr && "already set up");
+
+  // FIXME: This is from a later version of <objc/runtime.h>. Once the
+  // declaration is available in SDKs, we can access this directly instead of
+  // using dlsym.
+  if (void *setHookPtr = dlsym(RTLD_DEFAULT, "objc_setHook_getImageName")) {
+    auto setHook = reinterpret_cast<
+        void(*)(objc_hook_getImageName _Nonnull,
+                objc_hook_getImageName _Nullable * _Nonnull)>(setHookPtr);
+    setHook(getImageNameFromSwiftClass, &defaultGetImageNameFromClass);
+
+  } else {
+    // On older OSs, manually patch in our new implementation of
+    // class_getImageName, and set it up to chain to the original system 
+    // version.
+
+    // This assignment happens through a volatile pointer to make sure it occurs
+    // before the later call to _dyld_register_func_for_add_image. (More
+    // specifically, we need the original implementation of
+    // 'class_getImageName', not the replaced one.)
+    assert(unpatchedGetImageNameFromClass == nullptr);
+    volatile auto *originalImplementationPtr = &unpatchedGetImageNameFromClass;
+    *originalImplementationPtr = &class_getImageName;
+    defaultGetImageNameFromClass = callUnpatchedGetImageNameFromClass;
+
+    _dyld_register_func_for_add_image(&patchGetImageNameInImage);
+  }
+}
+
+#endif // SWIFT_OBJC_INTEROP
diff --git a/stdlib/public/runtime/ObjCRuntimeGetImageNameFromClass.h b/stdlib/public/runtime/ObjCRuntimeGetImageNameFromClass.h
new file mode 100644
index 0000000..96a1690
--- /dev/null
+++ b/stdlib/public/runtime/ObjCRuntimeGetImageNameFromClass.h
@@ -0,0 +1,29 @@
+//===--- ObjCRuntimeGetImageNameFromClass.h - ObjC hook setup ---*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2018 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_RUNTIME_OBJCRUNTIMEGETIMAGENAMEFROMCLASS_H
+#define SWIFT_RUNTIME_OBJCRUNTIMEGETIMAGENAMEFROMCLASS_H
+
+#include "swift/Runtime/Config.h"
+
+namespace swift {
+
+#if SWIFT_OBJC_INTEROP
+/// Set up class_getImageName so that it will understand Swift classes.
+///
+/// This function should only be called once per process.
+void setUpObjCRuntimeGetImageNameFromClass();
+#endif // SWIFT_OBJC_INTEROP
+
+} // end namespace swift
+
+#endif // SWIFT_RUNTIME_OBJCRUNTIMEGETIMAGENAMEFROMCLASS_H
diff --git a/test/Driver/linker-arclite.swift b/test/Driver/linker-arclite.swift
index 73e67aa..c320f9b 100644
--- a/test/Driver/linker-arclite.swift
+++ b/test/Driver/linker-arclite.swift
@@ -16,14 +16,12 @@
 // IOS_ARCLITE: -o {{[^ ]+}}
 
 
-// RUN: %swiftc_driver -driver-print-jobs -target x86_64-apple-macosx10.14 %S/../Inputs/empty.swift | %FileCheck -check-prefix NO_ARCLITE %s
-// RUN: %swiftc_driver -driver-print-jobs -target x86_64-apple-macosx10.13 %S/../Inputs/empty.swift | %FileCheck -check-prefix ANY_ARCLITE %s
-// RUN: %swiftc_driver -driver-print-jobs -target x86_64-apple-ios12 %S/../Inputs/empty.swift | %FileCheck -check-prefix NO_ARCLITE %s
-// RUN: %swiftc_driver -driver-print-jobs -target x86_64-apple-ios11 %S/../Inputs/empty.swift | %FileCheck -check-prefix ANY_ARCLITE %s
-// RUN: %swiftc_driver -driver-print-jobs -target arm64-apple-tvos12 %S/../Inputs/empty.swift | %FileCheck -check-prefix NO_ARCLITE %s
-// RUN: %swiftc_driver -driver-print-jobs -target arm64-apple-tvos11 %S/../Inputs/empty.swift | %FileCheck -check-prefix ANY_ARCLITE %s
-// RUN: %swiftc_driver -driver-print-jobs -target armv7k-apple-watchos5 %S/../Inputs/empty.swift | %FileCheck -check-prefix NO_ARCLITE %s
-// RUN: %swiftc_driver -driver-print-jobs -target armv7k-apple-watchos4 %S/../Inputs/empty.swift | %FileCheck -check-prefix ANY_ARCLITE %s
+// RUN: %swiftc_driver -driver-print-jobs -target x86_64-apple-macosx10.11 %S/../Inputs/empty.swift | %FileCheck -check-prefix NO_ARCLITE %s
+// RUN: %swiftc_driver -driver-print-jobs -target x86_64-apple-macosx10.10 %S/../Inputs/empty.swift | %FileCheck -check-prefix ANY_ARCLITE %s
+// RUN: %swiftc_driver -driver-print-jobs -target x86_64-apple-ios9 %S/../Inputs/empty.swift | %FileCheck -check-prefix NO_ARCLITE %s
+// RUN: %swiftc_driver -driver-print-jobs -target x86_64-apple-ios8 %S/../Inputs/empty.swift | %FileCheck -check-prefix ANY_ARCLITE %s
+// RUN: %swiftc_driver -driver-print-jobs -target arm64-apple-tvos9 %S/../Inputs/empty.swift | %FileCheck -check-prefix NO_ARCLITE %s
+// RUN: %swiftc_driver -driver-print-jobs -target armv7k-apple-watchos2 %S/../Inputs/empty.swift | %FileCheck -check-prefix NO_ARCLITE %s
 
 // NO_ARCLITE: bin/ld{{"? }}
 // NO_ARCLITE-NOT: arclite
diff --git a/test/Interpreter/SDK/class_getImageName.swift b/test/Interpreter/SDK/class_getImageName.swift
index c684d9d..6ab3af0 100644
--- a/test/Interpreter/SDK/class_getImageName.swift
+++ b/test/Interpreter/SDK/class_getImageName.swift
@@ -22,9 +22,18 @@
              "wrong library for \(cls)")
 }
 
-let isMissingObjCRuntimeHook =
-    (nil == dlsym(UnsafeMutableRawPointer(bitPattern: -2),
-                  "objc_setHook_getImageName"))
+var cannotUseObjCRuntimeHook = false
+#if targetEnvironment(device) && !os(macOS)
+if #available(iOS 12, tvOS 12, watchOS 5, *) {
+  // The only place these tests will fail is on early versions of the 2018 OSs.
+  // The final versions will have 'objc_setHook_getImageName'; anything earlier
+  // will be handled by manually overwriting the original implementation of
+  // 'class_getImageName'.
+  cannotUseObjCRuntimeHook =
+      (nil == dlsym(UnsafeMutableRawPointer(bitPattern: -2),
+                    "objc_setHook_getImageName"))
+}
+#endif // targetEnvironment(device) && !os(macOS)
 
 var testSuite = TestSuite("class_getImageName")
 
@@ -34,7 +43,7 @@
 }
 
 testSuite.test("Generic")
-    .xfail(.custom({ isMissingObjCRuntimeHook },
+    .xfail(.custom({ cannotUseObjCRuntimeHook },
                    reason: "hook for class_getImageName not present"))
     .code {
   check(GenericSwiftObject<Int>.self, in: "libGetImageNameHelper.dylib")
@@ -45,7 +54,7 @@
 }
 
 testSuite.test("GenericAncestry")
-    .xfail(.custom({ isMissingObjCRuntimeHook },
+    .xfail(.custom({ cannotUseObjCRuntimeHook },
                    reason: "hook for class_getImageName not present"))
     .code {
   check(GenericAncestrySwiftObject.self, in: "libGetImageNameHelper.dylib")