Refactor UnsafeListener and add a test for Shadowsocks Obfuscation

ios/MullvadRustRuntimeTests/TCPConnection.swift

@@ -11,14 +11,16 @@ import Network
 
 /// Minimal implementation of TCP connection capable of receiving data.
 /// > Warning: Do not use this implementation in production code. See the warning in `start()`.
-class TCPConnection {
+class TCPConnection: Connection {
     private let dispatchQueue = DispatchQueue(label: "TCPConnection")
     private let nwConnection: NWConnection
 
-    init(nwConnection: NWConnection) {
+    required init(nwConnection: NWConnection) {
         self.nwConnection = nwConnection
     }
 
+    static var connectionParameters: NWParameters { .tcp }
+
     deinit {
         cancel()
     }

ios/MullvadRustRuntimeTests/TunnelObfuscationTests.swift

@@ -11,16 +11,20 @@ import Network
 import XCTest
 
 final class TunnelObfuscationTests: XCTestCase {
-    func testRunningObfuscatorProxy() async throws {
+    func testRunningUdpOverTcpObfuscatorProxy() async throws {
         // Each packet is prefixed with u16 that contains a payload length.
         let preambleLength = MemoryLayout<UInt16>.size
         let markerData = Data([109, 117, 108, 108, 118, 97, 100])
         let packetLength = markerData.count + preambleLength
 
-        let tcpListener = try TCPUnsafeListener()
+        let tcpListener = try UnsafeListener<TCPConnection>()
         try await tcpListener.start()
 
-        let obfuscator = TunnelObfuscator(remoteAddress: IPv4Address.loopback, tcpPort: tcpListener.listenPort)
+        let obfuscator = TunnelObfuscator(
+            remoteAddress: IPv4Address.loopback,
+            tcpPort: tcpListener.listenPort,
+            obfuscationProtocol: .udpOverTcp
+        )
         obfuscator.start()
 
         // Accept incoming connections
@@ -46,4 +50,42 @@ final class TunnelObfuscationTests: XCTestCase {
         XCTAssert(receivedData.count == packetLength)
         XCTAssertEqual(receivedData[preambleLength...], markerData)
     }
+
+    func testRunningShadowsocksObfuscatorProxy() async throws {
+        let markerData = Data([109, 117, 108, 108, 118, 97, 100])
+
+        let localUdpListener = try UnsafeListener<UDPConnection>()
+        try await localUdpListener.start()
+
+        let localObfuscator = TunnelObfuscator(
+            remoteAddress: IPv4Address.loopback,
+            tcpPort: localUdpListener.listenPort,
+            obfuscationProtocol: .shadowsocks
+        )
+        localObfuscator.start()
+
+        // Accept incoming connections
+        let localConnectionDataTask = Task {
+            for await obfuscatedConnection in localUdpListener.newConnections {
+                try await obfuscatedConnection.start()
+
+                let readDatagram = try await obfuscatedConnection.readSingleDatagram()
+                // Write into the connection the unencrypted payload that was just read
+                try await obfuscatedConnection.sendData(readDatagram)
+                return readDatagram
+            }
+            throw POSIXError(.ECANCELED)
+        }
+
+        // Send marker data over UDP
+        let connection = UDPConnection(remote: IPv4Address.loopback, port: localObfuscator.localUdpPort)
+        try await connection.start()
+        try await connection.sendData(markerData)
+        let readDataFromObfuscator = try await connection.readSingleDatagram()
+
+        // As the connection data is encrypted and this test does not run a shadowsocks server to decrypt the payload
+        // The connection from the local UDP listener writes back what it read from the obfuscator, unencrypted
+        _ = try await localConnectionDataTask.value
+        XCTAssertEqual(readDataFromObfuscator, markerData)
+    }
 }

ios/MullvadRustRuntimeTests/UDPConnection.swift

@@ -9,20 +9,31 @@
 import Foundation
 import Network
 
+protocol Connection {
+    init(nwConnection: NWConnection)
+    static var connectionParameters: NWParameters { get }
+}
+
 /// Minimal implementation of UDP connection capable of sending data.
 /// > Warning: Do not use this implementation in production code. See the warning in `start()`.
-class UDPConnection {
+class UDPConnection: Connection {
     private let dispatchQueue = DispatchQueue(label: "UDPConnection")
     private let nwConnection: NWConnection
 
-    init(remote: IPAddress, port: UInt16) {
-        nwConnection = NWConnection(
+    convenience init(remote: IPAddress, port: UInt16) {
+        self.init(nwConnection: NWConnection(
             host: NWEndpoint.Host("\(remote)"),
             port: NWEndpoint.Port(integerLiteral: port),
             using: .udp
-        )
+        ))
+    }
+
+    required init(nwConnection: NWConnection) {
+        self.nwConnection = nwConnection
     }
 
+    static var connectionParameters: NWParameters { .udp }
+
     deinit {
         cancel()
     }
@@ -58,6 +69,22 @@ class UDPConnection {
         nwConnection.cancel()
     }
 
+    func readSingleDatagram() async throws -> Data {
+        return try await withCheckedThrowingContinuation { continuation in
+            nwConnection.receiveMessage { data, _, _, error in
+                guard let data else {
+                    continuation.resume(throwing: POSIXError(.EIO))
+                    return
+                }
+                if let error {
+                    continuation.resume(throwing: error)
+                    return
+                }
+                continuation.resume(with: .success(data))
+            }
+        }
+    }
+
     func sendData(_ data: Data) async throws {
         return try await withCheckedThrowingContinuation { continuation in
             nwConnection.send(content: data, completion: .contentProcessed { error in

ios/MullvadRustRuntimeTests/TCPUnsafeListener.swift → ios/MullvadRustRuntimeTests/UnsafeListener.swift

@@ -1,5 +1,5 @@
 //
-//  TCPUnsafeListener.swift
+//  UnsafeListener.swift
 //  MullvadRustRuntimeTests
 //
 //  Created by pronebird on 27/06/2023.
@@ -8,25 +8,23 @@
 
 import Network
 
-/// Minimal implementation of a TCP listener.
-/// > Warning: Do not use this implementation in production code. See the warning in `start()`.
-class TCPUnsafeListener {
-    private let dispatchQueue = DispatchQueue(label: "TCPListener")
+class UnsafeListener<T: Connection> {
+    private let dispatchQueue = DispatchQueue(label: "com.test.unsafeListener")
     private let listener: NWListener
 
-    /// A stream of new TCP connections.
-    /// The caller may iterate over this stream to accept new TCP connections.
+    /// A stream of new connections.
+    /// The caller may iterate over this stream to accept new connections.
     ///
-    /// `TCPConnection` objects are returned unopen, so the caller has to call `TCPConnection.start()` to accept the
+    /// `Connection` objects are returned unopen, so the caller has to call `Connection.start()` to accept the
     /// connection before initiating the data exchange.
-    let newConnections: AsyncStream<TCPConnection>
+    let newConnections: AsyncStream<T>
 
     init() throws {
-        let listener = try NWListener(using: .tcp)
+        let listener = try NWListener(using: T.connectionParameters)
 
         newConnections = AsyncStream { continuation in
             listener.newConnectionHandler = { nwConnection in
-                continuation.yield(TCPConnection(nwConnection: nwConnection))
+                continuation.yield(T(nwConnection: nwConnection))
             }
             continuation.onTermination = { @Sendable _ in
                 listener.newConnectionHandler = nil
@@ -40,7 +38,7 @@ class TCPUnsafeListener {
         cancel()
     }
 
-    /// Local TCP port bound by listener on which it accepts new connections.
+    /// Local port bound by listener on which it accepts new connections.
     var listenPort: UInt16 {
         return listener.port?.rawValue ?? 0
     }

ios/MullvadVPN.xcodeproj/project.pbxproj

@@ -838,7 +838,7 @@
 		A9D9A4B22C36D12D004088DD /* TunnelObfuscator.swift in Sources */ = {isa = PBXBuildFile; fileRef = 584023212A406BF5007B27AC /* TunnelObfuscator.swift */; };
 		A9D9A4BB2C36D397004088DD /* EphemeralPeerNegotiator.swift in Sources */ = {isa = PBXBuildFile; fileRef = A9EB4F9C2B7FAB21002A2D7A /* EphemeralPeerNegotiator.swift */; };
 		A9D9A4C42C36D53C004088DD /* MullvadRustRuntime.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = A992DA1D2C24709F00DE7CE5 /* MullvadRustRuntime.framework */; };
-		A9D9A4CC2C36D54E004088DD /* TCPUnsafeListener.swift in Sources */ = {isa = PBXBuildFile; fileRef = 585A02E82A4B283000C6CAFF /* TCPUnsafeListener.swift */; };
+		A9D9A4CC2C36D54E004088DD /* UnsafeListener.swift in Sources */ = {isa = PBXBuildFile; fileRef = 585A02E82A4B283000C6CAFF /* UnsafeListener.swift */; };
 		A9D9A4CD2C36D54E004088DD /* UDPConnection.swift in Sources */ = {isa = PBXBuildFile; fileRef = 585A02EA2A4B285800C6CAFF /* UDPConnection.swift */; };
 		A9D9A4CE2C36D54E004088DD /* TunnelObfuscationTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58695A9F2A4ADA9200328DB3 /* TunnelObfuscationTests.swift */; };
 		A9D9A4CF2C36D54E004088DD /* TCPConnection.swift in Sources */ = {isa = PBXBuildFile; fileRef = 585A02EC2A4B28F300C6CAFF /* TCPConnection.swift */; };
@@ -1519,7 +1519,7 @@
 		584D26C5270C8741004EA533 /* SettingsDNSTextCell.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SettingsDNSTextCell.swift; sourceTree = "<group>"; };
 		58561C98239A5D1500BD6B5E /* IPv4Endpoint.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = IPv4Endpoint.swift; sourceTree = "<group>"; };
 		5859A55429CD9DD800F66591 /* changes.txt */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = text; path = changes.txt; sourceTree = "<group>"; };
-		585A02E82A4B283000C6CAFF /* TCPUnsafeListener.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TCPUnsafeListener.swift; sourceTree = "<group>"; };
+		585A02E82A4B283000C6CAFF /* UnsafeListener.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = UnsafeListener.swift; sourceTree = "<group>"; };
 		585A02EA2A4B285800C6CAFF /* UDPConnection.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = UDPConnection.swift; sourceTree = "<group>"; };
 		585A02EC2A4B28F300C6CAFF /* TCPConnection.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TCPConnection.swift; sourceTree = "<group>"; };
 		585CA70E25F8C44600B47C62 /* UIMetrics.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = UIMetrics.swift; sourceTree = "<group>"; };
@@ -4067,12 +4067,12 @@
 		A9D9A4C12C36D53C004088DD /* MullvadRustRuntimeTests */ = {
 			isa = PBXGroup;
 			children = (
-				A9C308392C19DDA7008715F1 /* MullvadPostQuantum+Stubs.swift */,
 				A98F1B502C19C48D003C869E /* EphemeralPeerExchangeActorTests.swift */,
+				A9C308392C19DDA7008715F1 /* MullvadPostQuantum+Stubs.swift */,
 				585A02EC2A4B28F300C6CAFF /* TCPConnection.swift */,
-				585A02E82A4B283000C6CAFF /* TCPUnsafeListener.swift */,
 				58695A9F2A4ADA9200328DB3 /* TunnelObfuscationTests.swift */,
 				585A02EA2A4B285800C6CAFF /* UDPConnection.swift */,
+				585A02E82A4B283000C6CAFF /* UnsafeListener.swift */,
 			);
 			path = MullvadRustRuntimeTests;
 			sourceTree = "<group>";
@@ -6207,7 +6207,7 @@
 				F08B6B772C52878400D0A121 /* EphemeralPeerExchangeActorTests.swift in Sources */,
 				A9D9A4CF2C36D54E004088DD /* TCPConnection.swift in Sources */,
 				A9D9A4CE2C36D54E004088DD /* TunnelObfuscationTests.swift in Sources */,
-				A9D9A4CC2C36D54E004088DD /* TCPUnsafeListener.swift in Sources */,
+				A9D9A4CC2C36D54E004088DD /* UnsafeListener.swift in Sources */,
 				A9D9A4CD2C36D54E004088DD /* UDPConnection.swift in Sources */,
 				F0A1638A2C47B77300592300 /* ServerRelaysResponse+Stubs.swift in Sources */,
 			);

ios/PacketTunnelCore/Actor/ProtocolObfuscator.swift

@@ -58,7 +58,8 @@ public class ProtocolObfuscator<Obfuscator: TunnelObfuscation>: ProtocolObfuscat
         }
         let obfuscator = Obfuscator(
             remoteAddress: obfuscatedEndpoint.ipv4Relay.ip,
-            tcpPort: tcpPort.portValue ?? 0
+            tcpPort: tcpPort.portValue ?? 0,
+            obfuscationProtocol: .udpOverTcp
         )
         remotePort = tcpPort.portValue ?? 0
         obfuscator.start()

ios/PacketTunnelCoreTests/Mocks/TunnelObfuscationStub.swift

@@ -15,7 +15,7 @@ struct TunnelObfuscationStub: TunnelObfuscation {
     var transportLayer: TransportLayer { .udp }
 
     let remotePort: UInt16
-    init(remoteAddress: IPAddress, tcpPort: UInt16) {
+    init(remoteAddress: IPAddress, tcpPort: UInt16, obfuscationProtocol: TunnelObfuscationProtocol) {
         remotePort = tcpPort
     }