Update relay selector for shadowsocks obfuscation

ios/MullvadREST/Extensions/UInt+Counting.swift

@@ -0,0 +1,24 @@
+//
+//  UInt+Counting.swift
+//  MullvadVPN
+//
+//  Created by Jon Petersson on 2024-11-05.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import Foundation
+
+extension UInt {
+    /// Determines whether a number has a specific order in a given set.
+    /// Eg. `6.isOrdered(nth: 3, forEverySetOf: 4)` -> "Is a 6 ordered third in an arbitrary
+    /// amount of sets of four?". The result of this is `true`, since in a range of eg. 0-7 a six
+    /// would be considered third if the range was divided into sets of 4.
+    public func isOrdered(nth: UInt, forEverySetOf set: UInt) -> Bool {
+        guard nth > 0, set > 0 else {
+            assertionFailure("Both 'nth' and 'set' must be positive")
+            return false
+        }
+
+        return self % set == nth - 1
+    }
+}

ios/MullvadREST/Relay/ObfuscationMethodSelector.swift

@@ -0,0 +1,30 @@
+//
+//  ObfuscationMethodSelector.swift
+//  MullvadREST
+//
+//  Created by Jon Petersson on 2024-11-01.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import MullvadSettings
+
+public struct ObfuscationMethodSelector {
+    /// This retry logic used is explained at the following link:
+    /// https://github.com/mullvad/mullvadvpn-app/blob/main/docs/relay-selector.md#default-constraints-for-tunnel-endpoints
+    public static func obfuscationMethodBy(
+        connectionAttemptCount: UInt,
+        tunnelSettings: LatestTunnelSettings
+    ) -> WireGuardObfuscationState {
+        if tunnelSettings.wireGuardObfuscation.state == .automatic {
+            if connectionAttemptCount.isOrdered(nth: 3, forEverySetOf: 4) {
+                .shadowsocks
+            } else if connectionAttemptCount.isOrdered(nth: 4, forEverySetOf: 4) {
+                .udpOverTcp
+            } else {
+                .off
+            }
+        } else {
+            tunnelSettings.wireGuardObfuscation.state
+        }
+    }
+}

ios/MullvadREST/Relay/ObfuscatorPortSelector.swift

@@ -0,0 +1,128 @@
+//
+//  ObfuscatorPortSelector.swift
+//  MullvadVPN
+//
+//  Created by Jon Petersson on 2024-11-01.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import MullvadSettings
+import MullvadTypes
+
+struct ObfuscatorPortSelectorResult {
+    let relays: REST.ServerRelaysResponse
+    let port: RelayConstraint<UInt16>
+}
+
+struct ObfuscatorPortSelector {
+    let relays: REST.ServerRelaysResponse
+
+    func obfuscate(
+        tunnelSettings: LatestTunnelSettings,
+        connectionAttemptCount: UInt
+    ) throws -> ObfuscatorPortSelectorResult {
+        var relays = relays
+        var port = tunnelSettings.relayConstraints.port
+        let obfuscationMethod = ObfuscationMethodSelector.obfuscationMethodBy(
+            connectionAttemptCount: connectionAttemptCount,
+            tunnelSettings: tunnelSettings
+        )
+
+        switch obfuscationMethod {
+        case .udpOverTcp:
+            port = obfuscateUdpOverTcpPort(
+                tunnelSettings: tunnelSettings,
+                connectionAttemptCount: connectionAttemptCount
+            )
+        case .shadowsocks:
+            relays = obfuscateShadowsocksRelays(tunnelSettings: tunnelSettings)
+            port = obfuscateShadowsocksPort(
+                tunnelSettings: tunnelSettings,
+                shadowsocksPortRanges: relays.wireguard.shadowsocksPortRanges
+            )
+        default:
+            break
+        }
+
+        return ObfuscatorPortSelectorResult(relays: relays, port: port)
+    }
+
+    private func obfuscateShadowsocksRelays(tunnelSettings: LatestTunnelSettings) -> REST.ServerRelaysResponse {
+        let relays = relays
+        let wireGuardObfuscation = tunnelSettings.wireGuardObfuscation
+
+        return wireGuardObfuscation.state == .shadowsocks
+            ? filterShadowsocksRelays(from: relays, for: wireGuardObfuscation.shadowsocksPort)
+            : relays
+    }
+
+    private func filterShadowsocksRelays(
+        from relays: REST.ServerRelaysResponse,
+        for port: WireGuardObfuscationShadowsockPort
+    ) -> REST.ServerRelaysResponse {
+        let portRanges = RelaySelector.parseRawPortRanges(relays.wireguard.shadowsocksPortRanges)
+
+        // If the selected port is within the shadowsocks port ranges we can select from all relays.
+        guard
+            case let .custom(port) = port,
+            !portRanges.contains(where: { $0.contains(port) })
+        else {
+            return relays
+        }
+
+        let filteredRelays = relays.wireguard.relays.filter { relay in
+            relay.shadowsocksExtraAddrIn != nil
+        }
+
+        return REST.ServerRelaysResponse(
+            locations: relays.locations,
+            wireguard: REST.ServerWireguardTunnels(
+                ipv4Gateway: relays.wireguard.ipv4Gateway,
+                ipv6Gateway: relays.wireguard.ipv6Gateway,
+                portRanges: relays.wireguard.portRanges,
+                relays: filteredRelays,
+                shadowsocksPortRanges: relays.wireguard.shadowsocksPortRanges
+            ),
+            bridge: relays.bridge
+        )
+    }
+
+    private func obfuscateUdpOverTcpPort(
+        tunnelSettings: LatestTunnelSettings,
+        connectionAttemptCount: UInt
+    ) -> RelayConstraint<UInt16> {
+        switch tunnelSettings.wireGuardObfuscation.udpOverTcpPort {
+        case .automatic:
+            ((connectionAttemptCount % 2 == 0) ? .only(80) : .only(5001))
+        case .port5001:
+            .only(5001)
+        case .port80:
+            .only(80)
+        }
+    }
+
+    private func obfuscateShadowsocksPort(
+        tunnelSettings: LatestTunnelSettings,
+        shadowsocksPortRanges: [[UInt16]]
+    ) -> RelayConstraint<UInt16> {
+        let wireGuardObfuscation = tunnelSettings.wireGuardObfuscation
+
+        let shadowsockPort: () -> UInt16? = {
+            switch wireGuardObfuscation.shadowsocksPort {
+            case let .custom(port):
+                port
+            default:
+                RelaySelector.pickRandomPort(rawPortRanges: shadowsocksPortRanges)
+            }
+        }
+
+        guard
+            wireGuardObfuscation.state == .shadowsocks,
+            let port = shadowsockPort()
+        else {
+            return tunnelSettings.relayConstraints.port
+        }
+
+        return .only(port)
+    }
+}

ios/MullvadREST/Relay/RelaySelector.swift

@@ -30,36 +30,11 @@ public enum RelaySelector {
         }
     }
 
-    // MARK: - private
-
     static func pickRandomRelayByWeight<T: AnyRelay>(relays: [RelayWithLocation<T>])
         -> RelayWithLocation<T>? {
         rouletteSelection(relays: relays, weightFunction: { relayWithLocation in relayWithLocation.relay.weight })
     }
 
-    private static func pickRandomPort(rawPortRanges: [[UInt16]]) -> UInt16? {
-        let portRanges = parseRawPortRanges(rawPortRanges)
-        let portAmount = portRanges.reduce(0) { partialResult, closedRange in
-            partialResult + closedRange.count
-        }
-
-        guard var portIndex = (0 ..< portAmount).randomElement() else {
-            return nil
-        }
-
-        for range in portRanges {
-            if portIndex < range.count {
-                return UInt16(portIndex) + range.lowerBound
-            } else {
-                portIndex -= range.count
-            }
-        }
-
-        assertionFailure("Port selection algorithm is broken!")
-
-        return nil
-    }
-
     static func rouletteSelection<T>(relays: [T], weightFunction: (T) -> UInt64) -> T? {
         let totalWeight = relays.map { weightFunction($0) }.reduce(0) { accumulated, weight in
             accumulated + weight
@@ -97,40 +72,6 @@ public enum RelaySelector {
         }
     }
 
-    private static func makeRelayWithLocationFrom<T: AnyRelay>(
-        _ serverLocation: REST.ServerLocation,
-        relay: T
-    ) -> RelayWithLocation<T>? {
-        let locationComponents = relay.location.split(separator: "-")
-        guard locationComponents.count > 1 else { return nil }
-
-        let location = Location(
-            country: serverLocation.country,
-            countryCode: String(locationComponents[0]),
-            city: serverLocation.city,
-            cityCode: String(locationComponents[1]),
-            latitude: serverLocation.latitude,
-            longitude: serverLocation.longitude
-        )
-
-        return RelayWithLocation(relay: relay, serverLocation: location)
-    }
-
-    private static func parseRawPortRanges(_ rawPortRanges: [[UInt16]]) -> [ClosedRange<UInt16>] {
-        rawPortRanges.compactMap { inputRange -> ClosedRange<UInt16>? in
-            guard inputRange.count == 2 else { return nil }
-
-            let startPort = inputRange[0]
-            let endPort = inputRange[1]
-
-            if startPort <= endPort {
-                return startPort ... endPort
-            } else {
-                return nil
-            }
-        }
-    }
-
     /// Produce a list of `RelayWithLocation` items satisfying the given constraints
     static func applyConstraints<T: AnyRelay>(
         _ relayConstraint: RelayConstraint<UserSelectedRelays>,
@@ -166,6 +107,65 @@ public enum RelaySelector {
         }
     }
 
+    // MARK: - private
+
+    static func parseRawPortRanges(_ rawPortRanges: [[UInt16]]) -> [ClosedRange<UInt16>] {
+        rawPortRanges.compactMap { inputRange -> ClosedRange<UInt16>? in
+            guard inputRange.count == 2 else { return nil }
+
+            let startPort = inputRange[0]
+            let endPort = inputRange[1]
+
+            if startPort <= endPort {
+                return startPort ... endPort
+            } else {
+                return nil
+            }
+        }
+    }
+
+    static func pickRandomPort(rawPortRanges: [[UInt16]]) -> UInt16? {
+        let portRanges = parseRawPortRanges(rawPortRanges)
+        let portAmount = portRanges.reduce(0) { partialResult, closedRange in
+            partialResult + closedRange.count
+        }
+
+        guard var portIndex = (0 ..< portAmount).randomElement() else {
+            return nil
+        }
+
+        for range in portRanges {
+            if portIndex < range.count {
+                return UInt16(portIndex) + range.lowerBound
+            } else {
+                portIndex -= range.count
+            }
+        }
+
+        assertionFailure("Port selection algorithm is broken!")
+
+        return nil
+    }
+
+    private static func makeRelayWithLocationFrom<T: AnyRelay>(
+        _ serverLocation: REST.ServerLocation,
+        relay: T
+    ) -> RelayWithLocation<T>? {
+        let locationComponents = relay.location.split(separator: "-")
+        guard locationComponents.count > 1 else { return nil }
+
+        let location = Location(
+            country: serverLocation.country,
+            countryCode: String(locationComponents[0]),
+            city: serverLocation.city,
+            cityCode: String(locationComponents[1]),
+            latitude: serverLocation.latitude,
+            longitude: serverLocation.longitude
+        )
+
+        return RelayWithLocation(relay: relay, serverLocation: location)
+    }
+
     private static func filterByActive<T: AnyRelay>(
         relays: [RelayWithLocation<T>]
     ) throws -> [RelayWithLocation<T>] {

ios/MullvadREST/Relay/RelaySelectorWrapper.swift

@@ -20,20 +20,28 @@ public final class RelaySelectorWrapper: RelaySelectorProtocol {
         tunnelSettings: LatestTunnelSettings,
         connectionAttemptCount: UInt
     ) throws -> SelectedRelays {
-        let relays = try relayCache.read().relays
+        let obfuscationResult = 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: relays,
-                constraints: tunnelSettings.relayConstraints,
+                relays: obfuscationResult.relays,
+                constraints: constraints,
                 connectionAttemptCount: connectionAttemptCount,
                 daitaSettings: tunnelSettings.daita
             ).pick()
         case .on:
             try MultihopPicker(
-                relays: relays,
-                constraints: tunnelSettings.relayConstraints,
+                relays: obfuscationResult.relays,
+                constraints: constraints,
                 connectionAttemptCount: connectionAttemptCount,
                 daitaSettings: tunnelSettings.daita
             ).pick()

ios/MullvadSettings/WireGuardObfuscationSettings.swift

@@ -8,9 +8,9 @@
 
 import Foundation
 
-/// Whether obfuscation is enabled and which method is used
+/// Whether obfuscation is enabled and which method is used.
 ///
-/// `.automatic` means an algorithm will decide whether to use it or not.
+/// `.automatic` means an algorithm will decide whether to use obfuscation or not.
 public enum WireGuardObfuscationState: Codable {
     @available(*, deprecated, renamed: "udpOverTcp")
     case on