Smart routing should select the closest relay with Daitaclosest-daita-relay

8 files changed, 153 insertions, 68 deletions

diff --git a/ios/MullvadREST/Relay/RelayPicking.swift b/ios/MullvadREST/Relay/RelayPicking.swift
index 15ae2ee90c..1c54953ef1 100644
--- a/ios/MullvadREST/Relay/RelayPicking.swift
+++ b/ios/MullvadREST/Relay/RelayPicking.swift
@@ -13,18 +13,18 @@ protocol RelayPicking {
     var relays: REST.ServerRelaysResponse { get }
     var constraints: RelayConstraints { get }
     var connectionAttemptCount: UInt { get }
+    var preferClosest: Bool { get }
     func pick() throws -> SelectedRelays
 }
 
 extension RelayPicking {
-    func findBestMatch(
-        from candidates: [RelayWithLocation<REST.ServerRelay>]
-    ) throws -> SelectedRelay {
+    func findBestMatch(from candidates: [RelayWithLocation<REST.ServerRelay>]) throws -> SelectedRelay {
         let match = try RelaySelector.WireGuard.pickCandidate(
             from: candidates,
             relays: relays,
             portConstraint: constraints.port,
-            numberOfFailedAttempts: connectionAttemptCount
+            numberOfFailedAttempts: connectionAttemptCount,
+            preferClosest: preferClosest
         )
 
         return SelectedRelay(
@@ -40,6 +40,21 @@ struct SinglehopPicker: RelayPicking {
     let daitaSettings: DAITASettings
     let relays: REST.ServerRelaysResponse
     let connectionAttemptCount: UInt
+    let preferClosest: Bool
+
+    init(
+        constraints: RelayConstraints,
+        daitaSettings: DAITASettings,
+        relays: REST.ServerRelaysResponse,
+        connectionAttemptCount: UInt,
+        preferClosest: Bool = false
+    ) {
+        self.constraints = constraints
+        self.daitaSettings = daitaSettings
+        self.relays = relays
+        self.connectionAttemptCount = connectionAttemptCount
+        self.preferClosest = preferClosest
+    }
 
     func pick() throws -> SelectedRelays {
         var exitCandidates = [RelayWithLocation<REST.ServerRelay>]()
@@ -62,7 +77,8 @@ struct SinglehopPicker: RelayPicking {
                 constraints: constraints,
                 daitaSettings: daitaSettings,
                 relays: relays,
-                connectionAttemptCount: connectionAttemptCount
+                connectionAttemptCount: connectionAttemptCount,
+                preferClosest: true
             ).pick()
             #endif
         }
@@ -77,6 +93,21 @@ struct MultihopPicker: RelayPicking {
     let daitaSettings: DAITASettings
     let relays: REST.ServerRelaysResponse
     let connectionAttemptCount: UInt
+    let preferClosest: Bool
+
+    init(
+        constraints: RelayConstraints,
+        daitaSettings: DAITASettings,
+        relays: REST.ServerRelaysResponse,
+        connectionAttemptCount: UInt,
+        preferClosest: Bool = false
+    ) {
+        self.constraints = constraints
+        self.daitaSettings = daitaSettings
+        self.relays = relays
+        self.connectionAttemptCount = connectionAttemptCount
+        self.preferClosest = preferClosest
+    }
 
     func pick() throws -> SelectedRelays {
         let entryCandidates = try RelaySelector.WireGuard.findCandidates(
diff --git a/ios/MullvadREST/Relay/RelaySelector+Shadowsocks.swift b/ios/MullvadREST/Relay/RelaySelector+Shadowsocks.swift
index f529b9b924..bbc97c170f 100644
--- a/ios/MullvadREST/Relay/RelaySelector+Shadowsocks.swift
+++ b/ios/MullvadREST/Relay/RelaySelector+Shadowsocks.swift
@@ -36,7 +36,7 @@ extension RelaySelector {
         ///   - filter: The user filtered criteria
         ///   - relays: The list of relays to randomly select from.
         /// - Returns: A Shadowsocks relay or `nil` if no active relay were found.
-        public static func closestRelay(
+        public static func closestBridge(
             location: RelayConstraint<UserSelectedRelays>,
             port: RelayConstraint<UInt16>,
             filter: RelayConstraint<RelayFilter>,
@@ -49,43 +49,10 @@ extension RelaySelector {
                 daitaEnabled: false,
                 relays: mappedBridges
             )) ?? []
-            guard filteredRelays.isEmpty == false else { return relay(from: relaysResponse) }
-
-            // Compute the midpoint location from all the filtered relays
-            // Take *either* the first five relays, OR the relays below maximum bridge distance
-            // sort all of them by Haversine distance from the computed midpoint location
-            // then use the roulette selection to pick a bridge
-
-            let midpointDistance = Midpoint.location(in: filteredRelays.map { $0.serverLocation.geoCoordinate })
-            let maximumBridgeDistance = 1500.0
-            let relaysWithDistance = filteredRelays.map {
-                RelayWithDistance(
-                    relay: $0.relay,
-                    distance: Haversine.distance(
-                        midpointDistance.latitude,
-                        midpointDistance.longitude,
-                        $0.serverLocation.latitude,
-                        $0.serverLocation.longitude
-                    )
-                )
-            }.sorted {
-                $0.distance < $1.distance
-            }.filter {
-                $0.distance <= maximumBridgeDistance
-            }.prefix(5)
 
-            var greatestDistance = 0.0
-            relaysWithDistance.forEach {
-                if $0.distance > greatestDistance {
-                    greatestDistance = $0.distance
-                }
-            }
-
-            let randomRelay = rouletteSelection(relays: Array(relaysWithDistance), weightFunction: { relay in
-                UInt64(1 + greatestDistance - relay.distance)
-            })
+            guard filteredRelays.isEmpty == false else { return relay(from: relaysResponse) }
 
-            return randomRelay?.relay ?? filteredRelays.randomElement()?.relay
+            return closestRelay(from: filteredRelays) ?? filteredRelays.randomElement()?.relay
         }
     }
 }
diff --git a/ios/MullvadREST/Relay/RelaySelector+Wireguard.swift b/ios/MullvadREST/Relay/RelaySelector+Wireguard.swift
index 4c8561f38b..4c517a3b48 100644
--- a/ios/MullvadREST/Relay/RelaySelector+Wireguard.swift
+++ b/ios/MullvadREST/Relay/RelaySelector+Wireguard.swift
@@ -33,38 +33,69 @@ extension RelaySelector {
             from relayWithLocations: [RelayWithLocation<REST.ServerRelay>],
             relays: REST.ServerRelaysResponse,
             portConstraint: RelayConstraint<UInt16>,
-            numberOfFailedAttempts: UInt
+            numberOfFailedAttempts: UInt,
+            preferClosest: Bool = false
         ) throws -> RelaySelectorMatch {
-            let port = applyPortConstraint(
-                portConstraint,
-                rawPortRanges: relays.wireguard.portRanges,
+            let port = try evaluatePort(
+                relays: relays,
+                portConstraint: portConstraint,
                 numberOfFailedAttempts: numberOfFailedAttempts
             )
 
-            guard let port else {
-                throw NoRelaysSatisfyingConstraintsError(.invalidPort)
+            var relayWithLocation: RelayWithLocation<REST.ServerRelay>?
+            if preferClosest {
+                let relay = closestRelay(from: relayWithLocations)
+                relayWithLocation = relayWithLocations.first(where: { $0.relay == relay })
             }
 
-            guard let relayWithLocation = pickRandomRelayByWeight(relays: relayWithLocations) else {
+            guard
+                let relayWithLocation = relayWithLocation ?? pickRandomRelayByWeight(relays: relayWithLocations)
+            else {
                 throw NoRelaysSatisfyingConstraintsError(.relayConstraintNotMatching)
             }
 
-            let endpoint = MullvadEndpoint(
-                ipv4Relay: IPv4Endpoint(
-                    ip: relayWithLocation.relay.ipv4AddrIn,
-                    port: port
-                ),
-                ipv6Relay: nil,
-                ipv4Gateway: relays.wireguard.ipv4Gateway,
-                ipv6Gateway: relays.wireguard.ipv6Gateway,
-                publicKey: relayWithLocation.relay.publicKey
-            )
+            return createMatch(for: relayWithLocation, port: port, relays: relays)
+        }
+    }
 
-            return RelaySelectorMatch(
-                endpoint: endpoint,
-                relay: relayWithLocation.relay,
-                location: relayWithLocation.serverLocation
-            )
+    private static func evaluatePort(
+        relays: REST.ServerRelaysResponse,
+        portConstraint: RelayConstraint<UInt16>,
+        numberOfFailedAttempts: UInt
+    ) throws -> UInt16 {
+        let port = applyPortConstraint(
+            portConstraint,
+            rawPortRanges: relays.wireguard.portRanges,
+            numberOfFailedAttempts: numberOfFailedAttempts
+        )
+
+        guard let port else {
+            throw NoRelaysSatisfyingConstraintsError(.invalidPort)
         }
+
+        return port
+    }
+
+    private static func createMatch(
+        for relayWithLocation: RelayWithLocation<REST.ServerRelay>,
+        port: UInt16,
+        relays: REST.ServerRelaysResponse
+    ) -> RelaySelectorMatch {
+        let endpoint = MullvadEndpoint(
+            ipv4Relay: IPv4Endpoint(
+                ip: relayWithLocation.relay.ipv4AddrIn,
+                port: port
+            ),
+            ipv6Relay: nil,
+            ipv4Gateway: relays.wireguard.ipv4Gateway,
+            ipv6Gateway: relays.wireguard.ipv6Gateway,
+            publicKey: relayWithLocation.relay.publicKey
+        )
+
+        return RelaySelectorMatch(
+            endpoint: endpoint,
+            relay: relayWithLocation.relay,
+            location: relayWithLocation.serverLocation
+        )
     }
 }
diff --git a/ios/MullvadREST/Relay/RelaySelector.swift b/ios/MullvadREST/Relay/RelaySelector.swift
index 98279daca5..41803505ea 100644
--- a/ios/MullvadREST/Relay/RelaySelector.swift
+++ b/ios/MullvadREST/Relay/RelaySelector.swift
@@ -32,6 +32,44 @@ public enum RelaySelector {
 
     // MARK: - private
 
+    /// Computes the midpoint location from all the filtered relays.
+    ///
+    /// Take *either* the first five relays, OR the relays below maximum relay distance,
+    /// sort all of them by Haversine distance from the computed midpoint location,
+    /// then use the roulette selection to pick a bridge.
+    static func closestRelay<T: AnyRelay>(from relayWithLocations: [RelayWithLocation<T>]) -> T? {
+        let midpointDistance = Midpoint.location(in: relayWithLocations.map { $0.serverLocation.geoCoordinate })
+        let maximumRelayDistance = 1500.0
+        let relaysWithDistance = relayWithLocations.map {
+            RelayWithDistance(
+                relay: $0.relay,
+                distance: Haversine.distance(
+                    midpointDistance.latitude,
+                    midpointDistance.longitude,
+                    $0.serverLocation.latitude,
+                    $0.serverLocation.longitude
+                )
+            )
+        }.sorted {
+            $0.distance < $1.distance
+        }.filter {
+            $0.distance <= maximumRelayDistance
+        }.prefix(5)
+
+        var greatestDistance = 0.0
+        relaysWithDistance.forEach {
+            if $0.distance > greatestDistance {
+                greatestDistance = $0.distance
+            }
+        }
+
+        let closestRelay = rouletteSelection(relays: Array(relaysWithDistance), weightFunction: { relay in
+            UInt64(1 + greatestDistance - relay.distance)
+        })
+
+        return closestRelay?.relay
+    }
+
     static func pickRandomRelayByWeight<T: AnyRelay>(relays: [RelayWithLocation<T>])
         -> RelayWithLocation<T>? {
         rouletteSelection(relays: relays, weightFunction: { relayWithLocation in relayWithLocation.relay.weight })
diff --git a/ios/MullvadREST/Transport/Shadowsocks/ShadowsocksRelaySelector.swift b/ios/MullvadREST/Transport/Shadowsocks/ShadowsocksRelaySelector.swift
index c0d83bc701..f5dd4e433e 100644
--- a/ios/MullvadREST/Transport/Shadowsocks/ShadowsocksRelaySelector.swift
+++ b/ios/MullvadREST/Transport/Shadowsocks/ShadowsocksRelaySelector.swift
@@ -33,7 +33,7 @@ final public class ShadowsocksRelaySelector: ShadowsocksRelaySelectorProtocol {
         case .off: settings.relayConstraints.exitLocations
         }
 
-        return RelaySelector.Shadowsocks.closestRelay(
+        return RelaySelector.Shadowsocks.closestBridge(
             location: locationConstraint,
             port: settings.relayConstraints.port,
             filter: settings.relayConstraints.filter,
diff --git a/ios/MullvadVPNTests/MullvadREST/Relay/RelayPickingTests.swift b/ios/MullvadVPNTests/MullvadREST/Relay/RelayPickingTests.swift
index 9229d4ab8f..29f227e589 100644
--- a/ios/MullvadVPNTests/MullvadREST/Relay/RelayPickingTests.swift
+++ b/ios/MullvadVPNTests/MullvadREST/Relay/RelayPickingTests.swift
@@ -74,4 +74,22 @@ class RelayPickingTests: XCTestCase {
             XCTAssertEqual(error?.reason, .entryEqualsExit)
         }
     }
+
+    func testPickingClosestRelay() throws {
+        let constraints = RelayConstraints(
+            exitLocations: .only(UserSelectedRelays(locations: [.country("se")]))
+        )
+
+        let picker = SinglehopPicker(
+            constraints: constraints,
+            daitaSettings: DAITASettings(state: .off),
+            relays: sampleRelays,
+            connectionAttemptCount: 0,
+            preferClosest: true
+        )
+
+        let selectedRelays = try picker.pick()
+print(selectedRelays.exit.hostname)
+        XCTAssertTrue(["se6-wireguard", "se10-wireguard"].contains(selectedRelays.exit.hostname))
+    }
 }
diff --git a/ios/MullvadVPNTests/MullvadREST/Relay/RelaySelectorTests.swift b/ios/MullvadVPNTests/MullvadREST/Relay/RelaySelectorTests.swift
index 5b2a775a4f..3db93f3d53 100644
--- a/ios/MullvadVPNTests/MullvadREST/Relay/RelaySelectorTests.swift
+++ b/ios/MullvadVPNTests/MullvadREST/Relay/RelaySelectorTests.swift
@@ -136,12 +136,12 @@ class RelaySelectorTests: XCTestCase {
         XCTAssertTrue(allPorts.contains(result.endpoint.ipv4Relay.port))
     }
 
-    func testClosestShadowsocksRelay() throws {
+    func testClosestRelay() throws {
         let constraints = RelayConstraints(
             exitLocations: .only(UserSelectedRelays(locations: [.city("se", "sto")]))
         )
 
-        let selectedRelay = RelaySelector.Shadowsocks.closestRelay(
+        let selectedRelay = RelaySelector.Shadowsocks.closestBridge(
             location: constraints.exitLocations,
             port: constraints.port,
             filter: constraints.filter,
@@ -156,7 +156,7 @@ class RelaySelectorTests: XCTestCase {
             exitLocations: .only(UserSelectedRelays(locations: [.country("INVALID COUNTRY")]))
         )
 
-        let selectedRelay = try XCTUnwrap(RelaySelector.Shadowsocks.closestRelay(
+        let selectedRelay = try XCTUnwrap(RelaySelector.Shadowsocks.closestBridge(
             location: constraints.exitLocations,
             port: constraints.port,
             filter: constraints.filter,
diff --git a/ios/MullvadVPNTests/MullvadREST/Shadowsocks/ShadowsocksLoaderTests.swift b/ios/MullvadVPNTests/MullvadREST/Shadowsocks/ShadowsocksLoaderTests.swift
index c8cfc4308f..3654de5a6a 100644
--- a/ios/MullvadVPNTests/MullvadREST/Shadowsocks/ShadowsocksLoaderTests.swift
+++ b/ios/MullvadVPNTests/MullvadREST/Shadowsocks/ShadowsocksLoaderTests.swift
@@ -82,7 +82,7 @@ class ShadowsocksLoaderTests: XCTestCase {
         filter: RelayConstraint<RelayFilter>,
         in: REST.ServerRelaysResponse
     ) -> REST.BridgeRelay? {
-        RelaySelector.Shadowsocks.closestRelay(
+        RelaySelector.Shadowsocks.closestBridge(
             location: location,
             port: port,
             filter: filter,