Inline usbmux forwarding and migrate SSH transport

.gitmodules

@@ -40,3 +40,6 @@
 [submodule "vendor/zstd"]
 	path = vendor/zstd
 	url = https://github.com/facebook/zstd.git
+[submodule "vendor/swift-nio-ssh"]
+	path = vendor/swift-nio-ssh
+	url = https://github.com/apple/swift-nio-ssh.git

Package.swift

@@ -19,7 +19,9 @@ let package = Package(
         .package(path: "vendor/swift-subprocess"),
         .package(path: "vendor/swift-trustcache"),
         .package(path: "vendor/SWCompression"),
+        .package(path: "vendor/swift-nio-ssh"),
         .package(path: "vendor/zstd"),
+        .package(url: "https://github.com/apple/swift-nio.git", from: "2.81.0"),
     ],
     targets: [
         .target(
@@ -63,6 +65,9 @@ let package = Package(
                 .product(name: "Capstone", package: "libcapstone-spm"),
                 .product(name: "Dynamic", package: "Dynamic"),
                 .product(name: "libirecovery", package: "AppleMobileDeviceLibrary"),
+                .product(name: "NIOCore", package: "swift-nio"),
+                .product(name: "NIOPosix", package: "swift-nio"),
+                .product(name: "NIOSSH", package: "swift-nio-ssh"),
                 .product(name: "SWCompression", package: "SWCompression"),
                 .product(name: "Subprocess", package: "swift-subprocess"),
                 .product(name: "TrustCache", package: "swift-trustcache"),

sources/vphone-cli/VPhoneCFWInstallCLI.swift

@@ -1,5 +1,6 @@
 import ArgumentParser
 import Foundation
+import Subprocess
 
 struct CFWInstallCLI: AsyncParsableCommand {
     enum Variant: String, ExpressibleByArgument {
@@ -40,7 +41,7 @@ struct CFWInstallCLI: AsyncParsableCommand {
     }
 }
 
-private struct VPhoneCFWInstaller {
+private final class VPhoneCFWInstaller {
     let vmDirectory: URL
     let projectRoot: URL
     let variant: CFWInstallCLI.Variant
@@ -50,12 +51,12 @@ private struct VPhoneCFWInstaller {
     let sshPassword = "alpine"
     let sshUser = "root"
     let sshHost = "localhost"
-    let sshRetry = 3
     let scriptDirectory: URL
     let temporaryDirectory: URL
     let cfwInputDirectory: URL
     let jbInputDirectory: URL
     let patcherBinary: String
+    var sshClient: VPhoneSSHClient?
 
     init(vmDirectory: URL, projectRoot: URL, variant: CFWInstallCLI.Variant, sshPort: Int, skipHalt: Bool) throws {
         self.vmDirectory = vmDirectory
@@ -80,6 +81,13 @@ private struct VPhoneCFWInstaller {
         try await checkPrerequisites()
         try setupInputs()
         try await waitForSSHReady()
+        let sshClient = VPhoneSSHClient(host: sshHost, port: sshPort, username: sshUser, password: sshPassword)
+        try sshClient.connect()
+        self.sshClient = sshClient
+        defer {
+            try? sshClient.shutdown()
+            self.sshClient = nil
+        }
         try FileManager.default.createDirectory(at: temporaryDirectory, withIntermediateDirectories: true)
 
         switch variant {
@@ -101,7 +109,7 @@ private struct VPhoneCFWInstaller {
     }
 
     func checkPrerequisites() async throws {
-        var commands = ["ipsw", "aea", "ldid", patcherBinary, "ssh", "scp"]
+        var commands = ["ipsw", "aea", "ldid", patcherBinary]
         if variant == .jb {
             commands += ["xcrun"]
         }
@@ -156,8 +164,7 @@ private struct VPhoneCFWInstaller {
         print("[*] Waiting for ramdisk SSH on \(sshUser)@\(sshHost):\(sshPort)...")
         var elapsed = 0
         while elapsed < 60 {
-            let result = try await ssh("echo ready", requireSuccess: false)
-            if result.terminationStatus.isSuccess {
+            if VPhoneSSHClient.probe(host: sshHost, port: sshPort, username: sshUser, password: sshPassword) {
                 print("[+] Ramdisk SSH is reachable")
                 return
             }
@@ -268,74 +275,26 @@ private struct VPhoneCFWInstaller {
         return directory
     }
 
-    func sshBaseArguments() -> [String] {
-        [
-            "-o", "StrictHostKeyChecking=no",
-            "-o", "UserKnownHostsFile=/dev/null",
-            "-o", "PreferredAuthentications=password",
-            "-o", "NumberOfPasswordPrompts=1",
-            "-o", "ConnectTimeout=30",
-            "-q",
-            "-p", "\(sshPort)",
-            "\(sshUser)@\(sshHost)",
-        ]
-    }
-
-    func scpBaseArguments() -> [String] {
-        [
-            "-q",
-            "-o", "StrictHostKeyChecking=no",
-            "-o", "UserKnownHostsFile=/dev/null",
-            "-o", "PreferredAuthentications=password",
-            "-o", "NumberOfPasswordPrompts=1",
-            "-o", "ConnectTimeout=30",
-            "-P", "\(sshPort)",
-        ]
-    }
-
     func ssh(_ command: String, requireSuccess: Bool = true) async throws -> VPhoneCommandResult {
-        var lastError: Error?
-        for _ in 0 ..< sshRetry {
-            do {
-                let result = try await VPhoneHost.runCommand(
-                    "ssh",
-                    arguments: sshBaseArguments() + [command],
-                    environment: VPhoneHost.sshAskpassEnvironment(password: sshPassword, executablePath: patcherBinary),
-                    requireSuccess: requireSuccess
-                )
-                if !requireSuccess || result.terminationStatus.isSuccess || VPhoneHost.exitCode(from: result.terminationStatus) != 255 {
-                    return result
-                }
-            } catch {
-                lastError = error
-            }
-            try await Task.sleep(for: .seconds(3))
-        }
-        throw lastError ?? ValidationError("SSH command failed: \(command)")
+        let result = try requireSSHClient().execute(command, requireSuccess: requireSuccess)
+        return VPhoneCommandResult(
+            terminationStatus: .exited(result.exitStatus),
+            standardOutput: result.standardOutputString.trimmingCharacters(in: .newlines),
+            standardError: result.standardErrorString.trimmingCharacters(in: .newlines)
+        )
     }
 
     func scpTo(_ localPath: String, remotePath: String, recursive: Bool = false) async throws {
-        var arguments = scpBaseArguments()
+        let localURL = URL(fileURLWithPath: localPath).standardizedFileURL
         if recursive {
-            arguments.append("-r")
+            try requireSSHClient().uploadDirectory(localURL: localURL, remotePath: remotePath)
+        } else {
+            try requireSSHClient().uploadFile(localURL: localURL, remotePath: remotePath)
         }
-        arguments += [localPath, "\(sshUser)@\(sshHost):\(remotePath)"]
-        _ = try await VPhoneHost.runCommand(
-            "scp",
-            arguments: arguments,
-            environment: VPhoneHost.sshAskpassEnvironment(password: sshPassword, executablePath: patcherBinary),
-            requireSuccess: true
-        )
     }
 
     func scpFrom(_ remotePath: String, localPath: String) async throws {
-        let arguments = scpBaseArguments() + ["\(sshUser)@\(sshHost):\(remotePath)", localPath]
-        _ = try await VPhoneHost.runCommand(
-            "scp",
-            arguments: arguments,
-            environment: VPhoneHost.sshAskpassEnvironment(password: sshPassword, executablePath: patcherBinary),
-            requireSuccess: true
-        )
+        try requireSSHClient().downloadFile(remotePath: remotePath, localURL: URL(fileURLWithPath: localPath).standardizedFileURL)
     }
 
     func remoteFileExists(_ path: String) async throws -> Bool {
@@ -741,6 +700,13 @@ private struct VPhoneCFWInstaller {
         try data.write(to: launchdPlist)
     }
 
+    func requireSSHClient() throws -> VPhoneSSHClient {
+        guard let sshClient else {
+            throw ValidationError("SSH client is not connected")
+        }
+        return sshClient
+    }
+
     func unmount(paths: [String]) async throws {
         for path in paths {
             _ = try await ssh("/sbin/umount \(path) 2>/dev/null || true", requireSuccess: false)

sources/vphone-cli/VPhoneSSH.swift

@@ -0,0 +1,287 @@
+import Foundation
+import NIOCore
+import NIOPosix
+import NIOSSH
+
+struct VPhoneSSHCommandResult {
+    let exitStatus: Int32
+    let standardOutput: Data
+    let standardError: Data
+
+    var standardOutputString: String {
+        String(decoding: standardOutput, as: UTF8.self)
+    }
+
+    var standardErrorString: String {
+        String(decoding: standardError, as: UTF8.self)
+    }
+}
+
+enum VPhoneSSHError: Error, CustomStringConvertible {
+    case notConnected
+    case invalidChannelType
+    case commandFailed(String)
+
+    var description: String {
+        switch self {
+        case .notConnected:
+            return "SSH client is not connected"
+        case .invalidChannelType:
+            return "Invalid SSH channel type"
+        case let .commandFailed(message):
+            return message
+        }
+    }
+}
+
+final class VPhoneSSHClient: @unchecked Sendable {
+    let host: String
+    let port: Int
+    let username: String
+    let password: String
+
+    private let group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+    private var channel: Channel?
+    private var sshHandler: NIOSSHHandler?
+
+    init(host: String, port: Int, username: String, password: String) {
+        self.host = host
+        self.port = port
+        self.username = username
+        self.password = password
+    }
+
+    deinit {
+        try? shutdown()
+    }
+
+    func connect() throws {
+        let bootstrap = ClientBootstrap(group: group)
+            .channelInitializer { [username, password] channel in
+                channel.eventLoop.makeCompletedFuture {
+                    let ssh = NIOSSHHandler(
+                        role: .client(
+                            .init(
+                                userAuthDelegate: SimplePasswordDelegate(username: username, password: password),
+                                serverAuthDelegate: AcceptAllHostKeysDelegate()
+                            )
+                        ),
+                        allocator: channel.allocator,
+                        inboundChildChannelInitializer: nil
+                    )
+                    try channel.pipeline.syncOperations.addHandler(ssh)
+                    try channel.pipeline.syncOperations.addHandler(VPhoneSSHErrorHandler())
+                }
+            }
+            .channelOption(ChannelOptions.socket(SocketOptionLevel(SOL_SOCKET), SO_REUSEADDR), value: 1)
+            .channelOption(ChannelOptions.socket(SocketOptionLevel(IPPROTO_TCP), TCP_NODELAY), value: 1)
+
+        let connectedChannel = try bootstrap.connect(host: host, port: port).wait()
+        channel = connectedChannel
+        sshHandler = try connectedChannel.pipeline.syncOperations.handler(type: NIOSSHHandler.self)
+    }
+
+    func shutdown() throws {
+        if let channel {
+            try? channel.close().wait()
+            self.channel = nil
+        }
+        try group.syncShutdownGracefully()
+    }
+
+    func execute(_ command: String, stdin: Data? = nil, requireSuccess: Bool = true) throws -> VPhoneSSHCommandResult {
+        guard let channel, let sshHandler else {
+            throw VPhoneSSHError.notConnected
+        }
+
+        let resultPromise = channel.eventLoop.makePromise(of: VPhoneSSHCommandResult.self)
+        let childPromise = channel.eventLoop.makePromise(of: Channel.self)
+        sshHandler.createChannel(childPromise) { childChannel, channelType in
+            guard channelType == .session else {
+                return channel.eventLoop.makeFailedFuture(VPhoneSSHError.invalidChannelType)
+            }
+
+            return childChannel.eventLoop.makeCompletedFuture {
+                try childChannel.pipeline.syncOperations.addHandler(
+                    VPhoneSSHExecHandler(command: command, stdinData: stdin, resultPromise: resultPromise)
+                )
+            }
+        }
+
+        let childChannel = try childPromise.futureResult.wait()
+        try childChannel.closeFuture.wait()
+        let result = try resultPromise.futureResult.wait()
+        if requireSuccess, result.exitStatus != 0 {
+            let stderr = result.standardErrorString.trimmingCharacters(in: .whitespacesAndNewlines)
+            throw VPhoneSSHError.commandFailed(
+                stderr.isEmpty
+                    ? "SSH command failed with status \(result.exitStatus): \(command)"
+                    : "SSH command failed with status \(result.exitStatus): \(command)\n\(stderr)"
+            )
+        }
+        return result
+    }
+
+    func uploadFile(localURL: URL, remotePath: String) throws {
+        _ = try execute("/bin/cat > \(shellQuote(remotePath))", stdin: try Data(contentsOf: localURL))
+    }
+
+    func uploadData(_ data: Data, remotePath: String) throws {
+        _ = try execute("/bin/cat > \(shellQuote(remotePath))", stdin: data)
+    }
+
+    func downloadFile(remotePath: String, localURL: URL) throws {
+        let result = try execute("/bin/cat \(shellQuote(remotePath))")
+        try result.standardOutput.write(to: localURL)
+    }
+
+    func uploadDirectory(localURL: URL, remotePath: String) throws {
+        let archiveData = try VPhoneArchive.createTarArchive(from: localURL)
+        let temporaryRemotePath = "/tmp/vphone-upload-\(UUID().uuidString).tar"
+        try uploadData(archiveData, remotePath: temporaryRemotePath)
+        _ = try execute(
+            "/bin/mkdir -p \(shellQuote(remotePath)) && /usr/bin/tar --preserve-permissions --no-overwrite-dir -xf \(shellQuote(temporaryRemotePath)) -C \(shellQuote(remotePath)) && /bin/rm -f \(shellQuote(temporaryRemotePath))"
+        )
+    }
+
+    static func probe(host: String, port: Int, username: String, password: String) -> Bool {
+        do {
+            let client = VPhoneSSHClient(host: host, port: port, username: username, password: password)
+            defer { try? client.shutdown() }
+            try client.connect()
+            let result = try client.execute("echo ready", requireSuccess: false)
+            return result.exitStatus == 0
+        } catch {
+            return false
+        }
+    }
+
+    private func shellQuote(_ string: String) -> String {
+        "'" + string.replacingOccurrences(of: "'", with: "'\\''") + "'"
+    }
+}
+
+private final class VPhoneSSHExecHandler: ChannelDuplexHandler, @unchecked Sendable {
+    typealias InboundIn = SSHChannelData
+    typealias InboundOut = SSHChannelData
+    typealias OutboundIn = SSHChannelData
+    typealias OutboundOut = SSHChannelData
+
+    let command: String
+    let stdinData: Data?
+    let resultPromise: EventLoopPromise<VPhoneSSHCommandResult>
+
+    var standardOutput = Data()
+    var standardError = Data()
+    var exitStatus: Int32 = 0
+    var completed = false
+
+    init(command: String, stdinData: Data?, resultPromise: EventLoopPromise<VPhoneSSHCommandResult>) {
+        self.command = command
+        self.stdinData = stdinData
+        self.resultPromise = resultPromise
+    }
+
+    func handlerAdded(context: ChannelHandlerContext) {
+        let loopBoundContext = NIOLoopBound(context, eventLoop: context.eventLoop)
+        context.channel.setOption(ChannelOptions.allowRemoteHalfClosure, value: true).whenFailure { error in
+            self.fail(error, context: loopBoundContext.value)
+        }
+    }
+
+    func channelActive(context: ChannelHandlerContext) {
+        let loopBoundContext = NIOLoopBound(context, eventLoop: context.eventLoop)
+        let request = SSHChannelRequestEvent.ExecRequest(command: command, wantReply: true)
+        context.triggerUserOutboundEvent(request).whenComplete { result in
+            switch result {
+            case .success:
+                self.sendStandardInput(context: loopBoundContext.value)
+            case .failure(let error):
+                self.fail(error, context: loopBoundContext.value)
+            }
+        }
+    }
+
+    func channelRead(context: ChannelHandlerContext, data: NIOAny) {
+        let message = unwrapInboundIn(data)
+        guard case .byteBuffer(var bytes) = message.data,
+              let chunk = bytes.readData(length: bytes.readableBytes)
+        else {
+            return
+        }
+
+        switch message.type {
+        case .channel:
+            standardOutput.append(chunk)
+        case .stdErr:
+            standardError.append(chunk)
+        default:
+            break
+        }
+    }
+
+    func userInboundEventTriggered(context: ChannelHandlerContext, event: Any) {
+        if let exit = event as? SSHChannelRequestEvent.ExitStatus {
+            exitStatus = Int32(exit.exitStatus)
+        } else {
+            context.fireUserInboundEventTriggered(event)
+        }
+    }
+
+    func channelInactive(context: ChannelHandlerContext) {
+        succeedIfNeeded()
+        context.fireChannelInactive()
+    }
+
+    func errorCaught(context: ChannelHandlerContext, error: Error) {
+        fail(error, context: context)
+    }
+
+    private func sendStandardInput(context: ChannelHandlerContext) {
+        guard let stdinData, !stdinData.isEmpty else {
+            context.close(mode: .output, promise: nil)
+            return
+        }
+
+        var buffer = context.channel.allocator.buffer(capacity: stdinData.count)
+        buffer.writeBytes(stdinData)
+        let payload = SSHChannelData(type: .channel, data: .byteBuffer(buffer))
+        let loopBoundContext = NIOLoopBound(context, eventLoop: context.eventLoop)
+        context.writeAndFlush(wrapOutboundOut(payload)).whenComplete { _ in
+            loopBoundContext.value.close(mode: .output, promise: nil)
+        }
+    }
+
+    private func succeedIfNeeded() {
+        guard !completed else { return }
+        completed = true
+        resultPromise.succeed(
+            VPhoneSSHCommandResult(
+                exitStatus: exitStatus,
+                standardOutput: standardOutput,
+                standardError: standardError
+            )
+        )
+    }
+
+    private func fail(_ error: Error, context: ChannelHandlerContext) {
+        guard !completed else { return }
+        completed = true
+        resultPromise.fail(error)
+        context.close(promise: nil)
+    }
+}
+
+private final class AcceptAllHostKeysDelegate: NIOSSHClientServerAuthenticationDelegate {
+    func validateHostKey(hostKey: NIOSSHPublicKey, validationCompletePromise: EventLoopPromise<Void>) {
+        validationCompletePromise.succeed(())
+    }
+}
+
+private final class VPhoneSSHErrorHandler: ChannelInboundHandler {
+    typealias InboundIn = Any
+
+    func errorCaught(context: ChannelHandlerContext, error: Error) {
+        context.close(promise: nil)
+    }
+}

sources/vphone-cli/VPhoneSetupMachineCLI.swift

@@ -325,38 +325,18 @@ private struct SetupMachineRunner {
         return arguments
     }
 
-    func startUSBMuxForward(localPort: Int, serial: String) throws -> ManagedProcess {
-        let process = Process()
-        process.currentDirectoryURL = projectRoot
-        process.executableURL = URL(fileURLWithPath: patcherExecutable)
-        process.arguments = ["usbmux-forward", "--local-port", "\(localPort)", "--serial", serial, "--remote-port", "22"]
-        let output = Pipe()
-        process.standardOutput = output
-        process.standardError = output
-        try process.run()
-        return ManagedProcess(process: process, logHandle: nil)
+    func startUSBMuxForward(localPort: Int, serial: String) throws -> ManagedUSBMuxForward {
+        guard let local = UInt16(exactly: localPort) else {
+            throw ValidationError("Invalid local forward port: \(localPort)")
+        }
+        let service = try USBMuxForwarder.start(localPort: local, serial: serial, remotePort: 22)
+        return ManagedUSBMuxForward(service: service)
     }
 
     func waitForSSH(port: Int, timeout: TimeInterval = 90) async throws {
         let deadline = Date().addingTimeInterval(timeout)
         while Date() < deadline {
-            let result = try await VPhoneHost.runCommand(
-                "ssh",
-                arguments: [
-                    "-o", "StrictHostKeyChecking=no",
-                    "-o", "UserKnownHostsFile=/dev/null",
-                    "-o", "PreferredAuthentications=password",
-                    "-o", "NumberOfPasswordPrompts=1",
-                    "-o", "ConnectTimeout=5",
-                    "-q",
-                    "-p", "\(port)",
-                    "root@127.0.0.1",
-                    "echo ready",
-                ],
-                environment: VPhoneHost.sshAskpassEnvironment(password: "alpine"),
-                requireSuccess: false
-            )
-            if result.terminationStatus.isSuccess { return }
+            if VPhoneSSHClient.probe(host: "127.0.0.1", port: port, username: "root", password: "alpine") { return }
             try await Task.sleep(for: .seconds(2))
         }
         throw ValidationError("Timed out waiting for ramdisk SSH on port \(port)")
@@ -468,3 +448,15 @@ private final class ManagedProcess {
         try? logHandle?.close()
     }
 }
+
+private final class ManagedUSBMuxForward {
+    let service: USBMuxForwardingService
+
+    init(service: USBMuxForwardingService) {
+        self.service = service
+    }
+
+    func terminate() {
+        service.stop()
+    }
+}

sources/vphone-cli/VPhoneUSBMux.swift

@@ -252,33 +252,17 @@ enum USBMuxClient {
 }
 
 enum USBMuxForwarder {
-    static func run(localPort: UInt16, serial: String, remotePort: UInt16) throws {
+    static func start(localPort: UInt16, serial: String, remotePort: UInt16) throws -> USBMuxForwardingService {
         let device = try USBMuxClient.device(matching: serial)
         let listener = try makeListener(port: localPort)
-        defer { close(listener) }
+        let service = USBMuxForwardingService(listener: listener, device: device, remotePort: remotePort)
+        service.start()
+        return service
+    }
 
-        while true {
-            var address = sockaddr_storage()
-            var length = socklen_t(MemoryLayout<sockaddr_storage>.size)
-            let client = withUnsafeMutablePointer(to: &address) {
-                $0.withMemoryRebound(to: sockaddr.self, capacity: 1) {
-                    Darwin.accept(listener, $0, &length)
-                }
-            }
-            if client < 0 {
-                if errno == EINTR { continue }
-                throw USBMuxError.socketError("accept() failed: \(String(cString: strerror(errno)))")
-            }
-
-            DispatchQueue.global(qos: .userInitiated).async {
-                do {
-                    let remote = try USBMuxClient.connect(deviceID: device.deviceID, port: remotePort)
-                    relay(client: client, remote: remote)
-                } catch {
-                    close(client)
-                }
-            }
-        }
+    static func run(localPort: UInt16, serial: String, remotePort: UInt16) throws {
+        _ = try start(localPort: localPort, serial: serial, remotePort: remotePort)
+        dispatchMain()
     }
 
     static func makeListener(port: UInt16) throws -> Int32 {
@@ -349,3 +333,71 @@ enum USBMuxForwarder {
         }
     }
 }
+
+final class USBMuxForwardingService: @unchecked Sendable {
+    private let listener: Int32
+    private let device: USBMuxDevice
+    private let remotePort: UInt16
+    private let queue = DispatchQueue(label: "com.vphone.usbmux-forward", qos: .userInitiated)
+    private let lock = NSLock()
+    private var stopped = false
+
+    init(listener: Int32, device: USBMuxDevice, remotePort: UInt16) {
+        self.listener = listener
+        self.device = device
+        self.remotePort = remotePort
+    }
+
+    func start() {
+        queue.async { [self] in
+            acceptLoop()
+        }
+    }
+
+    func stop() {
+        lock.lock()
+        let shouldStop = !stopped
+        stopped = true
+        lock.unlock()
+
+        if shouldStop {
+            close(listener)
+        }
+    }
+
+    private func acceptLoop() {
+        while !isStopped {
+            var address = sockaddr_storage()
+            var length = socklen_t(MemoryLayout<sockaddr_storage>.size)
+            let client = withUnsafeMutablePointer(to: &address) {
+                $0.withMemoryRebound(to: sockaddr.self, capacity: 1) {
+                    Darwin.accept(listener, $0, &length)
+                }
+            }
+            if client < 0 {
+                if isStopped || errno == EBADF || errno == EINVAL {
+                    break
+                }
+                if errno == EINTR {
+                    continue
+                }
+                break
+            }
+
+            DispatchQueue.global(qos: .userInitiated).async { [device, remotePort] in
+                do {
+                    let remote = try USBMuxClient.connect(deviceID: device.deviceID, port: remotePort)
+                    USBMuxForwarder.relay(client: client, remote: remote)
+                } catch {
+                    close(client)
+                }
+            }
+        }
+    }
+
+    private var isStopped: Bool {
+        lock.lock()
+        defer { lock.unlock() }
+        return stopped
+    }
+}