Apply the obfuscation port to the entry configuration only

8 files changed, 115 insertions, 44 deletions

diff --git a/ios/MullvadREST/Relay/MultihopDecisionFlow.swift b/ios/MullvadREST/Relay/MultihopDecisionFlow.swift
index c8854bda3e..d543a54970 100644
--- a/ios/MullvadREST/Relay/MultihopDecisionFlow.swift
+++ b/ios/MullvadREST/Relay/MultihopDecisionFlow.swift
@@ -22,6 +22,7 @@ protocol MultihopDecisionFlow {
 struct OneToOne: MultihopDecisionFlow {
     let next: MultihopDecisionFlow?
     let relayPicker: RelayPicking
+
     init(next: (any MultihopDecisionFlow)?, relayPicker: RelayPicking) {
         self.next = next
         self.relayPicker = relayPicker
@@ -47,10 +48,11 @@ struct OneToOne: MultihopDecisionFlow {
             throw NoRelaysSatisfyingConstraintsError(.entryEqualsExit)
         }
 
-        let exitMatch = try relayPicker.findBestMatch(from: exitCandidates)
+        let exitMatch = try relayPicker.findBestMatch(from: exitCandidates, obfuscate: false)
         let entryMatch = try relayPicker.findBestMatch(
             from: entryCandidates,
-            closeTo: daitaAutomaticRouting ? exitMatch.location : nil
+            closeTo: daitaAutomaticRouting ? exitMatch.location : nil,
+            obfuscate: true
         )
 
         return SelectedRelays(entry: entryMatch, exit: exitMatch, retryAttempt: relayPicker.connectionAttemptCount)
@@ -95,8 +97,8 @@ struct OneToMany: MultihopDecisionFlow {
                 .pick(entryCandidates: entryCandidates, exitCandidates: exitCandidates, daitaAutomaticRouting: true)
         }
 
-        let entryMatch = try multihopPicker.findBestMatch(from: entryCandidates)
-        let exitMatch = try multihopPicker.exclude(relay: entryMatch, from: exitCandidates)
+        let entryMatch = try multihopPicker.findBestMatch(from: entryCandidates, obfuscate: true)
+        let exitMatch = try multihopPicker.exclude(relay: entryMatch, from: exitCandidates, obfuscate: false)
 
         return SelectedRelays(entry: entryMatch, exit: exitMatch, retryAttempt: relayPicker.connectionAttemptCount)
     }
@@ -135,11 +137,12 @@ struct ManyToOne: MultihopDecisionFlow {
             )
         }
 
-        let exitMatch = try multihopPicker.findBestMatch(from: exitCandidates)
+        let exitMatch = try multihopPicker.findBestMatch(from: exitCandidates, obfuscate: false)
         let entryMatch = try multihopPicker.exclude(
             relay: exitMatch,
             from: entryCandidates,
-            closeTo: daitaAutomaticRouting ? exitMatch.location : nil
+            closeTo: daitaAutomaticRouting ? exitMatch.location : nil,
+            obfuscate: true
         )
 
         return SelectedRelays(entry: entryMatch, exit: exitMatch, retryAttempt: relayPicker.connectionAttemptCount)
@@ -179,11 +182,12 @@ struct ManyToMany: MultihopDecisionFlow {
             )
         }
 
-        let exitMatch = try multihopPicker.findBestMatch(from: exitCandidates)
+        let exitMatch = try multihopPicker.findBestMatch(from: exitCandidates, obfuscate: false)
         let entryMatch = try multihopPicker.exclude(
             relay: exitMatch,
             from: entryCandidates,
-            closeTo: daitaAutomaticRouting ? exitMatch.location : nil
+            closeTo: daitaAutomaticRouting ? exitMatch.location : nil,
+            obfuscate: true
         )
 
         return SelectedRelays(entry: entryMatch, exit: exitMatch, retryAttempt: relayPicker.connectionAttemptCount)
diff --git a/ios/MullvadREST/Relay/ObfuscatorPortSelector.swift b/ios/MullvadREST/Relay/ObfuscatorPortSelector.swift
index 8cb2776d9b..c4d78becf9 100644
--- a/ios/MullvadREST/Relay/ObfuscatorPortSelector.swift
+++ b/ios/MullvadREST/Relay/ObfuscatorPortSelector.swift
@@ -9,7 +9,7 @@
 import MullvadSettings
 import MullvadTypes
 
-struct ObfuscatorPortSelectorResult {
+struct ObfuscatorPortSelection {
     let relays: REST.ServerRelaysResponse
     let port: RelayConstraint<UInt16>
 }
@@ -20,7 +20,7 @@ struct ObfuscatorPortSelector {
     func obfuscate(
         tunnelSettings: LatestTunnelSettings,
         connectionAttemptCount: UInt
-    ) throws -> ObfuscatorPortSelectorResult {
+    ) throws -> ObfuscatorPortSelection {
         var relays = relays
         var port = tunnelSettings.relayConstraints.port
         let obfuscationMethod = ObfuscationMethodSelector.obfuscationMethodBy(
@@ -44,7 +44,7 @@ struct ObfuscatorPortSelector {
             break
         }
 
-        return ObfuscatorPortSelectorResult(relays: relays, port: port)
+        return ObfuscatorPortSelection(relays: relays, port: port)
     }
 
     private func obfuscateShadowsocksRelays(tunnelSettings: LatestTunnelSettings) -> REST.ServerRelaysResponse {
diff --git a/ios/MullvadREST/Relay/RelayPicking.swift b/ios/MullvadREST/Relay/RelayPicking.swift
index 3e94b42d52..f89220c9e0 100644
--- a/ios/MullvadREST/Relay/RelayPicking.swift
+++ b/ios/MullvadREST/Relay/RelayPicking.swift
@@ -10,6 +10,7 @@ import MullvadSettings
 import MullvadTypes
 
 protocol RelayPicking {
+    var obfuscation: ObfuscatorPortSelection { get }
     var relays: REST.ServerRelaysResponse { get }
     var constraints: RelayConstraints { get }
     var connectionAttemptCount: UInt { get }
@@ -20,12 +21,13 @@ protocol RelayPicking {
 extension RelayPicking {
     func findBestMatch(
         from candidates: [RelayWithLocation<REST.ServerRelay>],
-        closeTo location: Location? = nil
+        closeTo location: Location? = nil,
+        obfuscate: Bool
     ) throws -> SelectedRelay {
         let match = try RelaySelector.WireGuard.pickCandidate(
             from: candidates,
             relays: relays,
-            portConstraint: constraints.port,
+            portConstraint: obfuscate ? obfuscation.port : constraints.port,
             numberOfFailedAttempts: connectionAttemptCount,
             closeTo: location
         )
@@ -39,11 +41,15 @@ extension RelayPicking {
 }
 
 struct SinglehopPicker: RelayPicking {
-    let relays: REST.ServerRelaysResponse
+    let obfuscation: ObfuscatorPortSelection
     let constraints: RelayConstraints
     let connectionAttemptCount: UInt
     let daitaSettings: DAITASettings
 
+    var relays: REST.ServerRelaysResponse {
+        obfuscation.relays
+    }
+
     func pick() throws -> SelectedRelays {
         do {
             let exitCandidates = try RelaySelector.WireGuard.findCandidates(
@@ -53,14 +59,14 @@ struct SinglehopPicker: RelayPicking {
                 daitaEnabled: daitaSettings.daitaState.isEnabled
             )
 
-            let match = try findBestMatch(from: exitCandidates)
+            let match = try findBestMatch(from: exitCandidates, obfuscate: true)
             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,
+                    obfuscation: obfuscation,
                     constraints: constraints,
                     connectionAttemptCount: connectionAttemptCount,
                     daitaSettings: daitaSettings
@@ -73,11 +79,15 @@ struct SinglehopPicker: RelayPicking {
 }
 
 struct MultihopPicker: RelayPicking {
-    let relays: REST.ServerRelaysResponse
+    let obfuscation: ObfuscatorPortSelection
     let constraints: RelayConstraints
     let connectionAttemptCount: UInt
     let daitaSettings: DAITASettings
 
+    var relays: REST.ServerRelaysResponse {
+        obfuscation.relays
+    }
+
     func pick() throws -> SelectedRelays {
         let exitCandidates = try RelaySelector.WireGuard.findCandidates(
             by: constraints.exitLocations,
@@ -129,12 +139,13 @@ struct MultihopPicker: RelayPicking {
     func exclude(
         relay: SelectedRelay,
         from candidates: [RelayWithLocation<REST.ServerRelay>],
-        closeTo location: Location? = nil
+        closeTo location: Location? = nil,
+        obfuscate: Bool
     ) throws -> SelectedRelay {
         let filteredCandidates = candidates.filter { relayWithLocation in
             relayWithLocation.relay.hostname != relay.hostname
         }
 
-        return try findBestMatch(from: filteredCandidates, closeTo: location)
+        return try findBestMatch(from: filteredCandidates, closeTo: location, obfuscate: obfuscate)
     }
 }
diff --git a/ios/MullvadREST/Relay/RelaySelectorWrapper.swift b/ios/MullvadREST/Relay/RelaySelectorWrapper.swift
index 48f5bf87d2..e7a15aa78f 100644
--- a/ios/MullvadREST/Relay/RelaySelectorWrapper.swift
+++ b/ios/MullvadREST/Relay/RelaySelectorWrapper.swift
@@ -20,28 +20,25 @@ public final class RelaySelectorWrapper: RelaySelectorProtocol {
         tunnelSettings: LatestTunnelSettings,
         connectionAttemptCount: UInt
     ) throws -> SelectedRelays {
-        let obfuscationResult = try ObfuscatorPortSelector(
+        let obfuscation = try ObfuscatorPortSelector(
             relays: try relayCache.read().relays
         ).obfuscate(
             tunnelSettings: tunnelSettings,
             connectionAttemptCount: connectionAttemptCount
         )
 
-        var constraints = tunnelSettings.relayConstraints
-        constraints.port = obfuscationResult.port
-
         return switch tunnelSettings.tunnelMultihopState {
         case .off:
             try SinglehopPicker(
-                relays: obfuscationResult.relays,
-                constraints: constraints,
+                obfuscation: obfuscation,
+                constraints: tunnelSettings.relayConstraints,
                 connectionAttemptCount: connectionAttemptCount,
                 daitaSettings: tunnelSettings.daita
             ).pick()
         case .on:
             try MultihopPicker(
-                relays: obfuscationResult.relays,
-                constraints: constraints,
+                obfuscation: obfuscation,
+                constraints: tunnelSettings.relayConstraints,
                 connectionAttemptCount: connectionAttemptCount,
                 daitaSettings: tunnelSettings.daita
             ).pick()
diff --git a/ios/MullvadSettings/TunnelSettingsStrategy.swift b/ios/MullvadSettings/TunnelSettingsStrategy.swift
index 1c8a08f212..22a3721cc4 100644
--- a/ios/MullvadSettings/TunnelSettingsStrategy.swift
+++ b/ios/MullvadSettings/TunnelSettingsStrategy.swift
@@ -18,9 +18,7 @@ public struct TunnelSettingsStrategy: TunnelSettingsStrategyProtocol {
         newSettings: LatestTunnelSettings
     ) -> Bool {
         switch (oldSettings, newSettings) {
-        case let (old, new) where old.relayConstraints != new.relayConstraints,
-             let (old, new) where old.tunnelMultihopState != new.tunnelMultihopState,
-             let (old, new) where old.daita != new.daita:
+        case let (old, new) where old != new:
             true
         default:
             false
diff --git a/ios/MullvadVPNTests/MullvadREST/Relay/MultihopDecisionFlowTests.swift b/ios/MullvadVPNTests/MullvadREST/Relay/MultihopDecisionFlowTests.swift
index 2919cff700..319a9eae16 100644
--- a/ios/MullvadVPNTests/MullvadREST/Relay/MultihopDecisionFlowTests.swift
+++ b/ios/MullvadVPNTests/MullvadREST/Relay/MultihopDecisionFlowTests.swift
@@ -160,13 +160,16 @@ class MultihopDecisionFlowTests: XCTestCase {
 
 extension MultihopDecisionFlowTests {
     var picker: MultihopPicker {
+        let obfuscation = try? ObfuscatorPortSelector(relays: sampleRelays)
+            .obfuscate(tunnelSettings: LatestTunnelSettings(), connectionAttemptCount: 0)
+
         let constraints = RelayConstraints(
             entryLocations: .only(UserSelectedRelays(locations: [.city("se", "sto")])),
             exitLocations: .only(UserSelectedRelays(locations: [.city("se", "sto")]))
         )
 
         return MultihopPicker(
-            relays: sampleRelays,
+            obfuscation: obfuscation.unsafelyUnwrapped,
             constraints: constraints,
             connectionAttemptCount: 0,
             daitaSettings: DAITASettings(daitaState: .off)
diff --git a/ios/MullvadVPNTests/MullvadREST/Relay/RelayPickingTests.swift b/ios/MullvadVPNTests/MullvadREST/Relay/RelayPickingTests.swift
index 41b65e3946..39970afc25 100644
--- a/ios/MullvadVPNTests/MullvadREST/Relay/RelayPickingTests.swift
+++ b/ios/MullvadVPNTests/MullvadREST/Relay/RelayPickingTests.swift
@@ -15,6 +15,12 @@ import XCTest
 
 class RelayPickingTests: XCTestCase {
     let sampleRelays = ServerRelaysResponseStubs.sampleRelays
+    var obfuscation: ObfuscatorPortSelection!
+
+    override func setUpWithError() throws {
+        obfuscation = try ObfuscatorPortSelector(relays: sampleRelays)
+            .obfuscate(tunnelSettings: LatestTunnelSettings(), connectionAttemptCount: 0)
+    }
 
     // MARK: Single-/multihop
 
@@ -25,7 +31,7 @@ class RelayPickingTests: XCTestCase {
         )
 
         let picker = SinglehopPicker(
-            relays: sampleRelays,
+            obfuscation: obfuscation,
             constraints: constraints,
             connectionAttemptCount: 0,
             daitaSettings: DAITASettings()
@@ -44,7 +50,7 @@ class RelayPickingTests: XCTestCase {
         )
 
         let picker = MultihopPicker(
-            relays: sampleRelays,
+            obfuscation: obfuscation,
             constraints: constraints,
             connectionAttemptCount: 0,
             daitaSettings: DAITASettings()
@@ -63,7 +69,7 @@ class RelayPickingTests: XCTestCase {
         )
 
         let picker = MultihopPicker(
-            relays: sampleRelays,
+            obfuscation: obfuscation,
             constraints: constraints,
             connectionAttemptCount: 0,
             daitaSettings: DAITASettings()
@@ -87,7 +93,7 @@ class RelayPickingTests: XCTestCase {
         )
 
         let picker = SinglehopPicker(
-            relays: sampleRelays,
+            obfuscation: obfuscation,
             constraints: constraints,
             connectionAttemptCount: 0,
             daitaSettings: DAITASettings(daitaState: .on, directOnlyState: .off)
@@ -107,7 +113,7 @@ class RelayPickingTests: XCTestCase {
         )
 
         let picker = SinglehopPicker(
-            relays: sampleRelays,
+            obfuscation: obfuscation,
             constraints: constraints,
             connectionAttemptCount: 0,
             daitaSettings: DAITASettings(daitaState: .on, directOnlyState: .on)
@@ -124,7 +130,7 @@ class RelayPickingTests: XCTestCase {
         )
 
         let picker = SinglehopPicker(
-            relays: sampleRelays,
+            obfuscation: obfuscation,
             constraints: constraints,
             connectionAttemptCount: 0,
             daitaSettings: DAITASettings(daitaState: .on, directOnlyState: .off)
@@ -144,7 +150,7 @@ class RelayPickingTests: XCTestCase {
         )
 
         let picker = SinglehopPicker(
-            relays: sampleRelays,
+            obfuscation: obfuscation,
             constraints: constraints,
             connectionAttemptCount: 0,
             daitaSettings: DAITASettings(daitaState: .on, directOnlyState: .on)
@@ -166,7 +172,7 @@ class RelayPickingTests: XCTestCase {
         )
 
         let picker = MultihopPicker(
-            relays: sampleRelays,
+            obfuscation: obfuscation,
             constraints: constraints,
             connectionAttemptCount: 0,
             daitaSettings: DAITASettings(daitaState: .on, directOnlyState: .off)
@@ -188,7 +194,7 @@ class RelayPickingTests: XCTestCase {
         )
 
         let picker = MultihopPicker(
-            relays: sampleRelays,
+            obfuscation: obfuscation,
             constraints: constraints,
             connectionAttemptCount: 0,
             daitaSettings: DAITASettings(daitaState: .on, directOnlyState: .off)
@@ -209,7 +215,7 @@ class RelayPickingTests: XCTestCase {
         )
 
         let picker = MultihopPicker(
-            relays: sampleRelays,
+            obfuscation: obfuscation,
             constraints: constraints,
             connectionAttemptCount: 0,
             daitaSettings: DAITASettings(daitaState: .on, directOnlyState: .on)
@@ -217,4 +223,56 @@ class RelayPickingTests: XCTestCase {
 
         XCTAssertThrowsError(try picker.pick())
     }
+
+    // MARK: Obfuscation
+
+    func testObfuscationOnSinglehop() throws {
+        let constraints = RelayConstraints(entryLocations: .any, exitLocations: .any, port: .only(5000))
+        let tunnelSettings = LatestTunnelSettings(
+            wireGuardObfuscation: WireGuardObfuscationSettings(
+                state: .udpOverTcp,
+                udpOverTcpPort: .port80
+            )
+        )
+
+        obfuscation = try ObfuscatorPortSelector(relays: sampleRelays)
+            .obfuscate(tunnelSettings: tunnelSettings, connectionAttemptCount: 0)
+
+        let picker = SinglehopPicker(
+            obfuscation: obfuscation,
+            constraints: constraints,
+            connectionAttemptCount: 0,
+            daitaSettings: DAITASettings()
+        )
+
+        let selectedRelays = try picker.pick()
+
+        XCTAssertNil(selectedRelays.entry?.endpoint.ipv4Relay.port)
+        XCTAssertEqual(selectedRelays.exit.endpoint.ipv4Relay.port, 80)
+    }
+
+    func testObfuscationOnMultihop() throws {
+        let constraints = RelayConstraints(entryLocations: .any, exitLocations: .any, port: .only(5000))
+        let tunnelSettings = LatestTunnelSettings(
+            wireGuardObfuscation: WireGuardObfuscationSettings(
+                state: .udpOverTcp,
+                udpOverTcpPort: .port80
+            )
+        )
+
+        obfuscation = try ObfuscatorPortSelector(relays: sampleRelays)
+            .obfuscate(tunnelSettings: tunnelSettings, connectionAttemptCount: 0)
+
+        let picker = MultihopPicker(
+            obfuscation: obfuscation,
+            constraints: constraints,
+            connectionAttemptCount: 0,
+            daitaSettings: DAITASettings()
+        )
+
+        let selectedRelays = try picker.pick()
+
+        XCTAssertEqual(selectedRelays.entry?.endpoint.ipv4Relay.port, 80)
+        XCTAssertEqual(selectedRelays.exit.endpoint.ipv4Relay.port, 5000)
+    }
 }
diff --git a/ios/MullvadVPNTests/MullvadVPN/TunnelManager/TunnelSettingsStrategyTests.swift b/ios/MullvadVPNTests/MullvadVPN/TunnelManager/TunnelSettingsStrategyTests.swift
index 97986de037..a10c42bdf0 100644
--- a/ios/MullvadVPNTests/MullvadVPN/TunnelManager/TunnelSettingsStrategyTests.swift
+++ b/ios/MullvadVPNTests/MullvadVPN/TunnelManager/TunnelSettingsStrategyTests.swift
@@ -52,7 +52,7 @@ final class TunnelSettingsStrategyTests: XCTestCase {
         TunnelSettingsUpdate.dnsSettings(dnsSettings).apply(to: &updatedSettings)
 
         let tunnelSettingsStrategy = TunnelSettingsStrategy()
-        XCTAssertFalse(tunnelSettingsStrategy.shouldReconnectToNewRelay(
+        XCTAssertTrue(tunnelSettingsStrategy.shouldReconnectToNewRelay(
             oldSettings: currentSettings,
             newSettings: updatedSettings
         ))
@@ -66,7 +66,7 @@ final class TunnelSettingsStrategyTests: XCTestCase {
         TunnelSettingsUpdate.quantumResistance(.on).apply(to: &updatedSettings)
 
         let tunnelSettingsStrategy = TunnelSettingsStrategy()
-        XCTAssertFalse(tunnelSettingsStrategy.shouldReconnectToNewRelay(
+        XCTAssertTrue(tunnelSettingsStrategy.shouldReconnectToNewRelay(
             oldSettings: currentSettings,
             newSettings: updatedSettings
         ))
@@ -88,7 +88,7 @@ final class TunnelSettingsStrategyTests: XCTestCase {
         .apply(to: &updatedSettings)
 
         let tunnelSettingsStrategy = TunnelSettingsStrategy()
-        XCTAssertFalse(tunnelSettingsStrategy.shouldReconnectToNewRelay(
+        XCTAssertTrue(tunnelSettingsStrategy.shouldReconnectToNewRelay(
             oldSettings: currentSettings,
             newSettings: updatedSettings
         ))