Allow relay selector to select an entry peer

ios/PacketTunnelCoreTests/Mocks/RelaySelectorStub.swift → ios/MullvadMockData/MullvadREST/RelaySelectorStub.swift

@@ -6,30 +6,29 @@
 //  Copyright © 2023 Mullvad VPN AB. All rights reserved.
 //
 
-import Foundation
+import MullvadREST
 import MullvadTypes
-import PacketTunnelCore
 import WireGuardKitTypes
 
 /// Relay selector stub that accepts a block that can be used to provide custom implementation.
-struct RelaySelectorStub: RelaySelectorProtocol {
-    let block: (RelayConstraints, UInt) throws -> SelectedRelay
+public struct RelaySelectorStub: RelaySelectorProtocol {
+    let block: (RelayConstraints, UInt) throws -> SelectedRelays
 
-    func selectRelay(
+    public func selectRelays(
         with constraints: RelayConstraints,
-        connectionAttemptFailureCount: UInt
-    ) throws -> SelectedRelay {
-        return try block(constraints, connectionAttemptFailureCount)
+        connectionAttemptCount: UInt
+    ) throws -> SelectedRelays {
+        return try block(constraints, connectionAttemptCount)
     }
 }
 
 extension RelaySelectorStub {
     /// Returns a relay selector that never fails.
-    static func nonFallible() -> RelaySelectorStub {
+    public static func nonFallible() -> RelaySelectorStub {
         let publicKey = PrivateKey().publicKey.rawValue
 
         return RelaySelectorStub { _, _ in
-            return SelectedRelay(
+            let cityRelay = SelectedRelay(
                 endpoint: MullvadEndpoint(
                     ipv4Relay: IPv4Endpoint(ip: .loopback, port: 1300),
                     ipv4Gateway: .loopback,
@@ -46,6 +45,11 @@ extension RelaySelectorStub {
                     longitude: 0
                 ), retryAttempts: 0
             )
+
+            return SelectedRelays(
+                entry: cityRelay,
+                exit: cityRelay
+            )
         }
     }
 }

ios/MullvadREST/Relay/MultihopDecisionFlow.swift

@@ -0,0 +1,116 @@
+//
+//  MultihopDecisionFlow.swift
+//  MullvadREST
+//
+//  Created by Jon Petersson on 2024-06-14.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import Foundation
+
+protocol MultihopDecisionFlow {
+    typealias RelayCandidate = RelayWithLocation<REST.ServerRelay>
+    init(next: MultihopDecisionFlow?, relayPicker: RelayPicking)
+    func canHandle(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) -> Bool
+    func pick(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) throws -> SelectedRelays
+}
+
+struct OneToOne: MultihopDecisionFlow {
+    let next: MultihopDecisionFlow?
+    let relayPicker: RelayPicking
+    init(next: (any MultihopDecisionFlow)?, relayPicker: RelayPicking) {
+        self.next = next
+        self.relayPicker = relayPicker
+    }
+
+    func pick(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) throws -> SelectedRelays {
+        guard canHandle(entryCandidates: entryCandidates, exitCandidates: exitCandidates) else {
+            guard let next else {
+                throw NoRelaysSatisfyingConstraintsError()
+            }
+            return try next.pick(entryCandidates: entryCandidates, exitCandidates: exitCandidates)
+        }
+
+        guard entryCandidates.first != exitCandidates.first else {
+            throw NoRelaysSatisfyingConstraintsError()
+        }
+
+        let entryMatch = try relayPicker.findBestMatch(from: entryCandidates)
+        let exitMatch = try relayPicker.findBestMatch(from: exitCandidates)
+        return SelectedRelays(entry: entryMatch, exit: exitMatch)
+    }
+
+    func canHandle(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) -> Bool {
+        entryCandidates.count == 1 && exitCandidates.count == 1
+    }
+}
+
+struct OneToMany: MultihopDecisionFlow {
+    let next: MultihopDecisionFlow?
+    let relayPicker: RelayPicking
+
+    init(next: (any MultihopDecisionFlow)?, relayPicker: RelayPicking) {
+        self.next = next
+        self.relayPicker = relayPicker
+    }
+
+    func pick(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) throws -> SelectedRelays {
+        guard let multihopPicker = relayPicker as? MultihopPicker else {
+            fatalError("Could not cast picker to MultihopPicker")
+        }
+
+        guard canHandle(entryCandidates: entryCandidates, exitCandidates: exitCandidates) else {
+            guard let next else {
+                throw NoRelaysSatisfyingConstraintsError()
+            }
+            return try next.pick(entryCandidates: entryCandidates, exitCandidates: exitCandidates)
+        }
+
+        switch (entryCandidates.count, exitCandidates.count) {
+        case let (1, count) where count > 1:
+            let entryMatch = try multihopPicker.findBestMatch(from: entryCandidates)
+            let exitMatch = try multihopPicker.exclude(relay: entryMatch, from: exitCandidates)
+            return SelectedRelays(entry: entryMatch, exit: exitMatch)
+        default:
+            let exitMatch = try multihopPicker.findBestMatch(from: exitCandidates)
+            let entryMatch = try multihopPicker.exclude(relay: exitMatch, from: entryCandidates)
+            return SelectedRelays(entry: entryMatch, exit: exitMatch)
+        }
+    }
+
+    func canHandle(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) -> Bool {
+        (entryCandidates.count == 1 && exitCandidates.count > 1) ||
+            (entryCandidates.count > 1 && exitCandidates.count == 1)
+    }
+}
+
+struct ManyToMany: MultihopDecisionFlow {
+    let next: MultihopDecisionFlow?
+    let relayPicker: RelayPicking
+
+    init(next: (any MultihopDecisionFlow)?, relayPicker: RelayPicking) {
+        self.next = next
+        self.relayPicker = relayPicker
+    }
+
+    func pick(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) throws -> SelectedRelays {
+        guard let multihopPicker = relayPicker as? MultihopPicker else {
+            fatalError("Could not cast picker to MultihopPicker")
+        }
+
+        guard canHandle(entryCandidates: entryCandidates, exitCandidates: exitCandidates) else {
+            guard let next else {
+                throw NoRelaysSatisfyingConstraintsError()
+            }
+            return try next.pick(entryCandidates: entryCandidates, exitCandidates: exitCandidates)
+        }
+
+        let exitMatch = try multihopPicker.findBestMatch(from: exitCandidates)
+        let entryMatch = try multihopPicker.exclude(relay: exitMatch, from: entryCandidates)
+        return SelectedRelays(entry: entryMatch, exit: exitMatch)
+    }
+
+    func canHandle(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) -> Bool {
+        entryCandidates.count > 1 && exitCandidates.count > 1
+    }
+}

ios/MullvadREST/Relay/NoRelaysSatisfyingConstraintsError.swift

@@ -9,6 +9,8 @@
 import Foundation
 
 public struct NoRelaysSatisfyingConstraintsError: LocalizedError {
+    public init() {}
+
     public var errorDescription: String? {
         "No relays satisfying constraints."
     }

ios/MullvadREST/Relay/RelayPicking.swift

@@ -0,0 +1,107 @@
+//
+//  RelaySelectorPicker.swift
+//  MullvadREST
+//
+//  Created by Jon Petersson on 2024-06-05.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import MullvadSettings
+import MullvadTypes
+
+protocol RelayPicking {
+    var relays: REST.ServerRelaysResponse { get }
+    var constraints: RelayConstraints { get }
+    var connectionAttemptCount: UInt { get }
+    func pick() throws -> SelectedRelays
+}
+
+extension RelayPicking {
+    func findBestMatch(
+        from candidates: [RelayWithLocation<REST.ServerRelay>]
+    ) throws -> SelectedRelay {
+        let match = try RelaySelector.WireGuard.pickCandidate(
+            from: candidates,
+            relays: relays,
+            portConstraint: constraints.port,
+            numberOfFailedAttempts: connectionAttemptCount
+        )
+
+        return SelectedRelay(
+            endpoint: match.endpoint,
+            hostname: match.relay.hostname,
+            location: match.location,
+            retryAttempts: connectionAttemptCount
+        )
+    }
+}
+
+struct SinglehopPicker: RelayPicking {
+    let constraints: RelayConstraints
+    let relays: REST.ServerRelaysResponse
+    let connectionAttemptCount: UInt
+
+    func pick() throws -> SelectedRelays {
+        let candidates = try RelaySelector.WireGuard.findCandidates(
+            by: constraints.exitLocations,
+            in: relays,
+            filterConstraint: constraints.filter
+        )
+
+        let match = try findBestMatch(from: candidates)
+
+        return SelectedRelays(entry: nil, exit: match)
+    }
+}
+
+struct MultihopPicker: RelayPicking {
+    let constraints: RelayConstraints
+    let relays: REST.ServerRelaysResponse
+    let connectionAttemptCount: UInt
+
+    func pick() throws -> SelectedRelays {
+        let entryCandidates = try RelaySelector.WireGuard.findCandidates(
+            by: constraints.entryLocations,
+            in: relays,
+            filterConstraint: constraints.filter
+        )
+
+        let exitCandidates = try RelaySelector.WireGuard.findCandidates(
+            by: constraints.exitLocations,
+            in: relays,
+            filterConstraint: constraints.filter
+        )
+
+        /*
+         Relay selection is prioritised in the following order:
+         1. Both entry and exit constraints match only a single relay. Both relays are selected.
+         2. Either entry or exit constraint matches only a single relay and the other multiple relays. The single relays
+         is selected and excluded from the list of multiple relays.
+         3. Both entry and exit constraints match multiple relays. Exit relay is picked first and then excluded from
+         the list of entry relays.
+         */
+        let decisionFlow = OneToOne(
+            next: OneToMany(
+                next: ManyToMany(
+                    next: nil,
+                    relayPicker: self
+                ),
+                relayPicker: self
+            ),
+            relayPicker: self
+        )
+
+        return try decisionFlow.pick(entryCandidates: entryCandidates, exitCandidates: exitCandidates)
+    }
+
+    func exclude(
+        relay: SelectedRelay,
+        from candidates: [RelayWithLocation<REST.ServerRelay>]
+    ) throws -> SelectedRelay {
+        let filteredCandidates = candidates.filter { relayWithLocation in
+            relayWithLocation.relay.hostname != relay.hostname
+        }
+
+        return try findBestMatch(from: filteredCandidates)
+    }
+}

ios/MullvadREST/Relay/RelaySelector+Shadowsocks.swift

@@ -45,7 +45,6 @@ extension RelaySelector {
             let mappedBridges = mapRelays(relays: relaysResponse.bridge.relays, locations: relaysResponse.locations)
             let filteredRelays = applyConstraints(
                 location,
-                portConstraint: port,
                 filterConstraint: filter,
                 relays: mappedBridges
             )

ios/MullvadREST/Relay/RelaySelector+Wireguard.swift

@@ -6,54 +6,39 @@
 //  Copyright © 2024 Mullvad VPN AB. All rights reserved.
 //
 
-import Foundation
 import MullvadTypes
 
 extension RelaySelector {
     public enum WireGuard {
-        /**
-         Filters relay list using given constraints and selects random relay for exit relay.
-         Throws an error if there are no relays satisfying the given constraints.
-         */
-        public static func evaluate(
-            by constraints: RelayConstraints,
-            in relaysResponse: REST.ServerRelaysResponse,
-            numberOfFailedAttempts: UInt
-        ) throws -> RelaySelectorResult {
-            let exitCandidates = try findBestMatch(
-                relays: relaysResponse,
-                relayConstraint: constraints.exitLocations,
-                portConstraint: constraints.port,
-                filterConstraint: constraints.filter,
-                numberOfFailedAttempts: numberOfFailedAttempts
-            )
+        /// Filters relay list using given constraints.
+        public static func findCandidates(
+            by relayConstraint: RelayConstraint<UserSelectedRelays>,
+            in relays: REST.ServerRelaysResponse,
+            filterConstraint: RelayConstraint<RelayFilter>
+        ) throws -> [RelayWithLocation<REST.ServerRelay>] {
+            let mappedRelays = mapRelays(relays: relays.wireguard.relays, locations: relays.locations)
 
-            return exitCandidates
+            return applyConstraints(
+                relayConstraint,
+                filterConstraint: filterConstraint,
+                relays: mappedRelays
+            )
         }
 
-        // MARK: - private functions
-
-        private static func findBestMatch(
+        /// Picks a random relay from a list.
+        public static func pickCandidate(
+            from relayWithLocations: [RelayWithLocation<REST.ServerRelay>],
             relays: REST.ServerRelaysResponse,
-            relayConstraint: RelayConstraint<UserSelectedRelays>,
             portConstraint: RelayConstraint<UInt16>,
-            filterConstraint: RelayConstraint<RelayFilter>,
             numberOfFailedAttempts: UInt
         ) throws -> RelaySelectorMatch {
-            let mappedRelays = mapRelays(relays: relays.wireguard.relays, locations: relays.locations)
-            let filteredRelays = applyConstraints(
-                relayConstraint,
-                portConstraint: portConstraint,
-                filterConstraint: filterConstraint,
-                relays: mappedRelays
-            )
             let port = applyPortConstraint(
                 portConstraint,
                 rawPortRanges: relays.wireguard.portRanges,
                 numberOfFailedAttempts: numberOfFailedAttempts
             )
 
-            guard let relayWithLocation = pickRandomRelayByWeight(relays: filteredRelays), let port else {
+            guard let port, let relayWithLocation = pickRandomRelayByWeight(relays: relayWithLocations) else {
                 throw NoRelaysSatisfyingConstraintsError()
             }
 

ios/MullvadREST/Relay/RelaySelector.swift

@@ -134,7 +134,6 @@ public enum RelaySelector {
     /// Produce a list of `RelayWithLocation` items satisfying the given constraints
     static func applyConstraints<T: AnyRelay>(
         _ relayConstraint: RelayConstraint<UserSelectedRelays>,
-        portConstraint: RelayConstraint<UInt16>,
         filterConstraint: RelayConstraint<RelayFilter>,
         relays: [RelayWithLocation<T>]
     ) -> [RelayWithLocation<T>] {