Extract method swizzling

3 files changed, 45 insertions, 33 deletions

diff --git a/ios/MullvadVPN.xcodeproj/project.pbxproj b/ios/MullvadVPN.xcodeproj/project.pbxproj
index 29b107b020..d3c509d4b6 100644
--- a/ios/MullvadVPN.xcodeproj/project.pbxproj
+++ b/ios/MullvadVPN.xcodeproj/project.pbxproj
@@ -98,6 +98,7 @@
 		586AA296234B696B00502875 /* WireguardAssociatedAddresses.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58B8743122B25A7600015324 /* WireguardAssociatedAddresses.swift */; };
 		586BD68322B7BBD800BB7F9F /* NetworkExtension.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 586BD68222B7BBD800BB7F9F /* NetworkExtension.framework */; };
 		586BD68422B7BBE400BB7F9F /* NetworkExtension.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 586BD68222B7BBD800BB7F9F /* NetworkExtension.framework */; };
+		5871FB8325498CA20051A0A4 /* Swizzle.swift in Sources */ = {isa = PBXBuildFile; fileRef = 5871FB8225498CA20051A0A4 /* Swizzle.swift */; };
 		5873884D239E6D7E00E96C4E /* EmbeddedViewContainerView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 5873884C239E6D7E00E96C4E /* EmbeddedViewContainerView.swift */; };
 		587425C12299833500CA2045 /* RootContainerViewController.swift in Sources */ = {isa = PBXBuildFile; fileRef = 587425C02299833500CA2045 /* RootContainerViewController.swift */; };
 		5877153023981F7B001F8237 /* WireguardKeysViewController.swift in Sources */ = {isa = PBXBuildFile; fileRef = 5877152F23981F7B001F8237 /* WireguardKeysViewController.swift */; };
@@ -304,6 +305,7 @@
 		5866F39B2243B82D00168AE5 /* MullvadVPN.entitlements */ = {isa = PBXFileReference; lastKnownFileType = text.plist.entitlements; path = MullvadVPN.entitlements; sourceTree = "<group>"; };
 		5868585424054096000B8131 /* AppButton.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = AppButton.swift; sourceTree = "<group>"; };
 		586BD68222B7BBD800BB7F9F /* NetworkExtension.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = NetworkExtension.framework; path = System/Library/Frameworks/NetworkExtension.framework; sourceTree = SDKROOT; };
+		5871FB8225498CA20051A0A4 /* Swizzle.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Swizzle.swift; sourceTree = "<group>"; };
 		5873884C239E6D7E00E96C4E /* EmbeddedViewContainerView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = EmbeddedViewContainerView.swift; sourceTree = "<group>"; };
 		587425C02299833500CA2045 /* RootContainerViewController.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = RootContainerViewController.swift; sourceTree = "<group>"; };
 		5877152F23981F7B001F8237 /* WireguardKeysViewController.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = WireguardKeysViewController.swift; sourceTree = "<group>"; };
@@ -598,6 +600,7 @@
 				58F19E34228C15BA00C7710B /* SpinnerActivityIndicatorView.swift */,
 				581CBCED229826FD00727D7F /* StaticTableViewDataSource.swift */,
 				5807E2BF2432038B00F5FF30 /* String+Split.swift */,
+				5871FB8225498CA20051A0A4 /* Swizzle.swift */,
 				5862805322428EF100F5A6E1 /* TranslucentButtonBlurView.swift */,
 				5835B7CB233B76CB0096D79F /* TunnelManager.swift */,
 				587AD7C523421D7000E93A53 /* TunnelSettings.swift */,
@@ -980,6 +983,7 @@
 				580EE20924B3224200F9D8A1 /* RetryOperation.swift in Sources */,
 				582AE3102440A6CA00E6733A /* AccountTokenInput.swift in Sources */,
 				58FAEDF7245088E100CB0F5B /* Keychain.swift in Sources */,
+				5871FB8325498CA20051A0A4 /* Swizzle.swift in Sources */,
 				58907D9524D17B4E00CFC3F5 /* DisconnectSplitButton.swift in Sources */,
 				5888AD87227B17950051EB06 /* SelectLocationViewController.swift in Sources */,
 				580EE20424B321EC00F9D8A1 /* OperationObserver.swift in Sources */,
diff --git a/ios/MullvadVPN/AlertPresenter.swift b/ios/MullvadVPN/AlertPresenter.swift
index 9718842c3c..8594c8211e 100644
--- a/ios/MullvadVPN/AlertPresenter.swift
+++ b/ios/MullvadVPN/AlertPresenter.swift
@@ -19,8 +19,15 @@ class AlertPresenter {
     private let operationQueue = OperationQueue()
     private lazy var exclusivityController = ExclusivityController<ExclusivityCategory>(operationQueue: operationQueue)
 
+    private static let initClass: Void = {
+        /// Swizzle `viewDidDisappear` on `UIAlertController` in order to be able to
+        /// detect when the controller disappears.
+        /// The event is broadcasted via `kUIAlertControllerDidDissmissNotification` notification.
+        swizzleMethod(aClass: UIAlertController.self, originalSelector: #selector(UIAlertController.viewDidDisappear(_:)), newSelector: #selector(UIAlertController.alertPresenter_viewDidDisappear(_:)))
+    }()
+
     init() {
-        _ = AlertPresenterUIKitHooks.once
+        _ = Self.initClass
     }
 
     func enqueue(_ alertController: UIAlertController, presentingController: UIViewController, presentCompletion: (() -> Void)? = nil) {
@@ -35,6 +42,19 @@ class AlertPresenter {
 
 }
 
+
+fileprivate extension UIAlertController {
+    @objc dynamic func alertPresenter_viewDidDisappear(_ animated: Bool) {
+        // Call super implementation
+        alertPresenter_viewDidDisappear(animated)
+
+        if presentingViewController == nil {
+            NotificationCenter.default.post(name: kUIAlertControllerDidDissmissNotification, object: self)
+        }
+    }
+}
+
+
 private class PresentAlertOperation: AsyncOperation {
     private let alertController: UIAlertController
     private let presentingController: UIViewController
@@ -63,35 +83,3 @@ private class PresentAlertOperation: AsyncOperation {
         }
     }
 }
-
-/// A helper struct that swizzles `viewDidDisappear` on `UIAlertController` in order to be able to
-/// detect when the controller disappears.
-/// The event is broadcasted via `kUIAlertControllerDidDissmissNotification` notification.
-private struct AlertPresenterUIKitHooks {
-    typealias MethodType = @convention(c) (UIAlertController, Selector, Bool) -> Void
-    typealias BlockImpType = @convention(block) (UIAlertController, Bool) -> Void
-
-    static let once = AlertPresenterUIKitHooks()
-
-    private init() {
-        let originalSelector = #selector(UIAlertController.viewDidDisappear(_:))
-        let originalMethod = class_getInstanceMethod(UIAlertController.self, originalSelector)!
-
-        var originalIMP: IMP? = nil
-        let swizzledBlockIMP: BlockImpType = { (receiver, animated) in
-            let superIMP = originalIMP.map { unsafeBitCast($0, to: MethodType.self) }
-            superIMP?(receiver, originalSelector, animated)
-
-            Self.handleViewDidDisappear(receiver, animated)
-        }
-
-        let swizzledIMP = imp_implementationWithBlock(unsafeBitCast(swizzledBlockIMP, to: AnyObject.self))
-        originalIMP = method_setImplementation(originalMethod, swizzledIMP)
-    }
-
-    private static func handleViewDidDisappear(_ alertController: UIAlertController, _ animated: Bool) {
-        if alertController.presentingViewController == nil {
-            NotificationCenter.default.post(name: kUIAlertControllerDidDissmissNotification, object: alertController)
-        }
-    }
-}
diff --git a/ios/MullvadVPN/Swizzle.swift b/ios/MullvadVPN/Swizzle.swift
new file mode 100644
index 0000000000..994285ce56
--- /dev/null
+++ b/ios/MullvadVPN/Swizzle.swift
@@ -0,0 +1,20 @@
+//
+//  Swizzle.swift
+//  MullvadVPN
+//
+//  Created by pronebird on 28/10/2020.
+//  Copyright © 2020 Mullvad VPN AB. All rights reserved.
+//
+
+import Foundation
+
+@inlinable func swizzleMethod(aClass: AnyClass, originalSelector: Selector, newSelector: Selector) {
+    guard let originalMethod = class_getInstanceMethod(aClass, originalSelector),
+          let newMethod = class_getInstanceMethod(aClass, newSelector) else { return }
+
+    if class_addMethod(aClass, originalSelector, method_getImplementation(newMethod), method_getTypeEncoding(newMethod)) {
+        class_replaceMethod(aClass, newSelector, method_getImplementation(originalMethod), method_getTypeEncoding(originalMethod))
+    } else {
+        method_exchangeImplementations(originalMethod, newMethod)
+    }
+}