Merge branch 'optimal-bridge-selection-ios-122'

8 files changed, 186 insertions, 96 deletions

diff --git a/ios/MullvadREST/ServerRelaysResponse.swift b/ios/MullvadREST/ServerRelaysResponse.swift
index 081802196b..7241d0271f 100644
--- a/ios/MullvadREST/ServerRelaysResponse.swift
+++ b/ios/MullvadREST/ServerRelaysResponse.swift
@@ -48,30 +48,6 @@ extension REST {
         public let ipv6AddrIn: IPv6Address
         public let publicKey: Data
         public let includeInCountry: Bool
-
-        public init(
-            hostname: String,
-            active: Bool,
-            owned: Bool,
-            location: String,
-            provider: String,
-            weight: UInt64,
-            ipv4AddrIn: IPv4Address,
-            ipv6AddrIn: IPv6Address,
-            publicKey: Data,
-            includeInCountry: Bool
-        ) {
-            self.hostname = hostname
-            self.active = active
-            self.owned = owned
-            self.location = location
-            self.provider = provider
-            self.weight = weight
-            self.ipv4AddrIn = ipv4AddrIn
-            self.ipv6AddrIn = ipv6AddrIn
-            self.publicKey = publicKey
-            self.includeInCountry = includeInCountry
-        }
     }
 
     public struct ServerWireguardTunnels: Codable, Equatable {
@@ -79,18 +55,6 @@ extension REST {
         public let ipv6Gateway: IPv6Address
         public let portRanges: [[UInt16]]
         public let relays: [ServerRelay]
-
-        public init(
-            ipv4Gateway: IPv4Address,
-            ipv6Gateway: IPv6Address,
-            portRanges: [[UInt16]],
-            relays: [REST.ServerRelay]
-        ) {
-            self.ipv4Gateway = ipv4Gateway
-            self.ipv6Gateway = ipv6Gateway
-            self.portRanges = portRanges
-            self.relays = relays
-        }
     }
 
     public struct ServerShadowsocks: Codable, Equatable {
@@ -98,38 +62,16 @@ extension REST {
         public let port: UInt16
         public let cipher: String
         public let password: String
-
-        public init(protocol: String, port: UInt16, cipher: String, password: String) {
-            self.protocol = `protocol`
-            self.port = port
-            self.cipher = cipher
-            self.password = password
-        }
     }
 
     public struct ServerBridges: Codable, Equatable {
         public let shadowsocks: [ServerShadowsocks]
         public let relays: [BridgeRelay]
-
-        public init(shadowsocks: [REST.ServerShadowsocks], relays: [BridgeRelay]) {
-            self.shadowsocks = shadowsocks
-            self.relays = relays
-        }
     }
 
     public struct ServerRelaysResponse: Codable, Equatable {
         public let locations: [String: ServerLocation]
         public let wireguard: ServerWireguardTunnels
         public let bridge: ServerBridges
-
-        public init(
-            locations: [String: REST.ServerLocation],
-            wireguard: REST.ServerWireguardTunnels,
-            bridge: ServerBridges
-        ) {
-            self.locations = locations
-            self.wireguard = wireguard
-            self.bridge = bridge
-        }
     }
 }
diff --git a/ios/MullvadTransport/TransportProvider.swift b/ios/MullvadTransport/TransportProvider.swift
index cb24300153..e741904457 100644
--- a/ios/MullvadTransport/TransportProvider.swift
+++ b/ios/MullvadTransport/TransportProvider.swift
@@ -23,19 +23,30 @@ public final class TransportProvider: RESTTransport {
 
     private var currentTransport: RESTTransport?
     private let parallelRequestsMutex = NSLock()
+    private var relayConstraints = RelayConstraints()
+    private let constraintsUpdater: RelayConstraintsUpdater
 
     public init(
         urlSessionTransport: URLSessionTransport,
         relayCache: RelayCache,
         addressCache: REST.AddressCache,
         shadowsocksCache: ShadowsocksConfigurationCache,
-        transportStrategy: TransportStrategy = .init()
+        transportStrategy: TransportStrategy = .init(),
+        constraintsUpdater: RelayConstraintsUpdater
     ) {
         self.urlSessionTransport = urlSessionTransport
         self.relayCache = relayCache
         self.addressCache = addressCache
         self.shadowsocksCache = shadowsocksCache
         self.transportStrategy = transportStrategy
+        self.constraintsUpdater = constraintsUpdater
+        constraintsUpdater.onNewConstraints = { [weak self] newConstraints in
+            self?.parallelRequestsMutex.lock()
+            defer {
+                self?.parallelRequestsMutex.unlock()
+            }
+            self?.relayConstraints = newConstraints
+        }
     }
 
     // MARK: -
@@ -75,10 +86,13 @@ public final class TransportProvider: RESTTransport {
     /// Returns a randomly selected shadowsocks configuration.
     private func makeNewShadowsocksConfiguration() throws -> ShadowsocksConfiguration {
         let cachedRelays = try relayCache.read()
-        let bridgeAddress = RelaySelector.getShadowsocksRelay(relays: cachedRelays.relays)?.ipv4AddrIn
-        let bridgeConfiguration = RelaySelector.getShadowsocksTCPBridge(relays: cachedRelays.relays)
+        let bridgeConfiguration = RelaySelector.shadowsocksTCPBridge(from: cachedRelays.relays)
+        let closestRelay = RelaySelector.closestShadowsocksRelayConstrained(
+            by: relayConstraints,
+            in: cachedRelays.relays
+        )
 
-        guard let bridgeAddress, let bridgeConfiguration else { throw POSIXError(.ENOENT) }
+        guard let bridgeAddress = closestRelay?.ipv4AddrIn, let bridgeConfiguration else { throw POSIXError(.ENOENT) }
 
         let newConfiguration = ShadowsocksConfiguration(
             bridgeAddress: bridgeAddress,
diff --git a/ios/MullvadTypes/RelayConstraints.swift b/ios/MullvadTypes/RelayConstraints.swift
index c37833823b..f0bc09bf25 100644
--- a/ios/MullvadTypes/RelayConstraints.swift
+++ b/ios/MullvadTypes/RelayConstraints.swift
@@ -8,6 +8,18 @@
 
 import Foundation
 
+public protocol ConstraintsPropagation {
+    var onNewConstraints: ((RelayConstraints) -> Void)? { get set }
+}
+
+public class RelayConstraintsUpdater: ConstraintsPropagation {
+    public var onNewConstraints: ((RelayConstraints) -> Void)?
+
+    public init(onNewConstraints: ((RelayConstraints) -> Void)? = nil) {
+        self.onNewConstraints = onNewConstraints
+    }
+}
+
 public struct RelayConstraints: Codable, Equatable, CustomDebugStringConvertible {
     public var location: RelayConstraint<RelayLocation>
 
diff --git a/ios/MullvadVPN/AppDelegate.swift b/ios/MullvadVPN/AppDelegate.swift
index 149348fbc7..c8003c068e 100644
--- a/ios/MullvadVPN/AppDelegate.swift
+++ b/ios/MullvadVPN/AppDelegate.swift
@@ -40,6 +40,7 @@ class AppDelegate: UIResponder, UIApplicationDelegate, UNUserNotificationCenterD
     private(set) var relayCacheTracker: RelayCacheTracker!
     private(set) var storePaymentManager: StorePaymentManager!
     private var transportMonitor: TransportMonitor!
+    private var relayConstraintsObserver: TunnelBlockObserver!
 
     // MARK: - Application lifecycle
 
@@ -85,6 +86,11 @@ class AppDelegate: UIResponder, UIApplicationDelegate, UNUserNotificationCenterD
             apiProxy: apiProxy
         )
 
+        let constraintsUpdater = RelayConstraintsUpdater()
+        relayConstraintsObserver = TunnelBlockObserver(didUpdateTunnelSettings: { _, settings in
+            constraintsUpdater.onNewConstraints?(settings.relayConstraints)
+        })
+
         storePaymentManager = StorePaymentManager(
             application: application,
             queue: .default(),
@@ -98,9 +104,12 @@ class AppDelegate: UIResponder, UIApplicationDelegate, UNUserNotificationCenterD
             urlSessionTransport: urlSessionTransport,
             relayCache: relayCache,
             addressCache: addressCache,
-            shadowsocksCache: shadowsocksCache
+            shadowsocksCache: shadowsocksCache,
+            constraintsUpdater: constraintsUpdater
         )
 
+        tunnelManager.addObserver(relayConstraintsObserver)
+
         transportMonitor = TransportMonitor(
             tunnelManager: tunnelManager,
             tunnelStore: tunnelStore,
diff --git a/ios/MullvadVPNTests/RelayCacheTests.swift b/ios/MullvadVPNTests/RelayCacheTests.swift
index 7e37eefe22..03899e543b 100644
--- a/ios/MullvadVPNTests/RelayCacheTests.swift
+++ b/ios/MullvadVPNTests/RelayCacheTests.swift
@@ -6,7 +6,7 @@
 //  Copyright © 2023 Mullvad VPN AB. All rights reserved.
 //
 
-import MullvadREST
+@testable import MullvadREST
 import MullvadTransport
 @testable import RelayCache
 import XCTest
diff --git a/ios/MullvadVPNTests/RelaySelectorTests.swift b/ios/MullvadVPNTests/RelaySelectorTests.swift
index fe7f16e4f6..0e97d09eff 100644
--- a/ios/MullvadVPNTests/RelaySelectorTests.swift
+++ b/ios/MullvadVPNTests/RelaySelectorTests.swift
@@ -6,7 +6,7 @@
 //  Copyright © 2019 Mullvad VPN AB. All rights reserved.
 //
 
-import MullvadREST
+@testable import MullvadREST
 import MullvadTypes
 import Network
 import RelaySelector
@@ -86,6 +86,14 @@ class RelaySelectorTests: XCTestCase {
         result = try RelaySelector.evaluate(relays: sampleRelays, constraints: constraints, numberOfFailedAttempts: 4)
         XCTAssertTrue(allPorts.contains(result.endpoint.ipv4Relay.port))
     }
+
+    func testClosestShadowsocksRelay() throws {
+        let constraints = RelayConstraints(location: .only(.city("se", "sto")))
+
+        let selectedRelay = RelaySelector.closestShadowsocksRelayConstrained(by: constraints, in: sampleRelays)
+
+        XCTAssertEqual(selectedRelay?.hostname, "se-sto-br-001")
+    }
 }
 
 private let sampleRelays = REST.ServerRelaysResponse(
@@ -108,6 +116,24 @@ private let sampleRelays = REST.ServerRelaysResponse(
             latitude: 59.3289,
             longitude: 18.0649
         ),
+        "ae-dxb": REST.ServerLocation(
+            country: "United Arab Emirates",
+            city: "Dubai",
+            latitude: 25.276987,
+            longitude: 55.296249
+        ),
+        "jp-tyo": REST.ServerLocation(
+            country: "Japan",
+            city: "Tokyo",
+            latitude: 35.685,
+            longitude: 139.751389
+        ),
+        "ca-tor": REST.ServerLocation(
+            country: "Canada",
+            city: "Toronto",
+            latitude: 43.666667,
+            longitude: -79.416667
+        ),
     ],
     wireguard: REST.ServerWireguardTunnels(
         ipv4Gateway: .loopback,
@@ -164,5 +190,48 @@ private let sampleRelays = REST.ServerRelaysResponse(
             ),
         ]
     ),
-    bridge: REST.ServerBridges(shadowsocks: [], relays: [])
+    bridge: REST.ServerBridges(shadowsocks: [
+        REST.ServerShadowsocks(protocol: "tcp", port: 443, cipher: "aes-256-gcm", password: "mullvad"),
+    ], relays: [
+        REST.BridgeRelay(
+            hostname: "se-sto-br-001",
+            active: true,
+            owned: true,
+            location: "se-sto",
+            provider: "31173",
+            ipv4AddrIn: .loopback,
+            weight: 100,
+            includeInCountry: true
+        ),
+        REST.BridgeRelay(
+            hostname: "jp-tyo-br-101",
+            active: true,
+            owned: true,
+            location: "jp-tyo",
+            provider: "M247",
+            ipv4AddrIn: .loopback,
+            weight: 1,
+            includeInCountry: true
+        ),
+        REST.BridgeRelay(
+            hostname: "ca-tor-ovpn-001",
+            active: false,
+            owned: false,
+            location: "ca-tor",
+            provider: "M247",
+            ipv4AddrIn: .loopback,
+            weight: 1,
+            includeInCountry: true
+        ),
+        REST.BridgeRelay(
+            hostname: "ae-dxb-ovpn-001",
+            active: true,
+            owned: false,
+            location: "ae-dxb",
+            provider: "M247",
+            ipv4AddrIn: .loopback,
+            weight: 100,
+            includeInCountry: true
+        ),
+    ])
 )
diff --git a/ios/PacketTunnel/PacketTunnelProvider.swift b/ios/PacketTunnel/PacketTunnelProvider.swift
index ecd7770d23..9d23cba3e8 100644
--- a/ios/PacketTunnel/PacketTunnelProvider.swift
+++ b/ios/PacketTunnel/PacketTunnelProvider.swift
@@ -105,6 +105,8 @@ class PacketTunnelProvider: NEPacketTunnelProvider, TunnelMonitorDelegate {
     /// Whether to use the cached device state.
     private var useCachedDeviceState = false
 
+    private let constraintsUpdater = RelayConstraintsUpdater()
+
     /// Returns `PacketTunnelStatus` used for sharing with main bundle process.
     private var packetTunnelStatus: PacketTunnelStatus {
         let errors: [PacketTunnelErrorWrapper?] = [
@@ -147,7 +149,8 @@ class PacketTunnelProvider: NEPacketTunnelProvider, TunnelMonitorDelegate {
             urlSessionTransport: urlSessionTransport,
             relayCache: relayCache,
             addressCache: addressCache,
-            shadowsocksCache: shadowsocksCache
+            shadowsocksCache: shadowsocksCache,
+            constraintsUpdater: constraintsUpdater
         )
 
         let proxyFactory = REST.ProxyFactory.makeProxyFactory(
@@ -570,6 +573,8 @@ class PacketTunnelProvider: NEPacketTunnelProvider, TunnelMonitorDelegate {
             selectorResult = aSelectorResult
         }
 
+        constraintsUpdater.onNewConstraints?(tunnelSettings.relayConstraints)
+
         return PacketTunnelConfiguration(
             deviceState: deviceState,
             tunnelSettings: tunnelSettings,
diff --git a/ios/RelaySelector/RelaySelector.swift b/ios/RelaySelector/RelaySelector.swift
index 53c2cc8995..5643f5d73f 100644
--- a/ios/RelaySelector/RelaySelector.swift
+++ b/ios/RelaySelector/RelaySelector.swift
@@ -16,17 +16,35 @@ public enum RelaySelector {
     /**
      Returns random shadowsocks TCP bridge, otherwise `nil` if there are no shadowdsocks bridges.
      */
-    public static func getShadowsocksTCPBridge(relays: REST.ServerRelaysResponse) -> REST.ServerShadowsocks? {
+    public static func shadowsocksTCPBridge(from relays: REST.ServerRelaysResponse) -> REST.ServerShadowsocks? {
         relays.bridge.shadowsocks.filter { $0.protocol == "tcp" }.randomElement()
     }
 
     /// Return a random Shadowsocks bridge relay, or `nil` if no relay were found.
     ///
     /// Non `active` relays are filtered out.
-    /// - Parameter relays: The list of relays to randomly select from
-    /// - Returns: A Shadowsocks relay or `nil` if no relay were found.
-    public static func getShadowsocksRelay(relays: REST.ServerRelaysResponse) -> REST.BridgeRelay? {
-        relays.bridge.relays.filter { $0.active }.randomElement()
+    /// - Parameter relays: The list of relays to randomly select from.
+    /// - Returns: A Shadowsocks relay or `nil` if no active relay were found.
+    public static func shadowsocksRelay(from relaysResponse: REST.ServerRelaysResponse) -> REST.BridgeRelay? {
+        relaysResponse.bridge.relays.filter { $0.active }.randomElement()
+    }
+
+    /// Returns the closest Shadowsocks relay using the given `constraints`, or a random relay if `constraints` were
+    /// unsatisfiable.
+    ///
+    /// - Parameters:
+    ///   - constraints: The user selected `constraints`
+    ///   - relays: The list of relays to randomly select from.
+    /// - Returns: A Shadowsocks relay or `nil` if no active relay were found.
+    public static func closestShadowsocksRelayConstrained(
+        by constraints: RelayConstraints,
+        in relaysResponse: REST.ServerRelaysResponse
+    ) -> REST.BridgeRelay? {
+        let mappedBridges = mapRelays(relays: relaysResponse.bridge.relays, locations: relaysResponse.locations)
+        let filteredRelays = applyConstraints(constraints, relays: mappedBridges)
+        let randomBridgeRelay = pickRandomRelay(relays: filteredRelays)
+
+        return randomBridgeRelay?.relay ?? shadowsocksRelay(from: relaysResponse)
     }
 
     /**
@@ -38,7 +56,8 @@ public enum RelaySelector {
         constraints: RelayConstraints,
         numberOfFailedAttempts: UInt
     ) throws -> RelaySelectorResult {
-        let filteredRelays = applyConstraints(constraints, relays: Self.parseRelaysResponse(relays))
+        let mappedRelays = mapRelays(relays: relays.wireguard.relays, locations: relays.locations)
+        let filteredRelays = applyConstraints(constraints, relays: mappedRelays)
         let port = applyConstraints(
             constraints,
             rawPortRanges: relays.wireguard.portRanges,
@@ -68,10 +87,10 @@ public enum RelaySelector {
     }
 
     /// Produce a list of `RelayWithLocation` items satisfying the given constraints
-    private static func applyConstraints(
+    private static func applyConstraints<T: AnyRelay>(
         _ constraints: RelayConstraints,
-        relays: [RelayWithLocation]
-    ) -> [RelayWithLocation] {
+        relays: [RelayWithLocation<T>]
+    ) -> [RelayWithLocation<T>] {
         relays.filter { relayWithLocation -> Bool in
             switch constraints.location {
             case .any:
@@ -117,7 +136,7 @@ public enum RelaySelector {
         }
     }
 
-    private static func pickRandomRelay(relays: [RelayWithLocation]) -> RelayWithLocation? {
+    private static func pickRandomRelay<T: AnyRelay>(relays: [RelayWithLocation<T>]) -> RelayWithLocation<T>? {
         let totalWeight = relays.reduce(0) { accummulatedWeight, relayWithLocation in
             accummulatedWeight + relayWithLocation.relay.weight
         }
@@ -127,7 +146,7 @@ public enum RelaySelector {
             return relays.randomElement()
         }
 
-        // Pick a random number in the range 1 - totalWeight. This choses the relay with a
+        // Pick a random number in the range 1 - totalWeight. This chooses the relay with a
         // non-zero weight.
         var i = (1 ... totalWeight).randomElement()!
 
@@ -183,24 +202,33 @@ public enum RelaySelector {
         }
     }
 
-    private static func parseRelaysResponse(_ response: REST.ServerRelaysResponse) -> [RelayWithLocation] {
-        response.wireguard.relays.compactMap { serverRelay -> RelayWithLocation? in
-            guard let serverLocation = response.locations[serverRelay.location] else { return nil }
+    private static func mapRelays<T: AnyRelay>(
+        relays: [T],
+        locations: [String: REST.ServerLocation]
+    ) -> [RelayWithLocation<T>] {
+        relays.compactMap { relay in
+            guard let serverLocation = locations[relay.location] else { return nil }
+            return makeRelayWithLocationFrom(serverLocation, relay: relay)
+        }
+    }
 
-            let locationComponents = serverRelay.location.split(separator: "-")
-            guard locationComponents.count > 1 else { return nil }
+    private static func makeRelayWithLocationFrom<T: AnyRelay>(
+        _ serverLocation: REST.ServerLocation,
+        relay: T
+    ) -> RelayWithLocation<T>? {
+        let locationComponents = relay.location.split(separator: "-")
+        guard locationComponents.count > 1 else { return nil }
 
-            let location = Location(
-                country: serverLocation.country,
-                countryCode: String(locationComponents[0]),
-                city: serverLocation.city,
-                cityCode: String(locationComponents[1]),
-                latitude: serverLocation.latitude,
-                longitude: serverLocation.longitude
-            )
+        let location = Location(
+            country: serverLocation.country,
+            countryCode: String(locationComponents[0]),
+            city: serverLocation.city,
+            cityCode: String(locationComponents[1]),
+            latitude: serverLocation.latitude,
+            longitude: serverLocation.longitude
+        )
 
-            return RelayWithLocation(relay: serverRelay, location: location)
-        }
+        return RelayWithLocation(relay: relay, location: location)
     }
 }
 
@@ -225,7 +253,18 @@ public struct RelaySelectorResult: Codable {
     }
 }
 
-private struct RelayWithLocation {
-    var relay: REST.ServerRelay
-    var location: Location
+protocol AnyRelay {
+    var hostname: String { get }
+    var location: String { get }
+    var weight: UInt64 { get }
+    var active: Bool { get }
+    var includeInCountry: Bool { get }
+}
+
+extension REST.ServerRelay: AnyRelay {}
+extension REST.BridgeRelay: AnyRelay {}
+
+fileprivate struct RelayWithLocation<T: AnyRelay> {
+    let relay: T
+    let location: Location
 }