Add LWO obfuscator to MullvadRustRuntimeadd-lwo-obfuscator-to-mullvad-ios-ios-1451

7 files changed, 176 insertions, 19 deletions

diff --git a/ios/MullvadRustRuntime/TunnelObfuscator.swift b/ios/MullvadRustRuntime/TunnelObfuscator.swift
index 6f74089265..8254d8ee7c 100644
--- a/ios/MullvadRustRuntime/TunnelObfuscator.swift
+++ b/ios/MullvadRustRuntime/TunnelObfuscator.swift
@@ -10,15 +10,22 @@ import Foundation
 import MullvadRustRuntimeProxy
 import MullvadTypes
 import Network
+import WireGuardKitTypes
 
 public enum TunnelObfuscationProtocol {
     case udpOverTcp
     case shadowsocks
     case quic(hostname: String, token: String)
+    case lwo(serverPublicKey: PublicKey)
 }
 
 public protocol TunnelObfuscation {
-    init(remoteAddress: IPAddress, tcpPort: UInt16, obfuscationProtocol: TunnelObfuscationProtocol)
+    init(
+        remoteAddress: IPAddress,
+        tcpPort: UInt16,
+        obfuscationProtocol: TunnelObfuscationProtocol,
+        clientPublicKey: PublicKey
+    )
     func start()
     func stop()
     var localUdpPort: UInt16 { get }
@@ -36,6 +43,7 @@ public final class TunnelObfuscator: TunnelObfuscation {
     private let remoteAddress: IPAddress
     internal let tcpPort: UInt16
     internal let obfuscationProtocol: TunnelObfuscationProtocol
+    private let clientPublicKey: PublicKey
 
     private var proxyHandle = ProxyHandle(context: nil, port: 0)
     private var isStarted = false
@@ -56,14 +64,22 @@ public final class TunnelObfuscator: TunnelObfuscation {
             .udp
         case .quic:
             .udp
+        case .lwo:
+            .udp
         }
     }
 
     /// Initialize tunnel obfuscator with remote server address and TCP port where udp2tcp is running.
-    public init(remoteAddress: IPAddress, tcpPort: UInt16, obfuscationProtocol: TunnelObfuscationProtocol) {
+    public init(
+        remoteAddress: IPAddress,
+        tcpPort: UInt16,
+        obfuscationProtocol: TunnelObfuscationProtocol,
+        clientPublicKey: PublicKey
+    ) {
         self.remoteAddress = remoteAddress
         self.tcpPort = tcpPort
         self.obfuscationProtocol = obfuscationProtocol
+        self.clientPublicKey = clientPublicKey
     }
 
     deinit {
@@ -101,6 +117,19 @@ public final class TunnelObfuscator: TunnelObfuscation {
                         token,
                         proxyHandlePointer
                     )
+                case let .lwo(serverPublicKey):
+                    clientPublicKey.rawValue.withUnsafeBytes { clientKeyPtr in
+                        serverPublicKey.rawValue.withUnsafeBytes { serverKeyPtr in
+                            start_lwo_obfuscator_proxy(
+                                addressData.map { $0 },
+                                UInt(addressData.count),
+                                tcpPort,
+                                clientKeyPtr.baseAddress?.assumingMemoryBound(to: UInt8.self),
+                                serverKeyPtr.baseAddress?.assumingMemoryBound(to: UInt8.self),
+                                proxyHandlePointer
+                            )
+                        }
+                    }
                 }
             }
 
diff --git a/ios/MullvadRustRuntimeTests/TunnelObfuscationTests.swift b/ios/MullvadRustRuntimeTests/TunnelObfuscationTests.swift
index 0d6ea2a871..07d03e7f52 100644
--- a/ios/MullvadRustRuntimeTests/TunnelObfuscationTests.swift
+++ b/ios/MullvadRustRuntimeTests/TunnelObfuscationTests.swift
@@ -8,9 +8,13 @@
 
 import MullvadRustRuntime
 import Network
+import WireGuardKitTypes
 import XCTest
 
 final class TunnelObfuscationTests: XCTestCase {
+    /// A test public key generated from a random private key.
+    private let testPublicKey = PrivateKey().publicKey
+
     func testRunningUdpOverTcpObfuscatorProxy() async throws {
         // Each packet is prefixed with u16 that contains a payload length.
         let preambleLength = MemoryLayout<UInt16>.size
@@ -23,7 +27,8 @@ final class TunnelObfuscationTests: XCTestCase {
         let obfuscator = TunnelObfuscator(
             remoteAddress: IPv4Address.loopback,
             tcpPort: tcpListener.listenPort,
-            obfuscationProtocol: .udpOverTcp
+            obfuscationProtocol: .udpOverTcp,
+            clientPublicKey: testPublicKey
         )
         obfuscator.start()
 
@@ -60,7 +65,8 @@ final class TunnelObfuscationTests: XCTestCase {
         let localObfuscator = TunnelObfuscator(
             remoteAddress: IPv4Address.loopback,
             tcpPort: localUdpListener.listenPort,
-            obfuscationProtocol: .shadowsocks
+            obfuscationProtocol: .shadowsocks,
+            clientPublicKey: testPublicKey
         )
         localObfuscator.start()
 
@@ -88,4 +94,96 @@ final class TunnelObfuscationTests: XCTestCase {
         _ = try await localConnectionDataTask.value
         XCTAssertEqual(readDataFromObfuscator, markerData)
     }
+
+    /// Tests that the LWO obfuscator proxy can be started and stopped correctly via FFI.
+    func testRunningLwoObfuscatorProxy() async throws {
+        let markerData = Data([109, 117, 108, 108, 118, 97, 100])
+
+        let localUdpListener = try UnsafeListener<UDPConnection>()
+        try await localUdpListener.start()
+
+        // Generate test keys for LWO
+        let clientPrivateKey = PrivateKey()
+        let serverPrivateKey = PrivateKey()
+
+        let obfuscator = TunnelObfuscator(
+            remoteAddress: IPv4Address.loopback,
+            tcpPort: localUdpListener.listenPort,
+            obfuscationProtocol: .lwo(serverPublicKey: serverPrivateKey.publicKey),
+            clientPublicKey: clientPrivateKey.publicKey
+        )
+        obfuscator.start()
+
+        // Verify the obfuscator has allocated a local UDP port
+        XCTAssertNotEqual(obfuscator.localUdpPort, 0, "LWO obfuscator should allocate a local UDP port")
+
+        // Accept incoming connections and echo back
+        let connectionDataTask = Task {
+            for await connection in localUdpListener.newConnections {
+                try await connection.start()
+                let readDatagram = try await connection.readSingleDatagram()
+                try await connection.sendData(readDatagram)
+                return readDatagram
+            }
+            throw POSIXError(.ECANCELED)
+        }
+
+        // Send marker data over UDP to the obfuscator's local port
+        let connection = UDPConnection(remote: IPv4Address.loopback, port: obfuscator.localUdpPort)
+        try await connection.start()
+        try await connection.sendData(markerData)
+
+        // Wait for the data to be received by the listener
+        // The data will be obfuscated, so we just verify something was received
+        let receivedData = try await connectionDataTask.value
+        XCTAssertFalse(receivedData.isEmpty, "LWO obfuscator should forward data to the remote endpoint")
+
+        // Stop the obfuscator - this tests that stop() doesn't crash or cause memory issues
+        obfuscator.stop()
+    }
+
+    /// Tests that the LWO obfuscator can be started and stopped multiple times without memory issues.
+    func testLwoObfuscatorStartStopMultipleTimes() async throws {
+        let localUdpListener = try UnsafeListener<UDPConnection>()
+        try await localUdpListener.start()
+
+        let clientPrivateKey = PrivateKey()
+        let serverPrivateKey = PrivateKey()
+
+        // Create and destroy obfuscators multiple times to check for memory leaks/issues
+        for iteration in 1...5 {
+            let obfuscator = TunnelObfuscator(
+                remoteAddress: IPv4Address.loopback,
+                tcpPort: localUdpListener.listenPort,
+                obfuscationProtocol: .lwo(serverPublicKey: serverPrivateKey.publicKey),
+                clientPublicKey: clientPrivateKey.publicKey
+            )
+
+            obfuscator.start()
+            XCTAssertNotEqual(obfuscator.localUdpPort, 0, "Iteration \(iteration): LWO obfuscator should allocate a port")
+
+            obfuscator.stop()
+        }
+    }
+
+    /// Tests that calling stop() on an already stopped obfuscator doesn't crash.
+    func testLwoObfuscatorDoubleStopSafe() async throws {
+        let localUdpListener = try UnsafeListener<UDPConnection>()
+        try await localUdpListener.start()
+
+        let clientPrivateKey = PrivateKey()
+        let serverPrivateKey = PrivateKey()
+
+        let obfuscator = TunnelObfuscator(
+            remoteAddress: IPv4Address.loopback,
+            tcpPort: localUdpListener.listenPort,
+            obfuscationProtocol: .lwo(serverPublicKey: serverPrivateKey.publicKey),
+            clientPublicKey: clientPrivateKey.publicKey
+        )
+
+        obfuscator.start()
+        obfuscator.stop()
+        // Second stop should be safe and not crash
+        obfuscator.stop()
+    }
 }
diff --git a/ios/PacketTunnelCore/Actor/PacketTunnelActor.swift b/ios/PacketTunnelCore/Actor/PacketTunnelActor.swift
index f043b5501d..65dfbbf7c6 100644
--- a/ios/PacketTunnelCore/Actor/PacketTunnelActor.swift
+++ b/ios/PacketTunnelCore/Actor/PacketTunnelActor.swift
@@ -447,7 +447,8 @@ extension PacketTunnelActor {
             connectionState.connectedEndpoint,
             relayFeatures: connectionState.selectedRelays.entry?.features
                 ?? connectionState.selectedRelays.exit
-                .features, obfuscationMethod: connectionState.obfuscationMethod
+                .features, obfuscationMethod: connectionState.obfuscationMethod,
+            clientPublicKey: settings.privateKey.publicKey
         )
         let transportLayer = protocolObfuscator.transportLayer.map { $0 } ?? .udp
 
diff --git a/ios/PacketTunnelCore/Actor/ProtocolObfuscator.swift b/ios/PacketTunnelCore/Actor/ProtocolObfuscator.swift
index 192fd57b16..5ffae5d051 100644
--- a/ios/PacketTunnelCore/Actor/ProtocolObfuscator.swift
+++ b/ios/PacketTunnelCore/Actor/ProtocolObfuscator.swift
@@ -11,6 +11,7 @@ import MullvadREST
 import MullvadRustRuntime
 import MullvadSettings
 import MullvadTypes
+import WireGuardKitTypes
 
 public struct ProtocolObfuscationResult {
     let endpoint: MullvadEndpoint
@@ -21,7 +22,8 @@ public protocol ProtocolObfuscation {
     func obfuscate(
         _ endpoint: MullvadEndpoint,
         relayFeatures: REST.ServerRelay.Features?,
-        obfuscationMethod: WireGuardObfuscationState
+        obfuscationMethod: WireGuardObfuscationState,
+        clientPublicKey: PublicKey
     ) -> ProtocolObfuscationResult
     var transportLayer: TransportLayer? { get }
     var remotePort: UInt16 { get }
@@ -46,7 +48,8 @@ public class ProtocolObfuscator<Obfuscator: TunnelObfuscation>: ProtocolObfuscat
     public func obfuscate(
         _ endpoint: MullvadEndpoint,
         relayFeatures: REST.ServerRelay.Features?,
-        obfuscationMethod: WireGuardObfuscationState
+        obfuscationMethod: WireGuardObfuscationState,
+        clientPublicKey: PublicKey
     ) -> ProtocolObfuscationResult {
         remotePort = endpoint.ipv4Relay.port
 
@@ -74,7 +77,8 @@ public class ProtocolObfuscator<Obfuscator: TunnelObfuscation>: ProtocolObfuscat
         let obfuscator = Obfuscator(
             remoteAddress: endpoint.ipv4Relay.ip,
             tcpPort: remotePort,
-            obfuscationProtocol: obfuscationProtocol
+            obfuscationProtocol: obfuscationProtocol,
+            clientPublicKey: clientPublicKey
         )
 
         obfuscator.start()
diff --git a/ios/PacketTunnelCoreTests/Mocks/ProtocolObfuscationStub.swift b/ios/PacketTunnelCoreTests/Mocks/ProtocolObfuscationStub.swift
index acbd721ca5..4f8a08c400 100644
--- a/ios/PacketTunnelCoreTests/Mocks/ProtocolObfuscationStub.swift
+++ b/ios/PacketTunnelCoreTests/Mocks/ProtocolObfuscationStub.swift
@@ -6,6 +6,8 @@
 //  Copyright © 2025 Mullvad VPN AB. All rights reserved.
 //
 
+import WireGuardKitTypes
+
 @testable import MullvadREST
 @testable import MullvadSettings
 @testable import MullvadTypes
@@ -17,7 +19,8 @@ struct ProtocolObfuscationStub: ProtocolObfuscation {
     func obfuscate(
         _ endpoint: MullvadEndpoint,
         relayFeatures: REST.ServerRelay.Features?,
-        obfuscationMethod: WireGuardObfuscationState
+        obfuscationMethod: WireGuardObfuscationState,
+        clientPublicKey: PublicKey
     ) -> ProtocolObfuscationResult {
         .init(endpoint: endpoint, method: .off)
     }
diff --git a/ios/PacketTunnelCoreTests/Mocks/TunnelObfuscationStub.swift b/ios/PacketTunnelCoreTests/Mocks/TunnelObfuscationStub.swift
index 7562114f35..48909105cc 100644
--- a/ios/PacketTunnelCoreTests/Mocks/TunnelObfuscationStub.swift
+++ b/ios/PacketTunnelCoreTests/Mocks/TunnelObfuscationStub.swift
@@ -8,6 +8,7 @@
 
 import Foundation
 import Network
+import WireGuardKitTypes
 
 @testable import MullvadRustRuntime
 @testable import MullvadTypes
@@ -16,7 +17,12 @@ struct TunnelObfuscationStub: TunnelObfuscation {
     var transportLayer: TransportLayer { .udp }
 
     let remotePort: UInt16
-    init(remoteAddress: IPAddress, tcpPort: UInt16, obfuscationProtocol: TunnelObfuscationProtocol) {
+    init(
+        remoteAddress: IPAddress,
+        tcpPort: UInt16,
+        obfuscationProtocol: TunnelObfuscationProtocol,
+        clientPublicKey: PublicKey
+    ) {
         remotePort = tcpPort
     }
 
diff --git a/ios/PacketTunnelCoreTests/ProtocolObfuscatorTests.swift b/ios/PacketTunnelCoreTests/ProtocolObfuscatorTests.swift
index bacdddf5c4..d7ddf55d81 100644
--- a/ios/PacketTunnelCoreTests/ProtocolObfuscatorTests.swift
+++ b/ios/PacketTunnelCoreTests/ProtocolObfuscatorTests.swift
@@ -7,6 +7,7 @@
 //
 
 import Network
+import WireGuardKitTypes
 import XCTest
 
 @testable import MullvadREST
@@ -17,6 +18,7 @@ import XCTest
 final class ProtocolObfuscatorTests: XCTestCase {
     var obfuscator: ProtocolObfuscator<TunnelObfuscationStub>!
     var endpoint: MullvadEndpoint!
+    let testPublicKey = PrivateKey().publicKey
 
     override func setUpWithError() throws {
         let address = try XCTUnwrap(IPv4Address("1.2.3.4"))
@@ -34,33 +36,46 @@ final class ProtocolObfuscatorTests: XCTestCase {
     }
 
     func testObfuscateOffDoesNotChangeEndpoint() {
-        let settings = settings(.off, obfuscationPort: .automatic)
-        let nonObfuscated = obfuscator.obfuscate(endpoint, relayFeatures: nil, obfuscationMethod: .off)
+        let nonObfuscated = obfuscator.obfuscate(
+            endpoint,
+            relayFeatures: nil,
+            obfuscationMethod: .off,
+            clientPublicKey: testPublicKey
+        )
 
         XCTAssertEqual(endpoint, nonObfuscated.endpoint)
     }
 
     func testObfuscateUdpOverTcp() throws {
-        let settings = settings(.udpOverTcp, obfuscationPort: .automatic)
-        let obfuscated = obfuscator.obfuscate(endpoint, relayFeatures: nil, obfuscationMethod: .udpOverTcp)
+        let obfuscated = obfuscator.obfuscate(
+            endpoint,
+            relayFeatures: nil,
+            obfuscationMethod: .udpOverTcp,
+            clientPublicKey: testPublicKey
+        )
         let obfuscationProtocol = try XCTUnwrap(obfuscator.tunnelObfuscator as? TunnelObfuscationStub)
 
         validate(obfuscated.endpoint, against: obfuscationProtocol)
     }
 
     func testObfuscateShadowsocks() throws {
-        let settings = settings(.shadowsocks, obfuscationPort: .automatic)
-        let obfuscated = obfuscator.obfuscate(endpoint, relayFeatures: nil, obfuscationMethod: .shadowsocks)
+        let obfuscated = obfuscator.obfuscate(
+            endpoint,
+            relayFeatures: nil,
+            obfuscationMethod: .shadowsocks,
+            clientPublicKey: testPublicKey
+        )
         let obfuscationProtocol = try XCTUnwrap(obfuscator.tunnelObfuscator as? TunnelObfuscationStub)
 
         validate(obfuscated.endpoint, against: obfuscationProtocol)
     }
 
     func testObfuscateQuic() throws {
-        let settings = settings(.quic, obfuscationPort: .automatic)
         let obfuscated = obfuscator.obfuscate(
             endpoint,
-            relayFeatures: .init(daita: nil, quic: .init(addrIn: [], domain: "", token: "")), obfuscationMethod: .quic
+            relayFeatures: .init(daita: nil, quic: .init(addrIn: [], domain: "", token: "")),
+            obfuscationMethod: .quic,
+            clientPublicKey: testPublicKey
         )
         let obfuscationProtocol = try XCTUnwrap(obfuscator.tunnelObfuscator as? TunnelObfuscationStub)
 
@@ -71,7 +86,8 @@ final class ProtocolObfuscatorTests: XCTestCase {
         let obfuscated = obfuscator.obfuscate(
             endpoint,
             relayFeatures: .init(daita: nil, quic: .init(addrIn: [], domain: "", token: "")),
-            obfuscationMethod: .automatic
+            obfuscationMethod: .automatic,
+            clientPublicKey: testPublicKey
         )
 
         XCTAssertEqual(endpoint, obfuscated.endpoint)