Enable custom Shadowsocks port selection on entry servers in multihop

12 files changed, 48 insertions, 41 deletions

diff --git a/ios/MullvadREST/Relay/ObfuscatorPortSelector.swift b/ios/MullvadREST/Relay/ObfuscatorPortSelector.swift
index e6ca92f4f7..4b2f2dca7f 100644
--- a/ios/MullvadREST/Relay/ObfuscatorPortSelector.swift
+++ b/ios/MullvadREST/Relay/ObfuscatorPortSelector.swift
@@ -10,9 +10,14 @@ import MullvadSettings
 import MullvadTypes
 
 struct ObfuscatorPortSelection {
-    let relays: REST.ServerRelaysResponse
+    let entryRelays: REST.ServerRelaysResponse
+    let exitRelays: REST.ServerRelaysResponse
     let port: RelayConstraint<UInt16>
     let method: WireGuardObfuscationState
+
+    var wireguard: REST.ServerWireguardTunnels {
+        exitRelays.wireguard
+    }
 }
 
 struct ObfuscatorPortSelector {
@@ -22,7 +27,9 @@ struct ObfuscatorPortSelector {
         tunnelSettings: LatestTunnelSettings,
         connectionAttemptCount: UInt
     ) throws -> ObfuscatorPortSelection {
-        var relays = relays
+        var entryRelays = relays
+        var exitRelays = relays
+
         var port = tunnelSettings.relayConstraints.port
         let obfuscationMethod = ObfuscationMethodSelector.obfuscationMethodBy(
             connectionAttemptCount: connectionAttemptCount,
@@ -36,7 +43,13 @@ struct ObfuscatorPortSelector {
                 connectionAttemptCount: connectionAttemptCount
             )
         case .shadowsocks:
-            relays = obfuscateShadowsocksRelays(tunnelSettings: tunnelSettings)
+            let filteredRelays = obfuscateShadowsocksRelays(tunnelSettings: tunnelSettings)
+            if tunnelSettings.tunnelMultihopState.isEnabled {
+                entryRelays = filteredRelays
+            } else {
+                exitRelays = filteredRelays
+            }
+
             port = obfuscateShadowsocksPort(
                 tunnelSettings: tunnelSettings,
                 shadowsocksPortRanges: relays.wireguard.shadowsocksPortRanges
@@ -45,7 +58,12 @@ struct ObfuscatorPortSelector {
             break
         }
 
-        return ObfuscatorPortSelection(relays: relays, port: port, method: obfuscationMethod)
+        return ObfuscatorPortSelection(
+            entryRelays: entryRelays,
+            exitRelays: exitRelays,
+            port: port,
+            method: obfuscationMethod
+        )
     }
 
     private func obfuscateShadowsocksRelays(tunnelSettings: LatestTunnelSettings) -> REST.ServerRelaysResponse {
diff --git a/ios/MullvadREST/Relay/RelayPicking.swift b/ios/MullvadREST/Relay/RelayPicking.swift
index 9ae31139f9..5d7460abec 100644
--- a/ios/MullvadREST/Relay/RelayPicking.swift
+++ b/ios/MullvadREST/Relay/RelayPicking.swift
@@ -12,7 +12,6 @@ import Network
 
 protocol RelayPicking {
     var obfuscation: ObfuscatorPortSelection { get }
-    var relays: REST.ServerRelaysResponse { get }
     var constraints: RelayConstraints { get }
     var connectionAttemptCount: UInt { get }
     var daitaSettings: DAITASettings { get }
@@ -27,7 +26,7 @@ extension RelayPicking {
     ) throws -> SelectedRelay {
         var match = try RelaySelector.WireGuard.pickCandidate(
             from: candidates,
-            relays: relays,
+            wireguard: obfuscation.wireguard,
             portConstraint: useObfuscatedPortIfAvailable ? obfuscation.port : constraints.port,
             numberOfFailedAttempts: connectionAttemptCount,
             closeTo: location
@@ -46,7 +45,7 @@ extension RelayPicking {
 
     private func applyShadowsocksIpAddress(in match: RelaySelectorMatch) -> RelaySelectorMatch {
         let port = match.endpoint.ipv4Relay.port
-        let portRanges = RelaySelector.parseRawPortRanges(relays.wireguard.shadowsocksPortRanges)
+        let portRanges = RelaySelector.parseRawPortRanges(obfuscation.wireguard.shadowsocksPortRanges)
         let portIsWithinRange = portRanges.contains(where: { $0.contains(port) })
 
         var endpoint = match.endpoint
@@ -76,15 +75,11 @@ struct SinglehopPicker: RelayPicking {
     let connectionAttemptCount: UInt
     let daitaSettings: DAITASettings
 
-    var relays: REST.ServerRelaysResponse {
-        obfuscation.relays
-    }
-
     func pick() throws -> SelectedRelays {
         do {
             let exitCandidates = try RelaySelector.WireGuard.findCandidates(
                 by: constraints.exitLocations,
-                in: relays,
+                in: obfuscation.exitRelays,
                 filterConstraint: constraints.filter,
                 daitaEnabled: daitaSettings.daitaState.isEnabled
             )
@@ -114,14 +109,10 @@ struct MultihopPicker: RelayPicking {
     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,
-            in: relays,
+            in: obfuscation.exitRelays,
             filterConstraint: constraints.filter,
             daitaEnabled: false
         )
@@ -153,7 +144,7 @@ struct MultihopPicker: RelayPicking {
         do {
             let entryCandidates = try RelaySelector.WireGuard.findCandidates(
                 by: daitaSettings.isAutomaticRouting ? .any : constraints.entryLocations,
-                in: relays,
+                in: obfuscation.entryRelays,
                 filterConstraint: constraints.filter,
                 daitaEnabled: daitaSettings.daitaState.isEnabled
             )
diff --git a/ios/MullvadREST/Relay/RelaySelector+Wireguard.swift b/ios/MullvadREST/Relay/RelaySelector+Wireguard.swift
index 0317788e23..1cb1d20ae5 100644
--- a/ios/MullvadREST/Relay/RelaySelector+Wireguard.swift
+++ b/ios/MullvadREST/Relay/RelaySelector+Wireguard.swift
@@ -31,14 +31,14 @@ extension RelaySelector {
         /// Picks a random relay from a list.
         public static func pickCandidate(
             from relayWithLocations: [RelayWithLocation<REST.ServerRelay>],
-            relays: REST.ServerRelaysResponse,
+            wireguard: REST.ServerWireguardTunnels,
             portConstraint: RelayConstraint<UInt16>,
             numberOfFailedAttempts: UInt,
             closeTo referenceLocation: Location? = nil
         ) throws -> RelaySelectorMatch {
             let port = try evaluatePort(
-                relays: relays,
                 portConstraint: portConstraint,
+                rawPortRanges: wireguard.portRanges,
                 numberOfFailedAttempts: numberOfFailedAttempts
             )
 
@@ -54,7 +54,7 @@ extension RelaySelector {
                 throw NoRelaysSatisfyingConstraintsError(.relayConstraintNotMatching)
             }
 
-            return createMatch(for: relayWithLocation, port: port, relays: relays)
+            return createMatch(for: relayWithLocation, port: port, wireguard: wireguard)
         }
 
         public static func closestRelay(
@@ -96,13 +96,13 @@ extension RelaySelector {
     }
 
     private static func evaluatePort(
-        relays: REST.ServerRelaysResponse,
         portConstraint: RelayConstraint<UInt16>,
+        rawPortRanges: [[UInt16]],
         numberOfFailedAttempts: UInt
     ) throws -> UInt16 {
         let port = applyPortConstraint(
             portConstraint,
-            rawPortRanges: relays.wireguard.portRanges,
+            rawPortRanges: rawPortRanges,
             numberOfFailedAttempts: numberOfFailedAttempts
         )
 
@@ -116,7 +116,7 @@ extension RelaySelector {
     private static func createMatch(
         for relayWithLocation: RelayWithLocation<REST.ServerRelay>,
         port: UInt16,
-        relays: REST.ServerRelaysResponse
+        wireguard: REST.ServerWireguardTunnels
     ) -> RelaySelectorMatch {
         let endpoint = MullvadEndpoint(
             ipv4Relay: IPv4Endpoint(
@@ -124,8 +124,8 @@ extension RelaySelector {
                 port: port
             ),
             ipv6Relay: nil,
-            ipv4Gateway: relays.wireguard.ipv4Gateway,
-            ipv6Gateway: relays.wireguard.ipv6Gateway,
+            ipv4Gateway: wireguard.ipv4Gateway,
+            ipv6Gateway: wireguard.ipv6Gateway,
             publicKey: relayWithLocation.relay.publicKey
         )
 
diff --git a/ios/MullvadVPN.xcodeproj/project.pbxproj b/ios/MullvadVPN.xcodeproj/project.pbxproj
index 8e60b70042..b80591e9f8 100644
--- a/ios/MullvadVPN.xcodeproj/project.pbxproj
+++ b/ios/MullvadVPN.xcodeproj/project.pbxproj
@@ -4251,12 +4251,12 @@
 				7ACE19102C1C349200260BB6 /* MultihopDecisionFlow.swift */,
 				F0F3161A2BF358590078DBCF /* NoRelaysSatisfyingConstraintsError.swift */,
 				7AD63A3A2CD5278900445268 /* ObfuscationMethodSelector.swift */,
+				7AD63A382CD520FD00445268 /* ObfuscatorPortSelector.swift */,
 				5820675A26E6576800655B05 /* RelayCache.swift */,
 				7A3AD5002C1068A800E9AD90 /* RelayPicking.swift */,
 				F0DDE4282B220A15006B57A7 /* RelaySelector.swift */,
 				F0B894F42BF7528700817A42 /* RelaySelector+Shadowsocks.swift */,
 				F0B894F22BF7526700817A42 /* RelaySelector+Wireguard.swift */,
-				7AD63A382CD520FD00445268 /* ObfuscatorPortSelector.swift */,
 				5824037F2A827DF300163DE8 /* RelaySelectorProtocol.swift */,
 				F0F316182BF3572B0078DBCF /* RelaySelectorResult.swift */,
 				7AEBA5292C2179F20018BEC5 /* RelaySelectorWrapper.swift */,
diff --git a/ios/MullvadVPN/View controllers/VPNSettings/VPNSettingsDataSource.swift b/ios/MullvadVPN/View controllers/VPNSettings/VPNSettingsDataSource.swift
index 6f5f59a240..2e3700039e 100644
--- a/ios/MullvadVPN/View controllers/VPNSettings/VPNSettingsDataSource.swift
+++ b/ios/MullvadVPN/View controllers/VPNSettings/VPNSettingsDataSource.swift
@@ -92,14 +92,12 @@ final class VPNSettingsDataSource: UITableViewDiffableDataSource<
         }
 
         static var wireGuardObfuscation: [Item] {
-            var items: [Item] = [
+            [
                 .wireGuardObfuscationAutomatic,
                 .wireGuardObfuscationShadowsocks,
                 .wireGuardObfuscationUdpOverTcp,
                 .wireGuardObfuscationOff,
             ]
-
-            return items
         }
 
         static var wireGuardObfuscationPort: [Item] {
diff --git a/ios/MullvadVPNTests/MullvadREST/Relay/ObfuscatorPortSelectorTests.swift b/ios/MullvadVPNTests/MullvadREST/Relay/ObfuscatorPortSelectorTests.swift
index f26bb149ae..70ad8244f5 100644
--- a/ios/MullvadVPNTests/MullvadREST/Relay/ObfuscatorPortSelectorTests.swift
+++ b/ios/MullvadVPNTests/MullvadREST/Relay/ObfuscatorPortSelectorTests.swift
@@ -153,7 +153,7 @@ final class ObfuscatorPortSelectorTests: XCTestCase {
             !relay.shadowsocksExtraAddrIn.isNil
         }
 
-        XCTAssertEqual(obfuscationResult.relays.wireguard.relays.count, relaysWithExtraAddresses.count)
+        XCTAssertEqual(obfuscationResult.wireguard.relays.count, relaysWithExtraAddresses.count)
     }
 
     func testObfuscateShadowsocksRelayFilteringWithPortInsideDefaultRanges() throws {
@@ -172,6 +172,6 @@ final class ObfuscatorPortSelectorTests: XCTestCase {
             connectionAttemptCount: 0
         )
 
-        XCTAssertEqual(obfuscationResult.relays.wireguard.relays.count, sampleRelays.wireguard.relays.count)
+        XCTAssertEqual(obfuscationResult.wireguard.relays.count, sampleRelays.wireguard.relays.count)
     }
 }
diff --git a/ios/MullvadVPNTests/MullvadREST/Relay/RelaySelectorTests.swift b/ios/MullvadVPNTests/MullvadREST/Relay/RelaySelectorTests.swift
index af14199e3b..0385cc5df0 100644
--- a/ios/MullvadVPNTests/MullvadREST/Relay/RelaySelectorTests.swift
+++ b/ios/MullvadVPNTests/MullvadREST/Relay/RelaySelectorTests.swift
@@ -304,7 +304,7 @@ extension RelaySelectorTests {
 
         return try RelaySelector.WireGuard.pickCandidate(
             from: candidates,
-            relays: relays,
+            wireguard: relays.wireguard,
             portConstraint: constraints.port,
             numberOfFailedAttempts: failedAttemptCount
         )
diff --git a/ios/PacketTunnelCoreTests/AppMessageHandlerTests.swift b/ios/PacketTunnelCoreTests/AppMessageHandlerTests.swift
index b300befc5a..b4398983ff 100644
--- a/ios/PacketTunnelCoreTests/AppMessageHandlerTests.swift
+++ b/ios/PacketTunnelCoreTests/AppMessageHandlerTests.swift
@@ -91,7 +91,7 @@ final class AppMessageHandlerTests: XCTestCase {
 
         let match = try RelaySelector.WireGuard.pickCandidate(
             from: candidates,
-            relays: ServerRelaysResponseStubs.sampleRelays,
+            wireguard: ServerRelaysResponseStubs.sampleRelays.wireguard,
             portConstraint: relayConstraints.port,
             numberOfFailedAttempts: 0
         )
diff --git a/ios/PacketTunnelCoreTests/EphemeralPeerExchangingPipelineTests.swift b/ios/PacketTunnelCoreTests/EphemeralPeerExchangingPipelineTests.swift
index e5a679ec58..2af86eedfe 100644
--- a/ios/PacketTunnelCoreTests/EphemeralPeerExchangingPipelineTests.swift
+++ b/ios/PacketTunnelCoreTests/EphemeralPeerExchangingPipelineTests.swift
@@ -31,7 +31,7 @@ final class EphemeralPeerExchangingPipelineTests: XCTestCase {
                 filterConstraint: relayConstraints.filter,
                 daitaEnabled: false
             ),
-            relays: ServerRelaysResponseStubs.sampleRelays,
+            wireguard: ServerRelaysResponseStubs.sampleRelays.wireguard,
             portConstraint: relayConstraints.port,
             numberOfFailedAttempts: 0
         )
@@ -43,7 +43,7 @@ final class EphemeralPeerExchangingPipelineTests: XCTestCase {
                 filterConstraint: relayConstraints.filter,
                 daitaEnabled: false
             ),
-            relays: ServerRelaysResponseStubs.sampleRelays,
+            wireguard: ServerRelaysResponseStubs.sampleRelays.wireguard,
             portConstraint: relayConstraints.port,
             numberOfFailedAttempts: 0
         )
diff --git a/ios/PacketTunnelCoreTests/MultiHopEphemeralPeerExchangerTests.swift b/ios/PacketTunnelCoreTests/MultiHopEphemeralPeerExchangerTests.swift
index 95436edb4c..a4c1d09155 100644
--- a/ios/PacketTunnelCoreTests/MultiHopEphemeralPeerExchangerTests.swift
+++ b/ios/PacketTunnelCoreTests/MultiHopEphemeralPeerExchangerTests.swift
@@ -30,7 +30,7 @@ final class MultiHopEphemeralPeerExchangerTests: XCTestCase {
                 filterConstraint: relayConstraints.filter,
                 daitaEnabled: false
             ),
-            relays: ServerRelaysResponseStubs.sampleRelays,
+            wireguard: ServerRelaysResponseStubs.sampleRelays.wireguard,
             portConstraint: relayConstraints.port,
             numberOfFailedAttempts: 0
         )
@@ -42,7 +42,7 @@ final class MultiHopEphemeralPeerExchangerTests: XCTestCase {
                 filterConstraint: relayConstraints.filter,
                 daitaEnabled: false
             ),
-            relays: ServerRelaysResponseStubs.sampleRelays,
+            wireguard: ServerRelaysResponseStubs.sampleRelays.wireguard,
             portConstraint: relayConstraints.port,
             numberOfFailedAttempts: 0
         )
diff --git a/ios/PacketTunnelCoreTests/ProtocolObfuscatorTests.swift b/ios/PacketTunnelCoreTests/ProtocolObfuscatorTests.swift
index 425c34705b..22ade73c7a 100644
--- a/ios/PacketTunnelCoreTests/ProtocolObfuscatorTests.swift
+++ b/ios/PacketTunnelCoreTests/ProtocolObfuscatorTests.swift
@@ -58,13 +58,13 @@ final class ProtocolObfuscatorTests: XCTestCase {
         let settings = settings(.automatic, obfuscationPort: .automatic)
 
         try (UInt(0) ... 3).forEach { attempt in
-            var obfuscatedEndpoint = obfuscator.obfuscate(endpoint, settings: settings, retryAttempts: attempt)
+            let obfuscatedEndpoint = obfuscator.obfuscate(endpoint, settings: settings, retryAttempts: attempt)
 
             switch attempt {
             case 0, 1:
                 XCTAssertEqual(endpoint, obfuscatedEndpoint)
             case 2, 3:
-                var obfuscationProtocol = try XCTUnwrap(obfuscator.tunnelObfuscator as? TunnelObfuscationStub)
+                let obfuscationProtocol = try XCTUnwrap(obfuscator.tunnelObfuscator as? TunnelObfuscationStub)
                 validate(obfuscatedEndpoint, against: obfuscationProtocol)
             default:
                 XCTExpectFailure("Should not end up here, test setup is wrong")
diff --git a/ios/PacketTunnelCoreTests/SingleHopEphemeralPeerExchangerTests.swift b/ios/PacketTunnelCoreTests/SingleHopEphemeralPeerExchangerTests.swift
index e94ca0889e..2ce3558fba 100644
--- a/ios/PacketTunnelCoreTests/SingleHopEphemeralPeerExchangerTests.swift
+++ b/ios/PacketTunnelCoreTests/SingleHopEphemeralPeerExchangerTests.swift
@@ -30,7 +30,7 @@ final class SingleHopEphemeralPeerExchangerTests: XCTestCase {
 
         let match = try RelaySelector.WireGuard.pickCandidate(
             from: candidates,
-            relays: ServerRelaysResponseStubs.sampleRelays,
+            wireguard: ServerRelaysResponseStubs.sampleRelays.wireguard,
             portConstraint: relayConstraints.port,
             numberOfFailedAttempts: 0
         )