Add RelaySelectorWrapper tests

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/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)
+    }
+}