Fix class availability watcher timing

.github/workflows/swiftlint_analyze.yml

@@ -29,6 +29,6 @@ jobs:
         run: xcodebuild -scheme InterposeKit -project InterposeKit.xcodeproj clean build-for-testing > xcodebuild.log
         shell: bash
       - name: Install SwiftLint
-        run: HOMEBREW_NO_AUTO_UPDATE=1 brew install https://raw.github.com/Homebrew/homebrew-core/master/Formula/swiftlint.rb
+        run: HOMEBREW_NO_AUTO_UPDATE=1 brew install swiftlint
       - name: Run SwiftLint Analyze
         run: swiftlint analyze --strict --compiler-log-path xcodebuild.log --reporter github-actions-logging

.github/workflows/xcodebuild.yml

@@ -7,19 +7,23 @@ on:
       - '.github/workflows/xcodebuild.yml'
       - 'InterposeKit.xcodeproj/**'
       - 'Sources/**/*.[ch]'
+      - 'Sources/**/*.m'
       - 'Sources/**/*.swift'
       - 'Example/**'
       - 'Tests/**/*.swift'
+      - 'Tests/verify_class_load_watcher.sh'
       - 'Tests/verify_release_linking.sh'
       - '!Tests/LinuxMain.swift'
   pull_request:
     paths:
       - '.github/workflows/xcodebuild.yml'
       - 'InterposeKit.xcodeproj/**'
       - 'Sources/**/*.[ch]'
+      - 'Sources/**/*.m'
       - 'Sources/**/*.swift'
       - 'Example/**'
       - 'Tests/**/*.swift'
+      - 'Tests/verify_class_load_watcher.sh'
       - 'Tests/verify_release_linking.sh'
       - '!Tests/LinuxMain.swift'
   workflow_dispatch:
@@ -38,6 +42,10 @@ jobs:
       - run: xcodebuild -version
       - name: Release device app preserves object-hook helper
         run: bash Tests/verify_release_linking.sh
+      - name: Class watcher observes newly loaded image
+        run: bash Tests/verify_class_load_watcher.sh
+      - name: Legacy class watcher waits for image loading
+        run: IKT_FORCE_DYLD_IMAGE_CALLBACK=1 bash Tests/verify_class_load_watcher.sh
       - name: macOS with UTF16
         if: always()
         run: YAMS_DEFAULT_ENCODING=UTF16 xcodebuild ${{ matrix.xcode_flags }} ${{ matrix.flags_for_test }} -destination "platform=macOS,arch=$(uname -m)" test | xcpretty

.gitignore

@@ -12,6 +12,7 @@ xcuserdata/
 ## compatibility with Xcode 3 and earlier (ignoring not required starting Xcode 4)
 build/
 DerivedData/
+*.profraw
 *.moved-aside
 *.pbxuser
 !default.pbxuser

CHANGELOG.md

@@ -13,6 +13,7 @@
 * Preserve trampoline argument registers in code-coverage builds.
 * Preserve object-hook super trampolines in optimized Release builds, fixing https://github.com/steipete/InterposeKit/issues/29. Thanks to [@Thomvis](https://github.com/Thomvis).
 * Preserve floating-point arguments when object hooks invoke original methods on arm64. Thanks to [@ishutinvv](https://github.com/ishutinvv).
+* Run class-availability hooks after Objective-C loads a new image, fixing https://github.com/steipete/InterposeKit/issues/26.
 
 ## 0.01
 

InterposeKit.podspec

@@ -6,7 +6,7 @@ Pod::Spec.new do |s|
   s.source                    = { :git => s.homepage + '.git', :tag => s.version }
   s.license                   = { :type => 'MIT', :file => 'LICENSE' }
   s.authors                   = { 'Peter Steinberger' => 'steipete@gmail.com' }
-  s.source_files              = 'Sources/**/*.{h,c,swift}'
+  s.source_files              = 'Sources/**/*.{h,m,c,swift}'
   s.swift_versions            = ['5.2']
   s.pod_target_xcconfig       = { 'APPLICATION_EXTENSION_API_ONLY' => 'YES' }
   s.ios.deployment_target     = '11.0'

InterposeKit.xcodeproj/project.pbxproj

@@ -25,7 +25,7 @@
 		78C39D912483165600B46395 /* InterposeKit.swift in Sources */ = {isa = PBXBuildFile; fileRef = 78C39D902483165600B46395 /* InterposeKit.swift */; };
 		78C39D932483169300B46395 /* InterposeKitTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 78C39D922483169300B46395 /* InterposeKitTests.swift */; };
 		78C5A4A82494D75100EE9756 /* MultipleInterposing.swift in Sources */ = {isa = PBXBuildFile; fileRef = 78C5A4A72494D75100EE9756 /* MultipleInterposing.swift */; };
-		78E20D952497B3480021552C /* ITKSuperBuilder.h in Headers */ = {isa = PBXBuildFile; fileRef = 78E20D922497B3470021552C /* ITKSuperBuilder.h */; };
+		78E20D952497B3480021552C /* ITKSuperBuilder.h in Headers */ = {isa = PBXBuildFile; fileRef = 78E20D922497B3470021552C /* ITKSuperBuilder.h */; settings = {ATTRIBUTES = (Public, ); }; };
 		78E20D962497B3480021552C /* ITKSuperBuilder.m in Sources */ = {isa = PBXBuildFile; fileRef = 78E20D942497B3470021552C /* ITKSuperBuilder.m */; };
 		78E20D9824981B2A0021552C /* InterposeSubclass.swift in Sources */ = {isa = PBXBuildFile; fileRef = 78E20D9724981B2A0021552C /* InterposeSubclass.swift */; };
 		78EDB8DA248BA9B300D2F6C1 /* TestClass.swift in Sources */ = {isa = PBXBuildFile; fileRef = 78EDB8D4248B9BB500D2F6C1 /* TestClass.swift */; };

Sources/InterposeKit/InterposeKit.h

@@ -6,6 +6,7 @@
 //
 
 #import <Foundation/Foundation.h>
+#import <InterposeKit/ITKSuperBuilder.h>
 
 //! Project version number for InterposeKit.
 FOUNDATION_EXPORT double InterposeKitVersionNumber;

Sources/InterposeKit/Watcher.swift

@@ -1,7 +1,7 @@
 import Foundation
 
-#if !os(Linux)
-import MachO.dyld
+#if SWIFT_PACKAGE && !os(Linux)
+import SuperBuilder
 #endif
 
 // MARK: Interpose Class Load Watcher
@@ -81,29 +81,34 @@ private struct InterposeWatcher {
         }
     }
 
-    // Register hook when dyld loads an image
+    // Register hook when the Objective-C runtime finishes loading an image.
     private static let globalWatcherQueue = DispatchQueue(label: "com.steipete.global-image-watcher")
     private static func installGlobalImageLoadWatcher() {
-        _dyld_register_func_for_add_image { _, _ in
-            InterposeWatcher.globalWatcherQueue.sync {
-                // this is called on the thread the image is loaded.
-                InterposeWatcher.globalWatchers = InterposeWatcher.globalWatchers.filter { waiter -> Bool in
-                    do {
-                        if try waiter.tryExecute() == false {
-                            return true // only collect if this fails because class is not there yet
-                        } else {
-                            Interpose.log("\(waiter.className) was successful.")
-                        }
-                    } catch {
-                        Interpose.log("Error while executing task: \(error).")
-                        // We can't bubble up the throw into the C context.
-                        #if DEBUG
-                        // Instead of silently eating, it's better to crash in DEBUG.
-                        fatalError("Error while executing task: \(error).")
-                        #endif
+        #if !os(Linux)
+        IKTRegisterImageDidLoadCallback {
+            InterposeWatcher.processWatchersAfterImageLoad()
+        }
+        #endif
+    }
+
+    private static func processWatchersAfterImageLoad() {
+        globalWatcherQueue.sync {
+            globalWatchers = globalWatchers.filter { waiter -> Bool in
+                do {
+                    if try waiter.tryExecute() == false {
+                        return true // only collect if this fails because class is not there yet
+                    } else {
+                        Interpose.log("\(waiter.className) was successful.")
                     }
-                    return false
+                } catch {
+                    Interpose.log("Error while executing task: \(error).")
+                    // We can't bubble up the throw into the C context.
+                    #if DEBUG
+                    // Instead of silently eating, it's better to crash in DEBUG.
+                    fatalError("Error while executing task: \(error).")
+                    #endif
                 }
+                return false
             }
         }
     }

Sources/SuperBuilder/include/ITKSuperBuilder.h

@@ -4,6 +4,11 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
+typedef void (*ITKImageDidLoadCallback)(void);
+
+/// Schedules `callback` asynchronously when an image load can be inspected safely.
+FOUNDATION_EXPORT void IKTRegisterImageDidLoadCallback(ITKImageDidLoadCallback callback);
+
 /**
 Adds an empty super implementation instance method to originalClass.
 If a method already exists, this will return NO and a descriptive error message.
@@ -63,7 +68,7 @@ There are a few important details:
 
 @end
 
-NSString *const SuperBuilderErrorDomain;
+FOUNDATION_EXPORT NSString *const SuperBuilderErrorDomain;
 
 typedef NS_ERROR_ENUM(SuperBuilderErrorDomain, SuperBuilderErrorCode) {
     SuperBuilderErrorCodeArchitectureNotSupported,

Sources/SuperBuilder/src/ITKSuperBuilder.m

@@ -1,13 +1,205 @@
 #if __APPLE__
 #import "ITKSuperBuilder.h"
 
+#import <dlfcn.h>
+#import <dispatch/dispatch.h>
+#import <mach-o/dyld.h>
+#import <mach-o/getsect.h>
+#import <os/lock.h>
+#import <stdlib.h>
+#import <string.h>
+#import <unistd.h>
+
+#if defined(IKT_TEST_DELAY_DYLD_CALLBACK)
+#import <stdatomic.h>
+#endif
+
 @import ObjectiveC.message;
 @import ObjectiveC.runtime;
 
 NS_ASSUME_NONNULL_BEGIN
 
 NSString *const SuperBuilderErrorDomain = @"com.steipete.superbuilder";
 
+static os_unfair_lock _imageDidLoadLock = OS_UNFAIR_LOCK_INIT;
+static ITKImageDidLoadCallback *_imageDidLoadCallbacks;
+static size_t _imageDidLoadCallbackCount;
+static dispatch_queue_t _imageDidLoadQueue;
+static dispatch_queue_t _legacyImageLoadQueue;
+static BOOL _imageDidLoadHookInstalled;
+static const NSUInteger IKTLegacyImageLoadRetryCount = 5000;
+#if defined(IKT_TEST_DELAY_DYLD_CALLBACK)
+static atomic_int _delayFutureDyldCallbacks;
+#endif
+
+typedef struct {
+    BOOL containsObjCClasses;
+    char *imagePath;
+} IKTLegacyImageLoadContext;
+
+static void IKTInvokeImageDidLoadCallback(void *context) {
+    os_unfair_lock_lock(&_imageDidLoadLock);
+    const size_t callbackCount = _imageDidLoadCallbackCount;
+    ITKImageDidLoadCallback *callbacks = malloc(callbackCount * sizeof(*callbacks));
+    if (callbacks != NULL) {
+        memcpy(callbacks, _imageDidLoadCallbacks, callbackCount * sizeof(*callbacks));
+    }
+    os_unfair_lock_unlock(&_imageDidLoadLock);
+
+    if (callbacks != NULL) {
+        for (size_t index = 0; index < callbackCount; index++) {
+            callbacks[index]();
+        }
+        free(callbacks);
+    }
+    if (context != NULL) {
+        dlclose(context);
+    }
+}
+
+static void IKTScheduleImageDidLoadCallback(void) {
+    dispatch_async_f(_imageDidLoadQueue, NULL, IKTInvokeImageDidLoadCallback);
+}
+
+#if !defined(IKT_FORCE_DYLD_IMAGE_CALLBACK)
+static void IKTObjCImageDidLoad(__unused const struct mach_header *header) {
+    IKTScheduleImageDidLoadCallback();
+}
+#endif
+
+static BOOL IKTImageContainsObjCClasses(const struct mach_header *header) {
+    unsigned long sectionSize = 0;
+    #if __LP64__
+    const struct mach_header_64 *header64 = (const struct mach_header_64 *)header;
+    getsectiondata(header64, "__DATA", "__objc_classlist", &sectionSize);
+    if (sectionSize == 0) {
+        getsectiondata(header64, "__DATA_CONST", "__objc_classlist", &sectionSize);
+    }
+    if (sectionSize == 0) {
+        getsectiondata(header64, "__DATA_DIRTY", "__objc_classlist", &sectionSize);
+    }
+    if (sectionSize == 0) {
+        getsectiondata(header64, "__AUTH_CONST", "__objc_classlist", &sectionSize);
+    }
+    #else
+    getsectiondata(header, "__DATA", "__objc_classlist", &sectionSize);
+    if (sectionSize == 0) {
+        getsectiondata(header, "__DATA_CONST", "__objc_classlist", &sectionSize);
+    }
+    if (sectionSize == 0) {
+        getsectiondata(header, "__DATA_DIRTY", "__objc_classlist", &sectionSize);
+    }
+    #endif
+    return sectionSize > 0;
+}
+
+static void IKTWaitForObjCImageLoad(void *context) {
+    IKTLegacyImageLoadContext *imageContext = context;
+    void *image = NULL;
+    if (imageContext->containsObjCClasses && imageContext->imagePath != NULL) {
+        // The add-image callback can run before RTLD_NOLOAD can retain the image.
+        for (NSUInteger attempt = 0; attempt < IKTLegacyImageLoadRetryCount && image == NULL; attempt++) {
+            image = dlopen(imageContext->imagePath, RTLD_LAZY | RTLD_NOLOAD);
+            if (image == NULL) {
+                usleep(1000);
+            }
+        }
+
+        unsigned int classCount = 0;
+        if (image != NULL) {
+            // Invalid weak-superclass classes may never register, so bound the wait.
+            for (NSUInteger attempt = 0; attempt < IKTLegacyImageLoadRetryCount && classCount == 0; attempt++) {
+                const char **classNames = objc_copyClassNamesForImage(imageContext->imagePath, &classCount);
+                free(classNames);
+                if (classCount == 0) {
+                    usleep(1000);
+                }
+            }
+        }
+    }
+    free(imageContext->imagePath);
+    free(imageContext);
+    dispatch_async_f(_imageDidLoadQueue, image, IKTInvokeImageDidLoadCallback);
+}
+
+static void IKTDyldImageDidLoad(const struct mach_header *header, __unused intptr_t slide) {
+    IKTLegacyImageLoadContext *context = calloc(1, sizeof(*context));
+    if (context != NULL) {
+        context->containsObjCClasses = IKTImageContainsObjCClasses(header);
+        Dl_info imageInfo = {};
+        if (dladdr(header, &imageInfo) != 0 && imageInfo.dli_fname != NULL) {
+            context->imagePath = strdup(imageInfo.dli_fname);
+        }
+        dispatch_async_f(_legacyImageLoadQueue, context, IKTWaitForObjCImageLoad);
+    } else {
+        IKTScheduleImageDidLoadCallback();
+    }
+    #if defined(IKT_TEST_DELAY_DYLD_CALLBACK)
+    if (atomic_load_explicit(&_delayFutureDyldCallbacks, memory_order_relaxed)) {
+        usleep(100000);
+    }
+    #endif
+}
+
+static void IKTRegisterDyldImageLoadCallback(void) {
+    _dyld_register_func_for_add_image(IKTDyldImageDidLoad);
+    #if defined(IKT_TEST_DELAY_DYLD_CALLBACK)
+    atomic_store_explicit(&_delayFutureDyldCallbacks, 1, memory_order_relaxed);
+    #endif
+}
+
+static void IKTReplayImageDidLoadCallback(ITKImageDidLoadCallback callback) {
+    dispatch_async(_imageDidLoadQueue, ^{
+        callback();
+    });
+}
+
+void IKTRegisterImageDidLoadCallback(ITKImageDidLoadCallback callback) {
+    if (callback == NULL) {
+        return;
+    }
+
+    BOOL shouldInstallHook = NO;
+    os_unfair_lock_lock(&_imageDidLoadLock);
+    for (size_t index = 0; index < _imageDidLoadCallbackCount; index++) {
+        if (_imageDidLoadCallbacks[index] == callback) {
+            os_unfair_lock_unlock(&_imageDidLoadLock);
+            return;
+        }
+    }
+
+    ITKImageDidLoadCallback *callbacks = realloc(
+        _imageDidLoadCallbacks, (_imageDidLoadCallbackCount + 1) * sizeof(*callbacks));
+    if (callbacks == NULL) {
+        os_unfair_lock_unlock(&_imageDidLoadLock);
+        return;
+    }
+    _imageDidLoadCallbacks = callbacks;
+    _imageDidLoadCallbacks[_imageDidLoadCallbackCount++] = callback;
+
+    if (!_imageDidLoadHookInstalled) {
+        _imageDidLoadHookInstalled = YES;
+        _imageDidLoadQueue = dispatch_queue_create("com.steipete.image-watcher", DISPATCH_QUEUE_SERIAL);
+        _legacyImageLoadQueue = dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0);
+        shouldInstallHook = YES;
+    }
+    os_unfair_lock_unlock(&_imageDidLoadLock);
+
+    if (!shouldInstallHook) {
+        IKTReplayImageDidLoadCallback(callback);
+        return;
+    }
+    #if defined(IKT_FORCE_DYLD_IMAGE_CALLBACK)
+    IKTRegisterDyldImageLoadCallback();
+    #else
+    if (@available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)) {
+        objc_addLoadImageFunc(IKTObjCImageDidLoad);
+    } else {
+        IKTRegisterDyldImageLoadCallback();
+    }
+    #endif
+}
+
 void msgSendSuperTrampoline(void);
 void msgSendSuperStretTrampoline(void);