Merge branch 'fix-transport-provider'

12 files changed, 95 insertions, 81 deletions

diff --git a/ios/MullvadREST/RESTNetworkOperation.swift b/ios/MullvadREST/RESTNetworkOperation.swift
index 27428a2755..447be5776a 100644
--- a/ios/MullvadREST/RESTNetworkOperation.swift
+++ b/ios/MullvadREST/RESTNetworkOperation.swift
@@ -17,7 +17,7 @@ extension REST {
         private let responseHandler: AnyResponseHandler<Success>
 
         private let logger: Logger
-        private let transportProvider: () -> RESTTransport?
+        private let transportProvider: RESTTransportProvider
         private let addressCacheStore: AddressCache
 
         private var networkTask: Cancellable?
@@ -137,8 +137,7 @@ extension REST {
         private func didReceiveURLRequest(_ restRequest: REST.Request, endpoint: AnyIPEndpoint) {
             dispatchPrecondition(condition: .onQueue(dispatchQueue))
 
-            let transport = transportProvider()
-            guard let transport else {
+            guard let transport = transportProvider.makeTransport() else {
                 logger.error("Failed to obtain transport.")
                 finish(result: .failure(REST.Error.transport(InternalTransportError.noTransport)))
                 return
diff --git a/ios/MullvadREST/RESTProxy.swift b/ios/MullvadREST/RESTProxy.swift
index 5c8a14df81..0206bef0df 100644
--- a/ios/MullvadREST/RESTProxy.swift
+++ b/ios/MullvadREST/RESTProxy.swift
@@ -67,11 +67,11 @@ extension REST {
     }
 
     public class ProxyConfiguration {
-        public let transportProvider: () -> RESTTransport?
+        public let transportProvider: RESTTransportProvider
         public let addressCacheStore: AddressCache
 
         public init(
-            transportProvider: @escaping () -> RESTTransport?,
+            transportProvider: RESTTransportProvider,
             addressCacheStore: AddressCache
         ) {
             self.transportProvider = transportProvider
diff --git a/ios/MullvadREST/RESTProxyFactory.swift b/ios/MullvadREST/RESTProxyFactory.swift
index 215c9c61e9..2436cbf3e8 100644
--- a/ios/MullvadREST/RESTProxyFactory.swift
+++ b/ios/MullvadREST/RESTProxyFactory.swift
@@ -13,7 +13,7 @@ extension REST {
         public let configuration: AuthProxyConfiguration
 
         public class func makeProxyFactory(
-            transportProvider: @escaping () -> RESTTransport?,
+            transportProvider: RESTTransportProvider,
             addressCache: AddressCache
         ) -> ProxyFactory {
             let basicConfiguration = REST.ProxyConfiguration(
diff --git a/ios/MullvadREST/RESTTransport.swift b/ios/MullvadREST/RESTTransport.swift
index 54b5e561d6..eb87b6db99 100644
--- a/ios/MullvadREST/RESTTransport.swift
+++ b/ios/MullvadREST/RESTTransport.swift
@@ -11,8 +11,6 @@ import MullvadTypes
 
 public protocol RESTTransport {
     var name: String { get }
-    func sendRequest(
-        _ request: URLRequest,
-        completion: @escaping (Data?, URLResponse?, Error?) -> Void
-    ) -> Cancellable
+
+    func sendRequest(_ request: URLRequest, completion: @escaping (Data?, URLResponse?, Error?) -> Void) -> Cancellable
 }
diff --git a/ios/MullvadREST/RESTTransportProvider.swift b/ios/MullvadREST/RESTTransportProvider.swift
new file mode 100644
index 0000000000..ed683fa163
--- /dev/null
+++ b/ios/MullvadREST/RESTTransportProvider.swift
@@ -0,0 +1,25 @@
+//
+//  RESTTransportProvider.swift
+//  MullvadREST
+//
+//  Created by pronebird on 24/08/2023.
+//  Copyright © 2023 Mullvad VPN AB. All rights reserved.
+//
+
+import Foundation
+
+public protocol RESTTransportProvider {
+    func makeTransport() -> RESTTransport?
+}
+
+public struct AnyTransportProvider: RESTTransportProvider {
+    private let block: () -> RESTTransport?
+
+    public init(_ block: @escaping () -> RESTTransport?) {
+        self.block = block
+    }
+
+    public func makeTransport() -> RESTTransport? {
+        return block()
+    }
+}
diff --git a/ios/MullvadTransport/TransportProvider.swift b/ios/MullvadTransport/TransportProvider.swift
index e741904457..c2c825f000 100644
--- a/ios/MullvadTransport/TransportProvider.swift
+++ b/ios/MullvadTransport/TransportProvider.swift
@@ -13,7 +13,7 @@ import MullvadTypes
 import RelayCache
 import RelaySelector
 
-public final class TransportProvider: RESTTransport {
+public final class TransportProvider: RESTTransportProvider {
     private let urlSessionTransport: URLSessionTransport
     private let relayCache: RelayCache
     private let logger = Logger(label: "TransportProvider")
@@ -49,9 +49,20 @@ public final class TransportProvider: RESTTransport {
         }
     }
 
-    // MARK: -
+    public func makeTransport() -> RESTTransport? {
+        parallelRequestsMutex.withLock {
+            guard let actualTransport = makeTransportInner() else { return nil }
 
-    // MARK: RESTTransport implementation
+            let currentStrategy = transportStrategy
+            return TransportWrapper(wrapped: actualTransport) { [weak self] error in
+                if let error = error as? URLError, error.shouldResetNetworkTransport {
+                    self?.resetTransportMatching(currentStrategy)
+                }
+            }
+        }
+    }
+
+    // MARK: -
 
     private func shadowsocks() -> RESTTransport? {
         do {
@@ -112,35 +123,6 @@ public final class TransportProvider: RESTTransport {
 
     // MARK: -
 
-    // MARK: RESTTransport implementation
-
-    public var name: String { currentTransport?.name ?? "TransportProvider" }
-
-    public func sendRequest(
-        _ request: URLRequest,
-        completion: @escaping (Data?, URLResponse?, Error?) -> Void
-    ) -> Cancellable {
-        parallelRequestsMutex.lock()
-        defer {
-            parallelRequestsMutex.unlock()
-        }
-
-        let currentStrategy = transportStrategy
-        guard let transport = makeTransport() else { return AnyCancellable() }
-
-        let failureCompletionHandler: (Data?, URLResponse?, Error?)
-            -> Void = { [weak self] data, response, error in
-                guard let self else { return }
-
-                if let error = error as? URLError, error.shouldResetNetworkTransport {
-                    resetTransportMatching(currentStrategy)
-                }
-                completion(data, response, error)
-            }
-
-        return transport.sendRequest(request, completion: failureCompletionHandler)
-    }
-
     /// When several requests fail at the same time,  prevents the `transportStrategy` from switching multiple times.
     ///
     /// The `strategy` is checked against the `transportStrategy`. When several requests are made and fail in parallel,
@@ -162,7 +144,7 @@ public final class TransportProvider: RESTTransport {
     /// > Warning: Do not  lock the `parallelRequestsMutex` in this method
     ///
     /// - Returns: A `RESTTransport` object to make a connection
-    private func makeTransport() -> RESTTransport? {
+    private func makeTransportInner() -> RESTTransport? {
         if currentTransport == nil {
             switch transportStrategy.connectionTransport() {
             case .useShadowsocks:
@@ -190,3 +172,23 @@ private extension URLError {
             code != .callIsActive
     }
 }
+
+/// Interstitial implementation of `RESTTransport` that intercepts the completion of the wrapped transport.
+private struct TransportWrapper: RESTTransport {
+    let wrapped: RESTTransport
+    let onComplete: (Error?) -> Void
+
+    var name: String {
+        return wrapped.name
+    }
+
+    func sendRequest(
+        _ request: URLRequest,
+        completion: @escaping (Data?, URLResponse?, Error?) -> Void
+    ) -> Cancellable {
+        return wrapped.sendRequest(request) { data, response, error in
+            onComplete(error)
+            completion(data, response, error)
+        }
+    }
+}
diff --git a/ios/MullvadVPN.xcodeproj/project.pbxproj b/ios/MullvadVPN.xcodeproj/project.pbxproj
index 0726d2b1b1..8bb63c8de1 100644
--- a/ios/MullvadVPN.xcodeproj/project.pbxproj
+++ b/ios/MullvadVPN.xcodeproj/project.pbxproj
@@ -356,6 +356,7 @@
 		58E45A5729F12C5100281ECF /* Result+Extensions.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58F1311427E0B2AB007AC5BC /* Result+Extensions.swift */; };
 		58E511E628DDDEAC00B0BCDE /* CodingErrors+CustomErrorDescription.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58E511E528DDDEAC00B0BCDE /* CodingErrors+CustomErrorDescription.swift */; };
 		58E511E828DDDF2400B0BCDE /* CodingErrors+CustomErrorDescription.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58E511E528DDDEAC00B0BCDE /* CodingErrors+CustomErrorDescription.swift */; };
+		58E7BA192A975DF70068EC3A /* RESTTransportProvider.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58E7BA182A975DF70068EC3A /* RESTTransportProvider.swift */; };
 		58EC067A2A8D208D00BEB973 /* MockTunnelDeviceInfo.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58EC06792A8D208D00BEB973 /* MockTunnelDeviceInfo.swift */; };
 		58EC067C2A8D2A0B00BEB973 /* NetworkCounters.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58EC067B2A8D2A0B00BEB973 /* NetworkCounters.swift */; };
 		58ED3A142A7C199C0085CE65 /* StartOptions.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58ED3A132A7C199C0085CE65 /* StartOptions.swift */; };
@@ -1259,6 +1260,7 @@
 		58E511E328DDDE8900B0BCDE /* CustomErrorDescriptionProtocol.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = CustomErrorDescriptionProtocol.swift; sourceTree = "<group>"; };
 		58E511E528DDDEAC00B0BCDE /* CodingErrors+CustomErrorDescription.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "CodingErrors+CustomErrorDescription.swift"; sourceTree = "<group>"; };
 		58E511EA28DDE18400B0BCDE /* Error+Chain.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "Error+Chain.swift"; sourceTree = "<group>"; };
+		58E7BA182A975DF70068EC3A /* RESTTransportProvider.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = RESTTransportProvider.swift; sourceTree = "<group>"; };
 		58E973DD24850EB600096F90 /* AsyncOperation.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AsyncOperation.swift; sourceTree = "<group>"; };
 		58EC06792A8D208D00BEB973 /* MockTunnelDeviceInfo.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = MockTunnelDeviceInfo.swift; sourceTree = "<group>"; };
 		58EC067B2A8D2A0B00BEB973 /* NetworkCounters.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = NetworkCounters.swift; sourceTree = "<group>"; };
@@ -1627,6 +1629,7 @@
 				06FAE66528F83CA30033DD93 /* RESTURLSession.swift */,
 				06FAE67728F83CA40033DD93 /* ServerRelaysResponse.swift */,
 				06FAE66B28F83CA30033DD93 /* SSLPinningURLSessionDelegate.swift */,
+				58E7BA182A975DF70068EC3A /* RESTTransportProvider.swift */,
 			);
 			path = MullvadREST;
 			sourceTree = "<group>";
@@ -3495,6 +3498,7 @@
 				06799AF428F98E4800ACD94E /* RESTAuthorization.swift in Sources */,
 				06799AE228F98E4800ACD94E /* RESTRequestFactory.swift in Sources */,
 				06799AEC28F98E4800ACD94E /* RESTTaskIdentifier.swift in Sources */,
+				58E7BA192A975DF70068EC3A /* RESTTransportProvider.swift in Sources */,
 				06799ADE28F98E4800ACD94E /* RESTRequestHandler.swift in Sources */,
 				06799AEF28F98E4800ACD94E /* RESTRetryStrategy.swift in Sources */,
 				06799AE128F98E4800ACD94E /* SSLPinningURLSessionDelegate.swift in Sources */,
diff --git a/ios/MullvadVPN/AppDelegate.swift b/ios/MullvadVPN/AppDelegate.swift
index 0cb57e863b..930f700246 100644
--- a/ios/MullvadVPN/AppDelegate.swift
+++ b/ios/MullvadVPN/AppDelegate.swift
@@ -59,7 +59,9 @@ class AppDelegate: UIResponder, UIApplicationDelegate, UNUserNotificationCenterD
         addressCache.loadFromFile()
 
         proxyFactory = REST.ProxyFactory.makeProxyFactory(
-            transportProvider: { [weak self] in self?.transportMonitor },
+            transportProvider: AnyTransportProvider { [weak self] in
+                return self?.transportMonitor.makeTransport()
+            },
             addressCache: addressCache
         )
 
@@ -119,7 +121,10 @@ class AppDelegate: UIResponder, UIApplicationDelegate, UNUserNotificationCenterD
 
         #if targetEnvironment(simulator)
         // Configure mock tunnel provider on simulator
-        simulatorTunnelProviderHost = SimulatorTunnelProviderHost(relayCacheTracker: relayCacheTracker)
+        simulatorTunnelProviderHost = SimulatorTunnelProviderHost(
+            relayCacheTracker: relayCacheTracker,
+            transportProvider: transportProvider
+        )
         SimulatorTunnelProvider.shared.delegate = simulatorTunnelProviderHost
         #endif
 
diff --git a/ios/MullvadVPN/SimulatorTunnelProvider/SimulatorTunnelProviderHost.swift b/ios/MullvadVPN/SimulatorTunnelProvider/SimulatorTunnelProviderHost.swift
index b347d080d1..2921905ea6 100644
--- a/ios/MullvadVPN/SimulatorTunnelProvider/SimulatorTunnelProviderHost.swift
+++ b/ios/MullvadVPN/SimulatorTunnelProvider/SimulatorTunnelProviderHost.swift
@@ -26,15 +26,11 @@ final class SimulatorTunnelProviderHost: SimulatorTunnelProviderDelegate {
     private let providerLogger = Logger(label: "SimulatorTunnelProviderHost")
     private let dispatchQueue = DispatchQueue(label: "SimulatorTunnelProviderHostQueue")
 
-    init(relayCacheTracker: RelayCacheTracker) {
+    init(relayCacheTracker: RelayCacheTracker, transportProvider: TransportProvider) {
         self.relayCacheTracker = relayCacheTracker
-
-        let urlSession = REST.makeURLSession()
-        let urlSessionTransport = URLSessionTransport(urlSession: urlSession)
-
         self.urlRequestProxy = URLRequestProxy(
             dispatchQueue: dispatchQueue,
-            transportProvider: { urlSessionTransport }
+            transportProvider: transportProvider
         )
     }
 
diff --git a/ios/MullvadVPN/TransportMonitor/TransportMonitor.swift b/ios/MullvadVPN/TransportMonitor/TransportMonitor.swift
index 4abf5cf43b..99a59cc116 100644
--- a/ios/MullvadVPN/TransportMonitor/TransportMonitor.swift
+++ b/ios/MullvadVPN/TransportMonitor/TransportMonitor.swift
@@ -14,7 +14,7 @@ import MullvadTypes
 import RelayCache
 import RelaySelector
 
-final class TransportMonitor: RESTTransport {
+final class TransportMonitor: RESTTransportProvider {
     private let tunnelManager: TunnelManager
     private let tunnelStore: TunnelStore
     private let transportProvider: TransportProvider
@@ -29,39 +29,22 @@ final class TransportMonitor: RESTTransport {
         self.transportProvider = transportProvider
     }
 
-    var name: String { selectTransport(transportProvider).name }
-
-    func sendRequest(_ request: URLRequest, completion: @escaping (Data?, URLResponse?, Error?) -> Void) -> Cancellable
-    {
-        selectTransport(transportProvider).sendRequest(request, completion: completion)
-    }
-
-    // MARK: -
-
-    // MARK: Private API
-
     /// Selects a transport to use for sending an `URLRequest`
     ///
-    /// This method returns the appropriate transport layer based on whether a tunnel is available, and whether it
-    /// should be bypassed
-    /// whenever a transport is requested.
+    /// This method returns the appropriate transport layer based on whether a tunnel is available, and whether it should be bypassed whenever a transport is
+    /// requested.
     ///
-    /// - Parameters:
-    ///   - transport: The transport to use if there is no tunnel, or if it shouldn't be bypassed
     /// - Returns: A transport to use for sending an `URLRequest`
-    private func selectTransport(_ transport: RESTTransport) -> RESTTransport {
+    func makeTransport() -> RESTTransport? {
         let tunnel = tunnelStore.getPersistentTunnels().first { tunnel in
-            tunnel.status == .connecting ||
-                tunnel.status == .reasserting ||
-                tunnel.status == .connected
+            tunnel.status == .connecting || tunnel.status == .reasserting || tunnel.status == .connected
         }
 
         if let tunnel, shouldByPassVPN(tunnel: tunnel) {
-            return PacketTunnelTransport(
-                tunnel: tunnel
-            )
+            return PacketTunnelTransport(tunnel: tunnel)
+        } else {
+            return transportProvider.makeTransport()
         }
-        return transport
     }
 
     private func shouldByPassVPN(tunnel: Tunnel) -> Bool {
diff --git a/ios/PacketTunnel/PacketTunnelProvider.swift b/ios/PacketTunnel/PacketTunnelProvider.swift
index 11bde83717..09d7bd8992 100644
--- a/ios/PacketTunnel/PacketTunnelProvider.swift
+++ b/ios/PacketTunnel/PacketTunnelProvider.swift
@@ -155,13 +155,13 @@ class PacketTunnelProvider: NEPacketTunnelProvider {
         )
 
         let proxyFactory = REST.ProxyFactory.makeProxyFactory(
-            transportProvider: { transportProvider },
+            transportProvider: transportProvider,
             addressCache: addressCache
         )
 
         urlRequestProxy = URLRequestProxy(
             dispatchQueue: dispatchQueue,
-            transportProvider: { transportProvider }
+            transportProvider: transportProvider
         )
         accountsProxy = proxyFactory.createAccountsProxy()
         devicesProxy = proxyFactory.createDevicesProxy()
diff --git a/ios/TunnelProviderMessaging/URLRequestProxy.swift b/ios/TunnelProviderMessaging/URLRequestProxy.swift
index 7fe2e060cf..e7b92c34b2 100644
--- a/ios/TunnelProviderMessaging/URLRequestProxy.swift
+++ b/ios/TunnelProviderMessaging/URLRequestProxy.swift
@@ -15,14 +15,14 @@ public final class URLRequestProxy {
     /// Serial queue used for synchronizing access to class members.
     private let dispatchQueue: DispatchQueue
 
-    private let transportProvider: () -> RESTTransport?
+    private let transportProvider: RESTTransportProvider
 
     /// List of all proxied network requests bypassing VPN.
     private var proxiedRequests: [UUID: Cancellable] = [:]
 
     public init(
         dispatchQueue: DispatchQueue,
-        transportProvider: @escaping () -> RESTTransport?
+        transportProvider: RESTTransportProvider
     ) {
         self.dispatchQueue = dispatchQueue
         self.transportProvider = transportProvider
@@ -33,7 +33,8 @@ public final class URLRequestProxy {
         completionHandler: @escaping (ProxyURLResponse) -> Void
     ) {
         dispatchQueue.async {
-            guard let transportProvider = self.transportProvider() else { return }
+            guard let transportProvider = self.transportProvider.makeTransport() else { return }
+
             // The task sent by `transport.sendRequest` comes in an already resumed state
             let task = transportProvider.sendRequest(proxyRequest.urlRequest) { [weak self] data, response, error in
                 guard let self else { return }
@@ -46,6 +47,7 @@ public final class URLRequestProxy {
                     completionHandler(response)
                 }
             }
+
             // All tasks should have unique identifiers, but if not, cancel the task scheduled
             // earlier.
             let oldTask = self.addRequest(identifier: proxyRequest.id, task: task)