Merge branch 'fix-port-selector-randomizer'

4 files changed, 53 insertions, 10 deletions

diff --git a/ios/CHANGELOG.md b/ios/CHANGELOG.md
index 8cbe7d7993..4aa84066dd 100644
--- a/ios/CHANGELOG.md
+++ b/ios/CHANGELOG.md
@@ -37,6 +37,9 @@ Line wrap the file at 100 chars.                                              Th
   credits on accounts that no longer exist on our backend. Usually the case with newly created 
   accounts that went stale.
 
+### Fixed
+- Improve random port distribution. Should be less biased towards port 53.
+
 
 ## [2022.1] - 2022-02-15
 ### Added
diff --git a/ios/MullvadVPN/REST/ServerRelaysResponse.swift b/ios/MullvadVPN/REST/ServerRelaysResponse.swift
index 1f93954d33..a5606510dc 100644
--- a/ios/MullvadVPN/REST/ServerRelaysResponse.swift
+++ b/ios/MullvadVPN/REST/ServerRelaysResponse.swift
@@ -34,7 +34,7 @@ extension REST {
     struct ServerWireguardTunnels: Codable {
         let ipv4Gateway: IPv4Address
         let ipv6Gateway: IPv6Address
-        let portRanges: [ClosedRange<UInt16>]
+        let portRanges: [[UInt16]]
         let relays: [ServerRelay]
     }
 
diff --git a/ios/MullvadVPN/RelaySelector.swift b/ios/MullvadVPN/RelaySelector.swift
index fdf587e6d4..5e7492ca16 100644
--- a/ios/MullvadVPN/RelaySelector.swift
+++ b/ios/MullvadVPN/RelaySelector.swift
@@ -8,6 +8,7 @@
 
 import Foundation
 import Network
+import Logging
 
 struct RelaySelectorResult: Codable {
     var endpoint: MullvadEndpoint
@@ -38,16 +39,16 @@ extension RelaySelector {
     static func evaluate(relays: REST.ServerRelaysResponse, constraints: RelayConstraints) -> RelaySelectorResult? {
         let filteredRelays = applyConstraints(constraints, relays: Self.parseRelaysResponse(relays))
 
-        guard let relayWithLocation = pickRandomRelay(relays: filteredRelays) else {
-            return nil
-        }
-
-        guard let port = relays.wireguard.portRanges.randomElement()?.randomElement() else {
-            return nil
-        }
+        guard let relayWithLocation = pickRandomRelay(relays: filteredRelays),
+              let port = pickRandomPort(rawPortRanges: relays.wireguard.portRanges) else {
+                  return nil
+              }
 
         let endpoint = MullvadEndpoint(
-            ipv4Relay: IPv4Endpoint(ip: relayWithLocation.relay.ipv4AddrIn, port: port),
+            ipv4Relay: IPv4Endpoint(
+                ip: relayWithLocation.relay.ipv4AddrIn,
+                port: port
+            ),
             ipv6Relay: nil,
             ipv4Gateway: relays.wireguard.ipv4Gateway,
             ipv6Gateway: relays.wireguard.ipv6Gateway,
@@ -115,6 +116,45 @@ extension RelaySelector {
         return randomRelay
     }
 
+    private static func pickRandomPort(rawPortRanges: [[UInt16]]) -> UInt16? {
+        let portRanges = parseRawPortRanges(rawPortRanges)
+        let portAmount = portRanges.reduce(0) { partialResult, closedRange in
+            return 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
+            }
+        }
+
+        let logger = Logger(label: "RelaySelector")
+        logger.error("Port selection algorithm is broken!")
+
+        return nil
+    }
+
+    private static func parseRawPortRanges(_ rawPortRanges: [[UInt16]]) -> [ClosedRange<UInt16>] {
+        return 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
+            }
+        }
+    }
+
     private static func parseRelaysResponse(_ response: REST.ServerRelaysResponse) -> [RelayWithLocation] {
         return response.wireguard.relays.compactMap { (serverRelay) -> RelayWithLocation? in
             guard let serverLocation = response.locations[serverRelay.location] else { return nil }
diff --git a/ios/MullvadVPNTests/RelaySelectorTests.swift b/ios/MullvadVPNTests/RelaySelectorTests.swift
index caa9f82ab5..10dcb8502e 100644
--- a/ios/MullvadVPNTests/RelaySelectorTests.swift
+++ b/ios/MullvadVPNTests/RelaySelectorTests.swift
@@ -60,7 +60,7 @@ private let sampleRelays = REST.ServerRelaysResponse(
     wireguard: REST.ServerWireguardTunnels(
         ipv4Gateway: .loopback,
         ipv6Gateway: .loopback,
-        portRanges: [53...53],
+        portRanges: [[53, 53]],
         relays: [
             REST.ServerRelay(
                 hostname: "es1-wireguard",