Allow relay selector to filter DAITA enabled relays

mullvad · Aug 15, 2024 · adf5fe4 · adf5fe4
1 parent 229dc5a
commit adf5fe4
Show file tree

Hide file tree

Showing 27 changed files with 447 additions and 94 deletions.
diff --git a/ios/MullvadMockData/MullvadREST/RelaySelectorStub.swift b/ios/MullvadMockData/MullvadREST/RelaySelectorStub.swift
@@ -58,7 +58,7 @@ extension RelaySelectorStub {
     /// Returns a relay selector that cannot satisfy constraints .
     public static func unsatisfied() -> RelaySelectorStub {
         return RelaySelectorStub { _ in
-            throw NoRelaysSatisfyingConstraintsError()
+            throw NoRelaysSatisfyingConstraintsError(.relayConstraintNotMatching)
         }
     }
 }
diff --git a/ios/MullvadREST/ApiHandlers/ServerRelaysResponse.swift b/ios/MullvadREST/ApiHandlers/ServerRelaysResponse.swift
@@ -34,6 +34,7 @@ extension REST {
         public let ipv4AddrIn: IPv4Address
         public let weight: UInt64
         public let includeInCountry: Bool
+        public var daita: Bool?
 
         public func override(ipv4AddrIn: IPv4Address?) -> Self {
             return BridgeRelay(
@@ -60,6 +61,7 @@ extension REST {
         public let ipv6AddrIn: IPv6Address
         public let publicKey: Data
         public let includeInCountry: Bool
+        public let daita: Bool?
 
         public func override(ipv4AddrIn: IPv4Address?, ipv6AddrIn: IPv6Address?) -> Self {
             return ServerRelay(
@@ -72,7 +74,8 @@ extension REST {
                 ipv4AddrIn: ipv4AddrIn ?? self.ipv4AddrIn,
                 ipv6AddrIn: ipv6AddrIn ?? self.ipv6AddrIn,
                 publicKey: publicKey,
-                includeInCountry: includeInCountry
+                includeInCountry: includeInCountry,
+                daita: daita
             )
         }
     }

diff --git a/ios/MullvadREST/Relay/AnyRelay.swift b/ios/MullvadREST/Relay/AnyRelay.swift
@@ -17,6 +17,7 @@ public protocol AnyRelay {
     var weight: UInt64 { get }
     var active: Bool { get }
     var includeInCountry: Bool { get }
+    var daita: Bool? { get }
 
     func override(ipv4AddrIn: IPv4Address?, ipv6AddrIn: IPv6Address?) -> Self
 }

diff --git a/ios/MullvadREST/Relay/MultihopDecisionFlow.swift b/ios/MullvadREST/Relay/MultihopDecisionFlow.swift
@@ -26,13 +26,13 @@ struct OneToOne: MultihopDecisionFlow {
     func pick(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) throws -> SelectedRelays {
         guard canHandle(entryCandidates: entryCandidates, exitCandidates: exitCandidates) else {
             guard let next else {
-                throw NoRelaysSatisfyingConstraintsError()
+                throw NoRelaysSatisfyingConstraintsError(.multihopOther)
             }
             return try next.pick(entryCandidates: entryCandidates, exitCandidates: exitCandidates)
         }
 
         guard entryCandidates.first != exitCandidates.first else {
-            throw NoRelaysSatisfyingConstraintsError()
+            throw NoRelaysSatisfyingConstraintsError(.multihopEntryEqualsExit)
         }
 
         let entryMatch = try relayPicker.findBestMatch(from: entryCandidates)
@@ -61,7 +61,7 @@ struct OneToMany: MultihopDecisionFlow {
 
         guard canHandle(entryCandidates: entryCandidates, exitCandidates: exitCandidates) else {
             guard let next else {
-                throw NoRelaysSatisfyingConstraintsError()
+                throw NoRelaysSatisfyingConstraintsError(.multihopOther)
             }
             return try next.pick(entryCandidates: entryCandidates, exitCandidates: exitCandidates)
         }
@@ -100,7 +100,7 @@ struct ManyToMany: MultihopDecisionFlow {
 
         guard canHandle(entryCandidates: entryCandidates, exitCandidates: exitCandidates) else {
             guard let next else {
-                throw NoRelaysSatisfyingConstraintsError()
+                throw NoRelaysSatisfyingConstraintsError(.multihopOther)
             }
             return try next.pick(entryCandidates: entryCandidates, exitCandidates: exitCandidates)
         }

diff --git a/ios/MullvadREST/Relay/NoRelaysSatisfyingConstraintsError.swift b/ios/MullvadREST/Relay/NoRelaysSatisfyingConstraintsError.swift
@@ -8,10 +8,24 @@
 
 import Foundation
 
+public enum NoRelaysSatisfyingConstraintsReason {
+    case filterConstraintNotMatching
+    case invalidPort
+    case multihopEntryEqualsExit
+    case multihopOther
+    case noActiveRelaysFound
+    case noDaitaRelaysFound
+    case relayConstraintNotMatching
+}
+
 public struct NoRelaysSatisfyingConstraintsError: LocalizedError {
-    public init() {}
+    public let reason: NoRelaysSatisfyingConstraintsReason
 
     public var errorDescription: String? {
         "No relays satisfying constraints."
     }
+
+    public init(_ reason: NoRelaysSatisfyingConstraintsReason) {
+        self.reason = reason
+    }
 }
diff --git a/ios/MullvadREST/Relay/RelayPicking.swift b/ios/MullvadREST/Relay/RelayPicking.swift
@@ -37,38 +37,60 @@ extension RelayPicking {
 
 struct SinglehopPicker: RelayPicking {
     let constraints: RelayConstraints
+    let daitaSettings: DAITASettings
     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
-        )
+        var exitCandidates = [RelayWithLocation<REST.ServerRelay>]()
 
-        let match = try findBestMatch(from: candidates)
+        do {
+            exitCandidates = try RelaySelector.WireGuard.findCandidates(
+                by: constraints.exitLocations,
+                in: relays,
+                filterConstraint: constraints.filter,
+                daitaEnabled: daitaSettings.state.isEnabled
+            )
+        } catch let error as NoRelaysSatisfyingConstraintsError where error.reason == .noDaitaRelaysFound {
+            #if DEBUG
+            // If DAITA is enabled 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.
+            var constraints = constraints
+            constraints.entryLocations = .any
+
+            return try MultihopPicker(
+                constraints: constraints,
+                daitaSettings: daitaSettings,
+                relays: relays,
+                connectionAttemptCount: connectionAttemptCount
+            ).pick()
+            #endif
+        }
 
+        let match = try findBestMatch(from: exitCandidates)
         return SelectedRelays(entry: nil, exit: match, retryAttempt: connectionAttemptCount)
     }
 }
 
 struct MultihopPicker: RelayPicking {
     let constraints: RelayConstraints
+    let daitaSettings: DAITASettings
     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
+            filterConstraint: constraints.filter,
+            daitaEnabled: daitaSettings.state.isEnabled
         )
 
         let exitCandidates = try RelaySelector.WireGuard.findCandidates(
             by: constraints.exitLocations,
             in: relays,
-            filterConstraint: constraints.filter
+            filterConstraint: constraints.filter,
+            daitaEnabled: false
         )
 
         /*

diff --git a/ios/MullvadREST/Relay/RelaySelector+Shadowsocks.swift b/ios/MullvadREST/Relay/RelaySelector+Shadowsocks.swift
@@ -43,11 +43,12 @@ extension RelaySelector {
             in relaysResponse: REST.ServerRelaysResponse
         ) -> REST.BridgeRelay? {
             let mappedBridges = mapRelays(relays: relaysResponse.bridge.relays, locations: relaysResponse.locations)
-            let filteredRelays = applyConstraints(
+            let filteredRelays = (try? applyConstraints(
                 location,
                 filterConstraint: filter,
+                daitaEnabled: false,
                 relays: mappedBridges
-            )
+            )) ?? []
             guard filteredRelays.isEmpty == false else { return relay(from: relaysResponse) }
 
             // Compute the midpoint location from all the filtered relays

diff --git a/ios/MullvadREST/Relay/RelaySelector+Wireguard.swift b/ios/MullvadREST/Relay/RelaySelector+Wireguard.swift
@@ -6,6 +6,7 @@
 //  Copyright © 2024 Mullvad VPN AB. All rights reserved.
 //
 
+import MullvadSettings
 import MullvadTypes
 
 extension RelaySelector {
@@ -14,13 +15,15 @@ extension RelaySelector {
         public static func findCandidates(
             by relayConstraint: RelayConstraint<UserSelectedRelays>,
             in relays: REST.ServerRelaysResponse,
-            filterConstraint: RelayConstraint<RelayFilter>
+            filterConstraint: RelayConstraint<RelayFilter>,
+            daitaEnabled: Bool
         ) throws -> [RelayWithLocation<REST.ServerRelay>] {
             let mappedRelays = mapRelays(relays: relays.wireguard.relays, locations: relays.locations)
 
-            return applyConstraints(
+            return try applyConstraints(
                 relayConstraint,
                 filterConstraint: filterConstraint,
+                daitaEnabled: daitaEnabled,
                 relays: mappedRelays
             )
         }
@@ -38,8 +41,12 @@ extension RelaySelector {
                 numberOfFailedAttempts: numberOfFailedAttempts
             )
 
-            guard let port, let relayWithLocation = pickRandomRelayByWeight(relays: relayWithLocations) else {
-                throw NoRelaysSatisfyingConstraintsError()
+            guard let port else {
+                throw NoRelaysSatisfyingConstraintsError(.invalidPort)
+            }
+
+            guard let relayWithLocation = pickRandomRelayByWeight(relays: relayWithLocations) else {
+                throw NoRelaysSatisfyingConstraintsError(.relayConstraintNotMatching)
             }
 
             let endpoint = MullvadEndpoint(