Fix relay selector for smart routing

2 files changed, 34 insertions, 10 deletions

diff --git a/ios/MullvadREST/Relay/RelayPicking.swift b/ios/MullvadREST/Relay/RelayPicking.swift
index 4d5c98975b..3e94b42d52 100644
--- a/ios/MullvadREST/Relay/RelayPicking.swift
+++ b/ios/MullvadREST/Relay/RelayPicking.swift
@@ -46,6 +46,18 @@ struct SinglehopPicker: RelayPicking {
 
     func pick() throws -> SelectedRelays {
         do {
+            let exitCandidates = try RelaySelector.WireGuard.findCandidates(
+                by: constraints.exitLocations,
+                in: relays,
+                filterConstraint: constraints.filter,
+                daitaEnabled: daitaSettings.daitaState.isEnabled
+            )
+
+            let match = try findBestMatch(from: exitCandidates)
+            return SelectedRelays(entry: nil, exit: match, retryAttempt: connectionAttemptCount)
+        } catch let error as NoRelaysSatisfyingConstraintsError where error.reason == .noDaitaRelaysFound {
+            // If DAITA is on and Direct only is off, and no supported relays are found, we should try to find the nearest
+            // available relay that supports DAITA and use it as entry in a multihop selection.
             if daitaSettings.isAutomaticRouting {
                 return try MultihopPicker(
                     relays: relays,
@@ -54,15 +66,7 @@ struct SinglehopPicker: RelayPicking {
                     daitaSettings: daitaSettings
                 ).pick()
             } else {
-                let exitCandidates = try RelaySelector.WireGuard.findCandidates(
-                    by: constraints.exitLocations,
-                    in: relays,
-                    filterConstraint: constraints.filter,
-                    daitaEnabled: daitaSettings.daitaState.isEnabled
-                )
-
-                let match = try findBestMatch(from: exitCandidates)
-                return SelectedRelays(entry: nil, exit: match, retryAttempt: connectionAttemptCount)
+                throw error
             }
         }
     }
diff --git a/ios/MullvadVPNTests/MullvadREST/Relay/RelayPickingTests.swift b/ios/MullvadVPNTests/MullvadREST/Relay/RelayPickingTests.swift
index e65a1c5a17..41b65e3946 100644
--- a/ios/MullvadVPNTests/MullvadREST/Relay/RelayPickingTests.swift
+++ b/ios/MullvadVPNTests/MullvadREST/Relay/RelayPickingTests.swift
@@ -16,6 +16,8 @@ import XCTest
 class RelayPickingTests: XCTestCase {
     let sampleRelays = ServerRelaysResponseStubs.sampleRelays
 
+    // MARK: Single-/multihop
+
     func testSinglehopPicker() throws {
         let constraints = RelayConstraints(
             entryLocations: .only(UserSelectedRelays(locations: [.hostname("se", "sto", "se2-wireguard")])),
@@ -75,6 +77,10 @@ class RelayPickingTests: XCTestCase {
         }
     }
 
+    // MARK: DAITA/Direct only
+
+    // DAITA - ON, Direct only - OFF, Multihop - OFF, Exit supports DAITA - FALSE
+    // Direct only is off, so we should automatically pick the entry that is closest to exit.
     func testDirectOnlyOffDaitaOnForSinglehopWithoutDaitaRelay() throws {
         let constraints = RelayConstraints(
             exitLocations: .only(UserSelectedRelays(locations: [.hostname("se", "got", "se10-wireguard")]))
@@ -93,6 +99,8 @@ class RelayPickingTests: XCTestCase {
         XCTAssertEqual(selectedRelays.exit.hostname, "se10-wireguard")
     }
 
+    // DAITA - ON, Direct only - ON, Multihop - OFF, Exit supports DAITA - FALSE
+    // Go into blocked state since Direct only requires a DAITA entry.
     func testDirectOnlyOnDaitaOnForSinglehopWithoutDaitaRelay() throws {
         let constraints = RelayConstraints(
             exitLocations: .only(UserSelectedRelays(locations: [.hostname("se", "got", "se10-wireguard")]))
@@ -108,6 +116,8 @@ class RelayPickingTests: XCTestCase {
         XCTAssertThrowsError(try picker.pick())
     }
 
+    // DAITA - ON, Direct only - OFF, Multihop - OFF, Exit supports DAITA - TRUE
+    // Select the DAITA entry, no automatic routing needed.
     func testDirectOnlyOffDaitaOnForSinglehopWithDaitaRelay() throws {
         let constraints = RelayConstraints(
             exitLocations: .only(UserSelectedRelays(locations: [.hostname("es", "mad", "es1-wireguard")]))
@@ -122,10 +132,12 @@ class RelayPickingTests: XCTestCase {
 
         let selectedRelays = try picker.pick()
 
-        XCTAssertEqual(selectedRelays.entry?.hostname, "us-nyc-wg-301") // New York relay is closest to exit relay.
+        XCTAssertNil(selectedRelays.entry)
         XCTAssertEqual(selectedRelays.exit.hostname, "es1-wireguard")
     }
 
+    // DAITA - ON, Direct only - ON, Multihop - OFF, Exit supports DAITA - TRUE
+    // Select the DAITA entry.
     func testDirectOnlyOnDaitaOnForSinglehopWithDaitaRelay() throws {
         let constraints = RelayConstraints(
             exitLocations: .only(UserSelectedRelays(locations: [.hostname("es", "mad", "es1-wireguard")]))
@@ -144,6 +156,9 @@ class RelayPickingTests: XCTestCase {
         XCTAssertEqual(selectedRelays.exit.hostname, "es1-wireguard")
     }
 
+    // DAITA - ON, Direct only - OFF, Multihop - ON, Entry supports DAITA - TRUE
+    // Direct only is off, so we should automatically pick the entry that is closest to exit, ignoring
+    // selected multihop entry.
     func testDirectOnlyOffDaitaOnForMultihopWithDaitaRelay() throws {
         let constraints = RelayConstraints(
             entryLocations: .only(UserSelectedRelays(locations: [.hostname("us", "nyc", "us-nyc-wg-301")])),
@@ -163,6 +178,9 @@ class RelayPickingTests: XCTestCase {
         XCTAssertEqual(selectedRelays.exit.hostname, "se10-wireguard")
     }
 
+    // DAITA - ON, Direct only - OFF, Multihop - ON, Entry supports DAITA - FALSE
+    // Direct only is off, so we should automatically pick the entry that is closest to exit, ignoring
+    // selected multihop entry.
     func testDirectOnlyOffDaitaOnForMultihopWithoutDaitaRelay() throws {
         let constraints = RelayConstraints(
             entryLocations: .only(UserSelectedRelays(locations: [.hostname("se", "got", "se10-wireguard")])),
@@ -182,6 +200,8 @@ class RelayPickingTests: XCTestCase {
         XCTAssertEqual(selectedRelays.exit.hostname, "se10-wireguard")
     }
 
+    // DAITA - ON, Direct only - ON, Multihop - ON, Entry supports DAITA - FALSE
+    // Go into blocked state since Direct only requires a DAITA entry.
     func testDirectOnlyOnDaitaOnForMultihopWithoutDaitaRelay() throws {
         let constraints = RelayConstraints(
             entryLocations: .only(UserSelectedRelays(locations: [.hostname("se", "got", "se10-wireguard")])),