Implemented Sweep

Author: johnno1962

EvalApp/AppDelegate.swift

@@ -26,6 +26,10 @@ class AppDelegate: NSObject, NSApplicationDelegate {
         }
     }
 
+    @objc func injected() {
+        print("I've been injected!")
+    }
+
     func applicationDidFinishLaunching(_ aNotification: Notification) {
         // Insert code here to initialize your application
         evalError = {

README.md

@@ -52,8 +52,8 @@ be running a small server "'signer", included in this project to do this alas.
 Included is a simple implementation of ["code injection"](SwiftEval/SwiftInjection.swift).
 If you are stopped in a class, you can edit the class' implementation, save it and type
 "p inject()". Your changes will be applied without having to restart the application.
-To detect this in your code to reload a view controller for example, subscribe to the
-"INJECTION_BUNDLE_NOTIFICATION".
+To detect this in your code to reload a view controller for example, add an @objc
+injected() method or subscribe to the "INJECTION_BUNDLE_NOTIFICATION".
 
 ### But Why?
 

SwiftEval/AppDelegate.swift

@@ -22,6 +22,8 @@ class AppDelegate: UIResponder, UIApplicationDelegate, UISplitViewControllerDele
         let navigationController = splitViewController.viewControllers[splitViewController.viewControllers.count-1] as! UINavigationController
         navigationController.topViewController!.navigationItem.leftBarButtonItem = splitViewController.displayModeButtonItem
         splitViewController.delegate = self
+//        Bundle(path: "/Applications/Injection.app/Contents/Resources/InjectionLoader.bundle")?.load()
+        FSEventStreamCreate()
         return true
     }
 

SwiftEval/MasterViewController.swift

@@ -13,6 +13,9 @@ class MasterViewController: UITableViewController {
     var detailViewController: DetailViewController? = nil
     var objects = [Any]()
 
+    @objc func injected() {
+        print("I've been injected!")
+    }
 
     override func viewDidLoad() {
         super.viewDidLoad()

SwiftEval/SwiftEval.swift

@@ -5,7 +5,7 @@
 //  Created by John Holdsworth on 02/11/2017.
 //  Copyright © 2017 John Holdsworth. All rights reserved.
 //
-//  $Id: //depot/ResidentEval/SwiftEval/SwiftEval.swift#4 $
+//  $Id: //depot/ResidentEval/SwiftEval/SwiftEval.swift#5 $
 //
 //  Basic implementation of ra Swift "eval()" including the
 //  mechanics of recompiling a class and loading the new
@@ -21,7 +21,7 @@ private func debug(_ str: String) {
 
 /// Error handler
 public var evalError = {
-    (_ err: String) -> AnyClass? in
+    (_ err: String) -> [AnyClass]? in
     print("** \(err) **")
     return nil
 }
@@ -55,7 +55,7 @@ extension NSObject {
         // update evalImpl to implement expression
 
         if NSObject.lastEvalByClass[className] != expression,
-            let newClass = SwiftEval.rebuildClass(oldClass: oldClass, className: className, extra: extra) {
+            let newClass = SwiftEval.rebuildClass(oldClass: oldClass, className: className, extra: extra)?.first {
 
             // swizzle new version of evalImpl onto class
 
@@ -96,7 +96,7 @@ class SwiftEval {
     static var dylibNumber = 0
     static var compileByClass = [String: String]()
 
-    static func rebuildClass(oldClass: AnyClass, className: String, extra: String?) -> AnyClass? {
+    static func rebuildClass(oldClass: AnyClass?, className: String, extra: String?) -> [AnyClass]? {
         let sourceURL = URL(fileURLWithPath: #file)
         guard let derivedData = findDerivedData(url: sourceURL) else {
             return evalError("Could not locate derived data")
@@ -107,7 +107,7 @@ class SwiftEval {
 
         // locate compile command for class
 
-        let regexp = " -primary-file (\"([^\"]+?/\(className)\\.swift)\"|(\\S+?/\(className)\\.swift)) "
+        let regexp = " -primary-file (\"([^\"]*?/\(className)\\.swift)\"|(\\S*?/\(className)\\.swift)) "
 
         guard var compileCommand = compileByClass[className] ?? {
             () -> String? in
@@ -118,7 +118,7 @@ class SwiftEval {
                     echo "Scanning $log"
                     # grep log for build of class source
                     /usr/bin/gunzip <"$log" | perl -lpe 's/\\r/\\n/g' | \
-                    /usr/bin/grep -E '\(regexp)' >/tmp/eval.sh && exit 0;
+                    time /usr/bin/grep -E '\(regexp)' >/tmp/eval.sh && exit 0;
                 done;
                 exit 1
                 """) else {
@@ -224,19 +224,31 @@ class SwiftEval {
             return evalError("dlopen() error: \(String(cString: dlerror()))")
         }
 
-        // find patched version of class using symbol for existing
+        if oldClass != nil {
+            // find patched version of class using symbol for existing
 
-        var info = Dl_info()
-        guard dladdr(unsafeBitCast(oldClass, to: UnsafeRawPointer.self), &info) != 0 else {
-            return evalError("Could not locate class symbol")
-        }
+            var info = Dl_info()
+            guard dladdr(unsafeBitCast(oldClass, to: UnsafeRawPointer.self), &info) != 0 else {
+                return evalError("Could not locate class symbol")
+            }
 
-        debug(String(cString: info.dli_sname))
-        guard let newSymbol = dlsym(dl, info.dli_sname) else {
-            return evalError("Could not locate newly loaded class symbol")
-        }
+            debug(String(cString: info.dli_sname))
+            guard let newSymbol = dlsym(dl, info.dli_sname) else {
+                return evalError("Could not locate newly loaded class symbol")
+            }
 
-        return unsafeBitCast(newSymbol, to: AnyClass.self)
+            return [unsafeBitCast(newSymbol, to: AnyClass.self)]
+        }
+        else {
+            guard shell(command: "\(xcode)/Toolchains/XcodeDefault.xctoolchain/usr/bin/nm /tmp/eval.o | grep 'S _OBJC_CLASS_$_' | awk '{print $3}' >/tmp/eval.classes") else {
+                return evalError("Could not list classes")
+            }
+            guard var symbols = (try? String(contentsOfFile: "/tmp/eval.classes"))?.components(separatedBy: "\n") else {
+                return evalError("Could not load class list")
+            }
+            symbols.removeLast()
+            return symbols.map { unsafeBitCast(dlsym(dl, String($0.dropFirst()))!, to: AnyClass.self) }
+        }
     }
 
     static func findDerivedData(url: URL) -> URL? {

SwiftEval/SwiftInjection.swift

@@ -5,7 +5,7 @@
 //  Created by John Holdsworth on 05/11/2017.
 //  Copyright © 2017 John Holdsworth. All rights reserved.
 //
-//  $Id: //depot/ResidentEval/SwiftEval/SwiftInjection.swift#4 $
+//  $Id: //depot/ResidentEval/SwiftEval/SwiftInjection.swift#8 $
 //
 //  Cut-down version of code injection in Swift. Uses code
 //  from SwiftEval.swift to recompile and reload class.
@@ -14,6 +14,10 @@
 #if arch(x86_64) // simulator/macOS only
 import Foundation
 
+@objc public protocol SwiftInjected {
+    @objc optional func injected()
+}
+
 #if os(iOS)
 import UIKit
 
@@ -33,39 +37,72 @@ extension UIViewController {
         viewDidLoad()
     }
 }
+#else
+import Cocoa
 #endif
 
 extension NSObject {
 
     public func inject() {
-        if let oldClass: AnyClass = object_getClass(self),
-            let newClass = SwiftEval.rebuildClass(oldClass: oldClass, className: "\(oldClass)", extra: nil) {
+        if let oldClass: AnyClass = object_getClass(self) {
+            SwiftInjection.inject(oldClass: oldClass, className: "\(oldClass)")
+        }
+    }
 
-            // old-school swizzle Objective-C class & instance methods
-            injection(swizzle: object_getClass(newClass), onto: object_getClass(oldClass))
-            injection(swizzle: newClass, onto: oldClass)
+    @objc
+    public class func inject(file: String) {
+        let path = URL(fileURLWithPath: file).deletingPathExtension().path
+        SwiftInjection.inject(oldClass: nil, className: String(path.dropFirst()))
+    }
+}
 
-            // overwrite Swift vtable of existing class with implementations from new class
-            let existingClass = unsafeBitCast(oldClass, to: UnsafeMutablePointer<ClassMetadataSwift>.self)
-            let classMetadata = unsafeBitCast(newClass, to: UnsafeMutablePointer<ClassMetadataSwift>.self)
+class SwiftInjection {
 
-            func byteAddr<T>(_ location: UnsafeMutablePointer<T>) -> UnsafeMutablePointer<UInt8> {
-                return location.withMemoryRebound(to: UInt8.self, capacity: 1) { $0 }
-            }
+    static func inject(oldClass: AnyClass?, className: String) {
+        if let newClasses = SwiftEval.rebuildClass(oldClass: oldClass, className: className, extra: nil) {
+            let oldClasses = //oldClass != nil ? [oldClass!] :
+                newClasses.map { objc_getClass(class_getName($0)) as! AnyClass }
+            for i in 0..<oldClasses.count {
+                let oldClass: AnyClass = oldClasses[i], newClass: AnyClass = newClasses[i]
+
+                // old-school swizzle Objective-C class & instance methods
+                injection(swizzle: object_getClass(newClass), onto: object_getClass(oldClass))
+                injection(swizzle: newClass, onto: oldClass)
+
+                // overwrite Swift vtable of existing class with implementations from new class
+                let existingClass = unsafeBitCast(oldClass, to: UnsafeMutablePointer<ClassMetadataSwift>.self)
+                let classMetadata = unsafeBitCast(newClass, to: UnsafeMutablePointer<ClassMetadataSwift>.self)
 
-            let vtableOffset = byteAddr(&existingClass.pointee.IVarDestroyer) - byteAddr(existingClass)
-            let vtableLength = Int(existingClass.pointee.ClassSize -
-                existingClass.pointee.ClassAddressPoint) - vtableOffset
+                func byteAddr<T>(_ location: UnsafeMutablePointer<T>) -> UnsafeMutablePointer<UInt8> {
+                    return location.withMemoryRebound(to: UInt8.self, capacity: 1) { $0 }
+                }
 
-            NSLog("\(unsafeBitCast(classMetadata, to: AnyClass.self)), vtable length: \(vtableLength)")
-            memcpy(byteAddr(existingClass) + vtableOffset, byteAddr(classMetadata) + vtableOffset, vtableLength)
+                let vtableOffset = byteAddr(&existingClass.pointee.IVarDestroyer) - byteAddr(existingClass)
+                let vtableLength = Int(existingClass.pointee.ClassSize -
+                    existingClass.pointee.ClassAddressPoint) - vtableOffset
+
+                NSLog("\(unsafeBitCast(classMetadata, to: AnyClass.self)), vtable length: \(vtableLength)")
+                memcpy(byteAddr(existingClass) + vtableOffset, byteAddr(classMetadata) + vtableOffset, vtableLength)
+
+                // implement -injected() method using sweep of objects in application
+                if class_getInstanceMethod(oldClass, #selector(SwiftInjected.injected)) != nil {
+                    #if os(iOS)
+                    let app = UIApplication.shared
+                    #else
+                    let app = NSApplication.shared
+                    #endif
+                    let seeds: [Any] = [app.delegate as Any] + app.windows
+                    sweepValue(seeds, for: oldClass)
+                    seen.removeAll()
+                }
+            }
 
             let notification = Notification.Name("INJECTION_BUNDLE_NOTIFICATION")
-            NotificationCenter.default.post(name: notification, object: [oldClass])
+            NotificationCenter.default.post(name: notification, object: oldClasses)
         }
     }
 
-    private func injection(swizzle newClass: AnyClass?, onto oldClass: AnyClass?) {
+    static func injection(swizzle newClass: AnyClass?, onto oldClass: AnyClass?) {
         var methodCount: UInt32 = 0
         if let methods = class_copyMethodList(newClass, &methodCount) {
             for i in 0 ..< Int(methodCount) {
@@ -76,6 +113,104 @@ extension NSObject {
             free(methods)
         }
     }
+
+    static func sweepValue(_ value: Any, for targetClass: AnyClass) {
+        let mirror = Mirror(reflecting: value)
+        if var style = mirror.displayStyle {
+            if _typeName(mirror.subjectType).hasPrefix("Swift.ImplicitlyUnwrappedOptional<") {
+                style = .optional
+            }
+            switch style {
+            case .set:
+                fallthrough
+            case .collection:
+                for (_, child) in mirror.children {
+                    sweepValue(child, for: targetClass)
+                }
+                return
+            case .dictionary:
+                for (_, child) in mirror.children {
+                    for (_, element) in Mirror(reflecting: child).children {
+                        sweepValue(element, for: targetClass)
+                    }
+                }
+                return
+            case .class:
+                sweepInstance(value as AnyObject, for: targetClass)
+                return
+            case .optional:
+                if let some = mirror.children.first?.value {
+                    sweepValue(some, for: targetClass)
+                }
+                return
+            default:
+                break
+            }
+        }
+
+        if let style = mirror.displayStyle {
+            switch style {
+            case .enum:
+                if let evals = mirror.children.first?.value {
+                    sweepValue(evals, for: targetClass)
+                }
+            case .tuple:
+                sweepMembers(value, for: targetClass)
+            case .struct:
+                sweepMembers(value, for: targetClass)
+            default:
+                break
+            }
+        }
+    }
+
+    static var seen = [UnsafeRawPointer: Bool]()
+
+    static func sweepInstance(_ instance: AnyObject, for targetClass: AnyClass) {
+        let reference = unsafeBitCast(instance, to: UnsafeRawPointer.self)
+        if seen[reference] == nil {
+            seen[reference] = true
+
+            if object_getClass(instance) == targetClass {
+                let proto = unsafeBitCast(instance, to: SwiftInjected.self)
+                proto.injected?()
+            }
+
+            sweepMembers(instance, for: targetClass)
+            sweepIvars(instance, for: targetClass)
+        }
+    }
+
+    static func sweepMembers(_ instance: Any, for targetClass: AnyClass) {
+        var mirror: Mirror? = Mirror(reflecting: instance)
+        while mirror != nil {
+            for (_, value) in mirror!.children {
+                sweepValue(value, for: targetClass)
+            }
+            mirror = mirror!.superclassMirror
+        }
+    }
+
+    static func sweepIvars(_ instance: AnyObject, for targetClass: AnyClass) {
+        var icnt: UInt32 = 0, cls: AnyClass? = object_getClass(instance)!
+        let object = "@".utf16.first!
+        while cls != nil && cls != NSURL.self {
+            if let ivars = class_copyIvarList(cls, &icnt) {
+                for i in 0 ..< Int(icnt) {
+                    if let type = ivar_getTypeEncoding(ivars[i]), type[0] == object {
+                        (unsafeBitCast(instance, to: UnsafePointer<Int8>.self) + ivar_getOffset(ivars[i]))
+                            .withMemoryRebound(to: AnyObject?.self, capacity: 1) {
+                                if let obj = $0.pointee {
+                                    sweepValue(obj, for: targetClass)
+                                }
+                        }
+                    }
+                }
+                free(ivars)
+            }
+            cls = class_getSuperclass(cls)
+        }
+    }
 }
 
 /**