Add LWO to automatic obfuscation rotationadd-lwo-to-automatic-obfuscation-rotation-ios-1458

8 files changed, 75 insertions, 108 deletions

diff --git a/docs/relay-selector.md b/docs/relay-selector.md
index 26e9af0d59..0894d9d8bb 100644
--- a/docs/relay-selector.md
+++ b/docs/relay-selector.md
@@ -60,6 +60,7 @@ As such, the above algorithm is simplified to the following version:
 - The second attempt will connect to a relay on a random port using Shadowsocks for obfuscation
 - The third attempt will connect to a relay using QUIC for obfuscation
 - The fourth attempt will connect to a relay on a random port using [UDP2TCP obfuscation](https://github.com/mullvad/udp-over-tcp)
+- The fifth attempt will connect to a relay using LWO
 
 ### Random Ports for UDP2TCP and Shadowsocks
 
diff --git a/ios/MullvadREST/Extensions/UInt+Counting.swift b/ios/MullvadREST/Extensions/UInt+Counting.swift
deleted file mode 100644
index 4158b27f10..0000000000
--- a/ios/MullvadREST/Extensions/UInt+Counting.swift
+++ /dev/null
@@ -1,24 +0,0 @@
-//
-//  UInt+Counting.swift
-//  MullvadVPN
-//
-//  Created by Jon Petersson on 2024-11-05.
-//  Copyright © 2025 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
-    }
-}
diff --git a/ios/MullvadREST/Relay/ObfuscationMethodSelector.swift b/ios/MullvadREST/Relay/ObfuscationMethodSelector.swift
index 6bae282368..193d9e6ab6 100644
--- a/ios/MullvadREST/Relay/ObfuscationMethodSelector.swift
+++ b/ios/MullvadREST/Relay/ObfuscationMethodSelector.swift
@@ -14,6 +14,21 @@ public protocol ObfuscationProviding {
 }
 
 public struct ObfuscationMethodSelector {
+    public static var obfuscationOrder: [WireGuardObfuscationState] {
+        var methods: [WireGuardObfuscationState] = [
+            .off,
+            .shadowsocks,
+            .quic,
+            .udpOverTcp,
+        ]
+
+        #if DEBUG
+            methods.append(.lwo)
+        #endif
+
+        return methods
+    }
+
     /// 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
     ///
@@ -24,19 +39,9 @@ public struct ObfuscationMethodSelector {
         obfuscationBypass: any ObfuscationProviding
     ) -> WireGuardObfuscationState {
         if tunnelSettings.wireGuardObfuscation.state == .automatic {
-            let selectedObfuscation: WireGuardObfuscationState =
-                if connectionAttemptCount.isOrdered(
-                    nth: 2,
-                    forEverySetOf: 4
-                ) {
-                    .shadowsocks
-                } else if connectionAttemptCount.isOrdered(nth: 3, forEverySetOf: 4) {
-                    .quic
-                } else if connectionAttemptCount.isOrdered(nth: 4, forEverySetOf: 4) {
-                    .udpOverTcp
-                } else {
-                    .off
-                }
+            let attemptIndex = Int(connectionAttemptCount) % obfuscationOrder.count
+            let selectedObfuscation = obfuscationOrder[attemptIndex]
+
             return obfuscationBypass.bypassUnsupportedObfuscation(selectedObfuscation)
         }
         return tunnelSettings.wireGuardObfuscation.state
diff --git a/ios/MullvadVPN.xcodeproj/project.pbxproj b/ios/MullvadVPN.xcodeproj/project.pbxproj
index dedad28af8..88fa957f0a 100644
--- a/ios/MullvadVPN.xcodeproj/project.pbxproj
+++ b/ios/MullvadVPN.xcodeproj/project.pbxproj
@@ -659,8 +659,6 @@
 		7AD63A3B2CD5278900445268 /* ObfuscationMethodSelector.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7AD63A3A2CD5278900445268 /* ObfuscationMethodSelector.swift */; };
 		7AD63A3D2CD9065D00445268 /* RelayObfuscatorTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7AD63A3C2CD9065100445268 /* RelayObfuscatorTests.swift */; };
 		7AD63A3F2CDA53F600445268 /* ObfuscationMethodSelectorTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7AD63A3E2CDA53E900445268 /* ObfuscationMethodSelectorTests.swift */; };
-		7AD63A442CDA663300445268 /* UInt+Counting.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7AD63A432CDA662900445268 /* UInt+Counting.swift */; };
-		7AD63A472CDA666100445268 /* UIntTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7AD63A462CDA665A00445268 /* UIntTests.swift */; };
 		7ADCB2D82B6A6EB300C88F89 /* AnyRelay.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7ADCB2D72B6A6EB300C88F89 /* AnyRelay.swift */; };
 		7ADCB2DA2B6A730400C88F89 /* IPOverrideRepositoryStub.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7ADCB2D92B6A730400C88F89 /* IPOverrideRepositoryStub.swift */; };
 		7AE044BB2A935726003915D8 /* Routing.h in Headers */ = {isa = PBXBuildFile; fileRef = 7A88DCD02A8FABBE00D2FF0E /* Routing.h */; settings = {ATTRIBUTES = (Public, ); }; };
@@ -2214,8 +2212,6 @@
 		7AD63A3A2CD5278900445268 /* ObfuscationMethodSelector.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ObfuscationMethodSelector.swift; sourceTree = "<group>"; };
 		7AD63A3C2CD9065100445268 /* RelayObfuscatorTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = RelayObfuscatorTests.swift; sourceTree = "<group>"; };
 		7AD63A3E2CDA53E900445268 /* ObfuscationMethodSelectorTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ObfuscationMethodSelectorTests.swift; sourceTree = "<group>"; };
-		7AD63A432CDA662900445268 /* UInt+Counting.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "UInt+Counting.swift"; sourceTree = "<group>"; };
-		7AD63A462CDA665A00445268 /* UIntTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = UIntTests.swift; sourceTree = "<group>"; };
 		7ADCB2D72B6A6EB300C88F89 /* AnyRelay.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AnyRelay.swift; sourceTree = "<group>"; };
 		7ADCB2D92B6A730400C88F89 /* IPOverrideRepositoryStub.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = IPOverrideRepositoryStub.swift; sourceTree = "<group>"; };
 		7AE90B672C2D726000375A60 /* NSParagraphStyle+Extensions.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "NSParagraphStyle+Extensions.swift"; sourceTree = "<group>"; };
@@ -2775,7 +2771,6 @@
 			children = (
 				F06045F02B2324DA00B2D37A /* ApiHandlers */,
 				062B45A228FD4C0F00746E77 /* Assets */,
-				7AD63A422CDA661B00445268 /* Extensions */,
 				582FFA82290A84E700895745 /* Info.plist */,
 				7A2C0E872D82E450003D8048 /* MullvadAPI */,
 				06799ABE28F98E1D00ACD94E /* MullvadREST.h */,
@@ -2825,7 +2820,6 @@
 			isa = PBXGroup;
 			children = (
 				F0F146932D9462BA00BF78E7 /* MullvadApi */,
-				7AD63A452CDA665200445268 /* Extensions */,
 				F072D3D02C071A9100906F64 /* Shadowsocks */,
 				440E9EF42BDA943B00B1FD11 /* ApiHandlers */,
 				440E9EF52BDA954000B1FD11 /* Relay */,
@@ -4431,22 +4425,6 @@
 			path = SelectLocation;
 			sourceTree = "<group>";
 		};
-		7AD63A422CDA661B00445268 /* Extensions */ = {
-			isa = PBXGroup;
-			children = (
-				7AD63A432CDA662900445268 /* UInt+Counting.swift */,
-			);
-			path = Extensions;
-			sourceTree = "<group>";
-		};
-		7AD63A452CDA665200445268 /* Extensions */ = {
-			isa = PBXGroup;
-			children = (
-				7AD63A462CDA665A00445268 /* UIntTests.swift */,
-			);
-			path = Extensions;
-			sourceTree = "<group>";
-		};
 		7AE241492C20682B0076CE33 /* Presentation controllers */ = {
 			isa = PBXGroup;
 			children = (
@@ -5920,7 +5898,6 @@
 				7A2E7B702D6C9FCF009EF2C3 /* APITransport.swift in Sources */,
 				7A2E7B752D6CA0B1009EF2C3 /* APITransportProvider.swift in Sources */,
 				F0B894F12BF751E300817A42 /* RelayWithDistance.swift in Sources */,
-				7AD63A442CDA663300445268 /* UInt+Counting.swift in Sources */,
 				7AA564502F30D0DE001D1FB9 /* LwoObfuscator.swift in Sources */,
 				7A516C3A2B7111A700BBD33D /* IPOverrideWrapper.swift in Sources */,
 				7A964BB92E699A3F00C6A4EC /* ShadowsocksObfuscator.swift in Sources */,
@@ -6061,7 +6038,6 @@
 				A9A5FA0E2ACB05160083449F /* StorePaymentObserver.swift in Sources */,
 				7A6811542DC8EC6E009CB61A /* UIFont+Weight.swift in Sources */,
 				A9A5FA102ACB05160083449F /* PacketTunnelAPITransport.swift in Sources */,
-				7AD63A472CDA666100445268 /* UIntTests.swift in Sources */,
 				A9A5FA112ACB05160083449F /* APITransportMonitor.swift in Sources */,
 				A9B6AC1A2ADE8FBB00F7802A /* InMemorySettingsStore.swift in Sources */,
 				A9A5FA132ACB05160083449F /* LoadTunnelConfigurationOperation.swift in Sources */,
diff --git a/ios/MullvadVPNTests/MullvadREST/Extensions/UIntTests.swift b/ios/MullvadVPNTests/MullvadREST/Extensions/UIntTests.swift
deleted file mode 100644
index 2f581d227e..0000000000
--- a/ios/MullvadVPNTests/MullvadREST/Extensions/UIntTests.swift
+++ /dev/null
@@ -1,32 +0,0 @@
-//
-//  UIntTests.swift
-//  MullvadVPN
-//
-//  Created by Jon Petersson on 2024-11-05.
-//  Copyright © 2025 Mullvad VPN AB. All rights reserved.
-//
-
-import XCTest
-
-@testable import MullvadREST
-
-class UIntTests: XCTestCase {
-    func testCountingSets() {
-        for setSize in UInt(1)..<20 {
-            let sampleSize: UInt = (setSize * 2) - 1
-
-            var count: UInt = 0
-            (UInt(0)...sampleSize).forEach { index in
-                count = count == setSize ? 1 : count + 1
-
-                let lowerHalfCount = count - 1
-                let upperHalfCount = lowerHalfCount + setSize
-
-                XCTAssertEqual(
-                    index.isOrdered(nth: count, forEverySetOf: setSize),
-                    index == lowerHalfCount || index == upperHalfCount
-                )
-            }
-        }
-    }
-}
diff --git a/ios/MullvadVPNTests/MullvadREST/Relay/ObfuscationMethodSelectorTests.swift b/ios/MullvadVPNTests/MullvadREST/Relay/ObfuscationMethodSelectorTests.swift
index ea073ee961..fb13ea2193 100644
--- a/ios/MullvadVPNTests/MullvadREST/Relay/ObfuscationMethodSelectorTests.swift
+++ b/ios/MullvadVPNTests/MullvadREST/Relay/ObfuscationMethodSelectorTests.swift
@@ -30,7 +30,9 @@ class ObfuscationMethodSelectorTests: XCTestCase {
                 connectionAttemptCount: attempt,
                 tunnelSettings: tunnelSettings, obfuscationBypass: IdentityObfuscationProvider()
             )
-            if attempt.isOrdered(nth: 1, forEverySetOf: 4) {
+
+            let attemptIndex = Int(attempt) % ObfuscationMethodSelector.obfuscationOrder.count
+            if attemptIndex == 0 {
                 XCTAssertEqual(method, .off)
             } else {
                 XCTAssertNotEqual(method, .off)
@@ -54,7 +56,9 @@ class ObfuscationMethodSelectorTests: XCTestCase {
                 connectionAttemptCount: attempt,
                 tunnelSettings: tunnelSettings, obfuscationBypass: IdentityObfuscationProvider()
             )
-            if attempt.isOrdered(nth: 2, forEverySetOf: 4) {
+
+            let attemptIndex = Int(attempt) % ObfuscationMethodSelector.obfuscationOrder.count
+            if attemptIndex == 1 {
                 XCTAssertEqual(method, .shadowsocks)
             } else {
                 XCTAssertNotEqual(method, .shadowsocks)
@@ -78,7 +82,9 @@ class ObfuscationMethodSelectorTests: XCTestCase {
                 connectionAttemptCount: attempt,
                 tunnelSettings: tunnelSettings, obfuscationBypass: IdentityObfuscationProvider()
             )
-            if attempt.isOrdered(nth: 3, forEverySetOf: 4) {
+
+            let attemptIndex = Int(attempt) % ObfuscationMethodSelector.obfuscationOrder.count
+            if attemptIndex == 2 {
                 XCTAssertEqual(method, .quic)
             } else {
                 XCTAssertNotEqual(method, .quic)
@@ -102,11 +108,39 @@ class ObfuscationMethodSelectorTests: XCTestCase {
                 connectionAttemptCount: attempt,
                 tunnelSettings: tunnelSettings, obfuscationBypass: IdentityObfuscationProvider()
             )
-            if attempt.isOrdered(nth: 4, forEverySetOf: 4) {
+
+            let attemptIndex = Int(attempt) % ObfuscationMethodSelector.obfuscationOrder.count
+            if attemptIndex == 3 {
                 XCTAssertEqual(method, .udpOverTcp)
             } else {
                 XCTAssertNotEqual(method, .udpOverTcp)
             }
         }
     }
+
+    func testMethodSelectionLwo() throws {
+        (UInt(0)...10).forEach { attempt in
+            tunnelSettings.wireGuardObfuscation = WireGuardObfuscationSettings(state: .lwo)
+
+            var method = ObfuscationMethodSelector.obfuscationMethodBy(
+                connectionAttemptCount: attempt,
+                tunnelSettings: tunnelSettings, obfuscationBypass: IdentityObfuscationProvider()
+            )
+            XCTAssertEqual(method, .lwo)
+
+            tunnelSettings.wireGuardObfuscation = WireGuardObfuscationSettings(state: .automatic)
+
+            method = ObfuscationMethodSelector.obfuscationMethodBy(
+                connectionAttemptCount: attempt,
+                tunnelSettings: tunnelSettings, obfuscationBypass: IdentityObfuscationProvider()
+            )
+
+            let attemptIndex = Int(attempt) % ObfuscationMethodSelector.obfuscationOrder.count
+            if attemptIndex == 4 {
+                XCTAssertEqual(method, .lwo)
+            } else {
+                XCTAssertNotEqual(method, .lwo)
+            }
+        }
+    }
 }
diff --git a/ios/MullvadVPNUITests/Pages/TunnelControlPage.swift b/ios/MullvadVPNUITests/Pages/TunnelControlPage.swift
index 1cd540ebcd..d03d934a71 100644
--- a/ios/MullvadVPNUITests/Pages/TunnelControlPage.swift
+++ b/ios/MullvadVPNUITests/Pages/TunnelControlPage.swift
@@ -130,12 +130,12 @@ class TunnelControlPage: Page {
     /// Verify that the app attempts to connect over UDP before switching to TCP. For testing blocked UDP traffic.
     @discardableResult func verifyConnectingOverTCPAfterUDPAttempts() -> Self {
         // Number of connection attempts should be equal to the number of obfuscation methods (incl. "off").
-        let connectionAttempts = waitForConnectionAttempts(4, timeout: 80)
+        let connectionAttempts = waitForConnectionAttempts(5, timeout: 80)
 
         // Should do four connection attempts but due to UI bug sometimes only two are displayed, sometimes all three
-        if connectionAttempts.count == 4 {  // Expected retries flow
+        if connectionAttempts.count == 5 {  // Expected retries flow
             for (attemptIndex, attempt) in connectionAttempts.enumerated() {
-                if attemptIndex < 3 {
+                if attemptIndex != 3 {
                     XCTAssertEqual(attempt.protocolName, "UDP")
                 } else if attemptIndex == 3 {
                     XCTAssertEqual(attempt.protocolName, "TCP")
@@ -143,9 +143,9 @@ class TunnelControlPage: Page {
                     XCTFail("Unexpected connection attempt")
                 }
             }
-        } else if connectionAttempts.count == 3 {  // Most of the times this incorrect flow is shown
+        } else if connectionAttempts.count == 4 {  // Most of the times this incorrect flow is shown
             for (attemptIndex, attempt) in connectionAttempts.enumerated() {
-                if attemptIndex < 2 {
+                if attemptIndex != 2 {
                     XCTAssertEqual(attempt.protocolName, "UDP")
                 } else if attemptIndex == 2 {
                     XCTAssertEqual(attempt.protocolName, "TCP")
@@ -154,7 +154,7 @@ class TunnelControlPage: Page {
                 }
             }
         } else {
-            XCTFail("Unexpected number of connection attempts, expected 3~4, got \(connectionAttempts.count)")
+            XCTFail("Unexpected number of connection attempts, expected 4~5, got \(connectionAttempts.count)")
         }
 
         return self
@@ -163,9 +163,9 @@ class TunnelControlPage: Page {
     /// Verify that connection attempts are made in the correct order
     @discardableResult func verifyConnectionAttemptsOrder() -> Self {
         // Number of connection attempts should be equal to the number of obfuscation methods (incl. "off").
-        var connectionAttempts = waitForConnectionAttempts(4, timeout: 80)
+        var connectionAttempts = waitForConnectionAttempts(5, timeout: 80)
         var totalAttemptsOffset = 0
-        XCTAssertEqual(connectionAttempts.count, 4)
+        XCTAssertEqual(connectionAttempts.count, 5)
 
         /// Sometimes, the UI will only show an IP address for the first connection attempt, which gets skipped by
         /// `waitForConnectionAttempts`, and offsets expected attempts count by 1, but still counts towards
@@ -176,10 +176,10 @@ class TunnelControlPage: Page {
             totalAttemptsOffset = 1
         }
         for (attemptIndex, attempt) in connectionAttempts.enumerated() {
-            if attemptIndex < 3 - totalAttemptsOffset {
-                XCTAssertEqual(attempt.protocolName, "UDP")
-            } else {
+            if attemptIndex == 3 - totalAttemptsOffset {
                 XCTAssertEqual(attempt.protocolName, "TCP")
+            } else {
+                XCTAssertEqual(attempt.protocolName, "UDP")
             }
         }
 
diff --git a/ios/PacketTunnelCoreTests/ProtocolObfuscatorTests.swift b/ios/PacketTunnelCoreTests/ProtocolObfuscatorTests.swift
index 3aaca43689..8c8cf37608 100644
--- a/ios/PacketTunnelCoreTests/ProtocolObfuscatorTests.swift
+++ b/ios/PacketTunnelCoreTests/ProtocolObfuscatorTests.swift
@@ -66,6 +66,13 @@ final class ProtocolObfuscatorTests: XCTestCase {
         validate(obfuscated.endpoint, against: obfuscationProtocol)
     }
 
+    func testObfuscateLwo() throws {
+        let endpoint = try makeEndpoint(obfuscation: .lwo)
+        let obfuscated = obfuscator.obfuscate(endpoint)
+        let obfuscationProtocol = try XCTUnwrap(obfuscator.tunnelObfuscator as? TunnelObfuscationStub)
+
+        validate(obfuscated.endpoint, against: obfuscationProtocol)
+    }
 }
 
 extension ProtocolObfuscatorTests {