Pick random relay when all relays within the list have zero weight.

Fix roulette algorithm to skip relays with zero weight.

1 files changed, 31 insertions, 9 deletions

diff --git a/ios/MullvadVPN/RelaySelector.swift b/ios/MullvadVPN/RelaySelector.swift
index a3f9f7e474..f0bb12ce7f 100644
--- a/ios/MullvadVPN/RelaySelector.swift
+++ b/ios/MullvadVPN/RelaySelector.swift
@@ -36,16 +36,11 @@ enum RelaySelector {}
 extension RelaySelector {
 
     static func evaluate(relays: REST.ServerRelaysResponse, constraints: RelayConstraints) -> RelaySelectorResult? {
-        let filteredRelays = Self.applyConstraints(constraints, relays: Self.parseRelaysResponse(relays))
-        let totalWeight = filteredRelays.reduce(0) { $0 + $1.relay.weight }
+        let filteredRelays = applyConstraints(constraints, relays: Self.parseRelaysResponse(relays))
 
-        guard totalWeight > 0 else { return nil }
-        guard var i = (0...totalWeight).randomElement() else { return nil }
-
-        let relayWithLocation = filteredRelays.first { (relayWithLocation) -> Bool in
-            i -= relayWithLocation.relay.weight
-            return i <= 0
-        }.unsafelyUnwrapped
+        guard let relayWithLocation = pickRandomRelay(relays: filteredRelays) else {
+            return nil
+        }
 
         guard let port = relays.wireguard.portRanges.randomElement()?.randomElement() else {
             return nil
@@ -93,6 +88,33 @@ extension RelaySelector {
         }
     }
 
+    private static func pickRandomRelay(relays: [RelayWithLocation]) -> RelayWithLocation? {
+        let totalWeight = relays.reduce(0) { accummulatedWeight, relayWithLocation in
+            return accummulatedWeight + relayWithLocation.relay.weight
+        }
+
+        // Return random relay when all relays within the list have zero weight.
+        guard totalWeight > 0 else {
+            return relays.randomElement()
+        }
+
+        // Pick a random number in the range 1 - totalWeight. This choses the relay with a
+        // non-zero weight.
+        var i = (1...totalWeight).randomElement()!
+
+        let randomRelay = relays.first { (relayWithLocation) -> Bool in
+            let (result, isOverflow) = i.subtractingReportingOverflow(relayWithLocation.relay.weight)
+
+            i = isOverflow ? 0 : result
+
+            return i == 0
+        }
+
+        precondition(randomRelay != nil, "At least one relay must've had a weight above 0")
+
+        return randomRelay
+    }
+
     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 }