Fix ip address selection for ShadowSocks

5 files changed, 67 insertions, 18 deletions

diff --git a/ios/MullvadREST/ApiHandlers/ServerRelaysResponse.swift b/ios/MullvadREST/ApiHandlers/ServerRelaysResponse.swift
index 5ea32c0951..bab0962bba 100644
--- a/ios/MullvadREST/ApiHandlers/ServerRelaysResponse.swift
+++ b/ios/MullvadREST/ApiHandlers/ServerRelaysResponse.swift
@@ -37,7 +37,7 @@ extension REST {
         public var daita: Bool?
 
         public func override(ipv4AddrIn: IPv4Address?) -> Self {
-            return BridgeRelay(
+            BridgeRelay(
                 hostname: hostname,
                 active: active,
                 owned: owned,
@@ -65,7 +65,7 @@ extension REST {
         public let shadowsocksExtraAddrIn: [String]?
 
         public func override(ipv4AddrIn: IPv4Address?, ipv6AddrIn: IPv6Address?) -> Self {
-            return ServerRelay(
+            ServerRelay(
                 hostname: hostname,
                 active: active,
                 owned: owned,
@@ -82,7 +82,7 @@ extension REST {
         }
 
         public func override(daita: Bool) -> Self {
-            return ServerRelay(
+            ServerRelay(
                 hostname: hostname,
                 active: active,
                 owned: owned,
diff --git a/ios/MullvadREST/Relay/MultihopDecisionFlow.swift b/ios/MullvadREST/Relay/MultihopDecisionFlow.swift
index d543a54970..f36422f4dd 100644
--- a/ios/MullvadREST/Relay/MultihopDecisionFlow.swift
+++ b/ios/MullvadREST/Relay/MultihopDecisionFlow.swift
@@ -48,11 +48,11 @@ struct OneToOne: MultihopDecisionFlow {
             throw NoRelaysSatisfyingConstraintsError(.entryEqualsExit)
         }
 
-        let exitMatch = try relayPicker.findBestMatch(from: exitCandidates, obfuscate: false)
+        let exitMatch = try relayPicker.findBestMatch(from: exitCandidates, useObfuscatedPortIfAvailable: false)
         let entryMatch = try relayPicker.findBestMatch(
             from: entryCandidates,
             closeTo: daitaAutomaticRouting ? exitMatch.location : nil,
-            obfuscate: true
+            useObfuscatedPortIfAvailable: true
         )
 
         return SelectedRelays(entry: entryMatch, exit: exitMatch, retryAttempt: relayPicker.connectionAttemptCount)
@@ -97,8 +97,12 @@ struct OneToMany: MultihopDecisionFlow {
                 .pick(entryCandidates: entryCandidates, exitCandidates: exitCandidates, daitaAutomaticRouting: true)
         }
 
-        let entryMatch = try multihopPicker.findBestMatch(from: entryCandidates, obfuscate: true)
-        let exitMatch = try multihopPicker.exclude(relay: entryMatch, from: exitCandidates, obfuscate: false)
+        let entryMatch = try multihopPicker.findBestMatch(from: entryCandidates, useObfuscatedPortIfAvailable: true)
+        let exitMatch = try multihopPicker.exclude(
+            relay: entryMatch,
+            from: exitCandidates,
+            useObfuscatedPortIfAvailable: false
+        )
 
         return SelectedRelays(entry: entryMatch, exit: exitMatch, retryAttempt: relayPicker.connectionAttemptCount)
     }
@@ -137,12 +141,12 @@ struct ManyToOne: MultihopDecisionFlow {
             )
         }
 
-        let exitMatch = try multihopPicker.findBestMatch(from: exitCandidates, obfuscate: false)
+        let exitMatch = try multihopPicker.findBestMatch(from: exitCandidates, useObfuscatedPortIfAvailable: false)
         let entryMatch = try multihopPicker.exclude(
             relay: exitMatch,
             from: entryCandidates,
             closeTo: daitaAutomaticRouting ? exitMatch.location : nil,
-            obfuscate: true
+            useObfuscatedPortIfAvailable: true
         )
 
         return SelectedRelays(entry: entryMatch, exit: exitMatch, retryAttempt: relayPicker.connectionAttemptCount)
@@ -182,12 +186,12 @@ struct ManyToMany: MultihopDecisionFlow {
             )
         }
 
-        let exitMatch = try multihopPicker.findBestMatch(from: exitCandidates, obfuscate: false)
+        let exitMatch = try multihopPicker.findBestMatch(from: exitCandidates, useObfuscatedPortIfAvailable: false)
         let entryMatch = try multihopPicker.exclude(
             relay: exitMatch,
             from: entryCandidates,
             closeTo: daitaAutomaticRouting ? exitMatch.location : nil,
-            obfuscate: true
+            useObfuscatedPortIfAvailable: 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 c4d78becf9..152cf6b638 100644
--- a/ios/MullvadREST/Relay/ObfuscatorPortSelector.swift
+++ b/ios/MullvadREST/Relay/ObfuscatorPortSelector.swift
@@ -12,6 +12,7 @@ import MullvadTypes
 struct ObfuscatorPortSelection {
     let relays: REST.ServerRelaysResponse
     let port: RelayConstraint<UInt16>
+    let method: WireGuardObfuscationState
 }
 
 struct ObfuscatorPortSelector {
@@ -44,7 +45,7 @@ struct ObfuscatorPortSelector {
             break
         }
 
-        return ObfuscatorPortSelection(relays: relays, port: port)
+        return ObfuscatorPortSelection(relays: relays, 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 f89220c9e0..9ae31139f9 100644
--- a/ios/MullvadREST/Relay/RelayPicking.swift
+++ b/ios/MullvadREST/Relay/RelayPicking.swift
@@ -8,6 +8,7 @@
 
 import MullvadSettings
 import MullvadTypes
+import Network
 
 protocol RelayPicking {
     var obfuscation: ObfuscatorPortSelection { get }
@@ -22,22 +23,51 @@ extension RelayPicking {
     func findBestMatch(
         from candidates: [RelayWithLocation<REST.ServerRelay>],
         closeTo location: Location? = nil,
-        obfuscate: Bool
+        useObfuscatedPortIfAvailable: Bool
     ) throws -> SelectedRelay {
-        let match = try RelaySelector.WireGuard.pickCandidate(
+        var match = try RelaySelector.WireGuard.pickCandidate(
             from: candidates,
             relays: relays,
-            portConstraint: obfuscate ? obfuscation.port : constraints.port,
+            portConstraint: useObfuscatedPortIfAvailable ? obfuscation.port : constraints.port,
             numberOfFailedAttempts: connectionAttemptCount,
             closeTo: location
         )
 
+        if useObfuscatedPortIfAvailable && obfuscation.method == .shadowsocks {
+            match = applyShadowsocksIpAddress(in: match)
+        }
+
         return SelectedRelay(
             endpoint: match.endpoint,
             hostname: match.relay.hostname,
             location: match.location
         )
     }
+
+    private func applyShadowsocksIpAddress(in match: RelaySelectorMatch) -> RelaySelectorMatch {
+        let port = match.endpoint.ipv4Relay.port
+        let portRanges = RelaySelector.parseRawPortRanges(relays.wireguard.shadowsocksPortRanges)
+        let portIsWithinRange = portRanges.contains(where: { $0.contains(port) })
+
+        var endpoint = match.endpoint
+
+        // If the currently selected obfuscation port is not within the allowed range (as specified
+        // in the relay list), we should use one of the extra Shadowsocks IP addresses instead of
+        // the default one.
+        if !portIsWithinRange {
+            var ipv4Address = match.endpoint.ipv4Relay.ip
+            if let shadowsocksAddress = match.relay.shadowsocksExtraAddrIn?.randomElement() {
+                ipv4Address = IPv4Address(shadowsocksAddress) ?? ipv4Address
+            }
+
+            endpoint = match.endpoint.override(ipv4Relay: IPv4Endpoint(
+                ip: ipv4Address,
+                port: port
+            ))
+        }
+
+        return RelaySelectorMatch(endpoint: endpoint, relay: match.relay, location: match.location)
+    }
 }
 
 struct SinglehopPicker: RelayPicking {
@@ -59,7 +89,7 @@ struct SinglehopPicker: RelayPicking {
                 daitaEnabled: daitaSettings.daitaState.isEnabled
             )
 
-            let match = try findBestMatch(from: exitCandidates, obfuscate: true)
+            let match = try findBestMatch(from: exitCandidates, useObfuscatedPortIfAvailable: 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
@@ -140,12 +170,16 @@ struct MultihopPicker: RelayPicking {
         relay: SelectedRelay,
         from candidates: [RelayWithLocation<REST.ServerRelay>],
         closeTo location: Location? = nil,
-        obfuscate: Bool
+        useObfuscatedPortIfAvailable: Bool
     ) throws -> SelectedRelay {
         let filteredCandidates = candidates.filter { relayWithLocation in
             relayWithLocation.relay.hostname != relay.hostname
         }
 
-        return try findBestMatch(from: filteredCandidates, closeTo: location, obfuscate: obfuscate)
+        return try findBestMatch(
+            from: filteredCandidates,
+            closeTo: location,
+            useObfuscatedPortIfAvailable: useObfuscatedPortIfAvailable
+        )
     }
 }
diff --git a/ios/MullvadTypes/MullvadEndpoint.swift b/ios/MullvadTypes/MullvadEndpoint.swift
index 9c05111c8c..1361df2e46 100644
--- a/ios/MullvadTypes/MullvadEndpoint.swift
+++ b/ios/MullvadTypes/MullvadEndpoint.swift
@@ -30,4 +30,14 @@ public struct MullvadEndpoint: Equatable, Codable {
         self.ipv6Gateway = ipv6Gateway
         self.publicKey = publicKey
     }
+
+    public func override(ipv4Relay: IPv4Endpoint) -> Self {
+        MullvadEndpoint(
+            ipv4Relay: ipv4Relay,
+            ipv6Relay: ipv6Relay,
+            ipv4Gateway: ipv4Gateway,
+            ipv6Gateway: ipv6Gateway,
+            publicKey: publicKey
+        )
+    }
 }