From 51a7969b09c0c6e83365bce8d9b1ca23be82ca56 Mon Sep 17 00:00:00 2001 From: osy <50960678+osy@users.noreply.github.com> Date: Sat, 24 Feb 2024 18:00:41 -0800 Subject: [PATCH] remote: support takeover of existing session and auto-pause of orphaned sessions --- Platform/Shared/VMDetailsView.swift | 4 +-- Platform/UTMData.swift | 6 ++-- Platform/VMData.swift | 6 +++- Platform/iOS/UTMDataExtension.swift | 4 +-- Platform/iOS/VMSessionState.swift | 6 +++- Platform/macOS/UTMDataExtension.swift | 7 +++++ Platform/macOS/VMRemoteSessionState.swift | 4 +-- Remote/UTMRemoteClient.swift | 2 +- Remote/UTMRemoteMessage.swift | 2 ++ Remote/UTMRemoteServer.swift | 34 +++++++++++++++++++++-- Remote/UTMRemoteSpiceVirtualMachine.swift | 10 +++++-- 11 files changed, 68 insertions(+), 17 deletions(-) diff --git a/Platform/Shared/VMDetailsView.swift b/Platform/Shared/VMDetailsView.swift index 6e21a0d2..b1cbe8d5 100644 --- a/Platform/Shared/VMDetailsView.swift +++ b/Platform/Shared/VMDetailsView.swift @@ -162,7 +162,7 @@ struct Screenshot: View { .blendMode(.hardLight) #if os(visionOS) .overlay { - if vm.isStopped { + if vm.isStopped || vm.isTakeoverAllowed { Image(systemName: "play.circle.fill") .resizable() .frame(width: 100, height: 100) @@ -175,7 +175,7 @@ struct Screenshot: View { #endif if vm.isBusy { Spinner(size: .large) - } else if vm.isStopped { + } else if vm.isStopped || vm.isTakeoverAllowed { #if !os(visionOS) Button(action: { data.run(vm: vm) }, label: { Label("Run", systemImage: "play.circle.fill") diff --git a/Platform/UTMData.swift b/Platform/UTMData.swift index a45790c6..b034851f 100644 --- a/Platform/UTMData.swift +++ b/Platform/UTMData.swift @@ -829,8 +829,9 @@ struct AlertMessage: Identifiable { }) observers.insert(vm.$state.sink { state in Task { + let isTakeoverAllowed = self.vmWindows[vm] is VMRemoteSessionState && (state == .started || state == .paused) await self.remoteServer.broadcast { remote in - try await remote.virtualMachine(id: vm.id, didTransitionToState: state) + try await remote.virtualMachine(id: vm.id, didTransitionToState: state, isTakeoverAllowed: isTakeoverAllowed) } } }) @@ -1271,12 +1272,13 @@ class UTMRemoteData: UTMData { vm.updateMountedDrives(mountedDrives) } - func remoteVirtualMachineDidTransition(id: UUID, state: UTMVirtualMachineState) async { + func remoteVirtualMachineDidTransition(id: UUID, state: UTMVirtualMachineState, isTakeoverAllowed: Bool) async { guard let vm = virtualMachines.first(where: { $0.id == id }) else { return } let remoteVM = vm as! VMRemoteData let wrapped = remoteVM.wrapped as! UTMRemoteSpiceVirtualMachine + remoteVM.isTakeoverAllowed = isTakeoverAllowed await wrapped.updateRemoteState(state) } diff --git a/Platform/VMData.swift b/Platform/VMData.swift index c8351a1c..8afe9a05 100644 --- a/Platform/VMData.swift +++ b/Platform/VMData.swift @@ -68,7 +68,10 @@ import SwiftUI /// Copy from wrapped VM @Published var screenshot: UTMVirtualMachineScreenshot? - + + /// If true, it is possible to hijack the session. + @Published var isTakeoverAllowed: Bool = false + /// Allows changes in the config, registry, and VM to be reflected private var observers: [AnyCancellable] = [] @@ -450,6 +453,7 @@ class VMRemoteData: VMData { self._isShortcut = item.isShortcut self.initialState = item.state super.init() + self.isTakeoverAllowed = item.isTakeoverAllowed self.registryEntryWrapped = UTMRegistry.shared.entry(uuid: item.id, name: item.name, path: item.path) self.registryEntryWrapped!.isSuspended = item.isSuspended self.registryEntryWrapped!.externalDrives = item.mountedDrives.mapValues({ UTMRegistryEntry.File(dummyFromPath: $0) }) diff --git a/Platform/iOS/UTMDataExtension.swift b/Platform/iOS/UTMDataExtension.swift index 6d6c606a..3c109a48 100644 --- a/Platform/iOS/UTMDataExtension.swift +++ b/Platform/iOS/UTMDataExtension.swift @@ -30,9 +30,9 @@ extension UTMData { } if let session = VMSessionState.allActiveSessions.values.first(where: { $0.vm.id == wrapped.id }) { session.showWindow() - } else if vm.state == .stopped { + } else if vm.isStopped || vm.isTakeoverAllowed { let session = VMSessionState(for: wrapped as! (any UTMSpiceVirtualMachine)) - session.start() + session.start(options: options) } else { showErrorAlert(message: NSLocalizedString("This virtual machine is already running. In order to run it from this device, you must stop it first.", comment: "UTMDataExtension")) } diff --git a/Platform/iOS/VMSessionState.swift b/Platform/iOS/VMSessionState.swift index 1182e471..0805b635 100644 --- a/Platform/iOS/VMSessionState.swift +++ b/Platform/iOS/VMSessionState.swift @@ -434,7 +434,11 @@ extension VMSessionState { } Self.allActiveSessions[id] = self showWindow() - vm.requestVmStart(options: options) + if vm.state == .paused { + vm.requestVmResume() + } else { + vm.requestVmStart(options: options) + } } func showWindow() { diff --git a/Platform/macOS/UTMDataExtension.swift b/Platform/macOS/UTMDataExtension.swift index 631d198c..3acdb861 100644 --- a/Platform/macOS/UTMDataExtension.swift +++ b/Platform/macOS/UTMDataExtension.swift @@ -86,6 +86,13 @@ extension UTMData { guard let wrapped = vm.wrapped as? UTMQemuVirtualMachine, type(of: wrapped).capabilities.supportsRemoteSession else { throw UTMDataError.unsupportedBackend } + if let existingSession = vmWindows[vm] as? VMRemoteSessionState, let spiceServerInfo = wrapped.spiceServerInfo { + if wrapped.state == .paused { + try await wrapped.resume() + } + existingSession.client = client + return spiceServerInfo + } guard vmWindows[vm] == nil else { throw UTMDataError.virtualMachineUnavailable } diff --git a/Platform/macOS/VMRemoteSessionState.swift b/Platform/macOS/VMRemoteSessionState.swift index 30b2ed56..66a946c7 100644 --- a/Platform/macOS/VMRemoteSessionState.swift +++ b/Platform/macOS/VMRemoteSessionState.swift @@ -19,7 +19,7 @@ import IOKit.pwr_mgt /// Represents the UI state for a single headless VM session. class VMRemoteSessionState: VMHeadlessSessionState { - let client: UTMRemoteServer.Remote + public weak var client: UTMRemoteServer.Remote? init(for vm: any UTMVirtualMachine, client: UTMRemoteServer.Remote, onStop: (() -> Void)?) { self.client = client @@ -28,7 +28,7 @@ class VMRemoteSessionState: VMHeadlessSessionState { override func virtualMachine(_ vm: any UTMVirtualMachine, didErrorWithMessage message: String) { Task { - try? await client.virtualMachine(id: vm.id, didErrorWithMessage: message) + try? await client?.virtualMachine(id: vm.id, didErrorWithMessage: message) super.virtualMachine(vm, didErrorWithMessage: message) } } diff --git a/Remote/UTMRemoteClient.swift b/Remote/UTMRemoteClient.swift index 5ca01307..f2cc3d67 100644 --- a/Remote/UTMRemoteClient.swift +++ b/Remote/UTMRemoteClient.swift @@ -313,7 +313,7 @@ extension UTMRemoteClient { } private func _virtualMachineDidTransition(parameters: M.VirtualMachineDidTransition.Request) async throws -> M.VirtualMachineDidTransition.Reply { - await data.remoteVirtualMachineDidTransition(id: parameters.id, state: parameters.state) + await data.remoteVirtualMachineDidTransition(id: parameters.id, state: parameters.state, isTakeoverAllowed: parameters.isTakeoverAllowed) return .init() } diff --git a/Remote/UTMRemoteMessage.swift b/Remote/UTMRemoteMessage.swift index 92bfd800..9901d0be 100644 --- a/Remote/UTMRemoteMessage.swift +++ b/Remote/UTMRemoteMessage.swift @@ -74,6 +74,7 @@ extension UTMRemoteMessageServer { let path: String let isShortcut: Bool let isSuspended: Bool + let isTakeoverAllowed: Bool let backend: UTMBackend let state: UTMVirtualMachineState let mountedDrives: [String: String] @@ -360,6 +361,7 @@ extension UTMRemoteMessageClient { struct Request: Serializable, Codable { let id: UUID let state: UTMVirtualMachineState + let isTakeoverAllowed: Bool } struct Reply: Serializable, Codable {} diff --git a/Remote/UTMRemoteServer.swift b/Remote/UTMRemoteServer.swift index 76f88b8d..ee39e3be 100644 --- a/Remote/UTMRemoteServer.swift +++ b/Remote/UTMRemoteServer.swift @@ -228,11 +228,36 @@ actor UTMRemoteServer { if !connectedClients.contains(client) { if let remote = establishedConnections.removeValue(forKey: client) { remote.close() + Task { @MainActor in + await suspendSessions(for: remote) + } } } } } + @MainActor + private func suspendSessions(for remote: Remote) async { + let sessions = data.vmWindows.compactMap { + if let session = $0.value as? VMRemoteSessionState { + return ($0.key, session) + } else { + return nil + } + } + await withTaskGroup(of: Void.self) { group in + for (vm, session) in sessions { + if session.client?.id == remote.id { + session.client = nil + } + group.addTask { + try? await vm.wrapped?.pause() + } + } + await group.waitForAll() + } + } + private func establishConnection(_ connection: Connection) async { guard let fingerprint = connection.fingerprint else { connection.close() @@ -712,11 +737,13 @@ extension UTMRemoteServer { try parameters.ids.map { id in let vm = try findVM(withId: id) let mountedDrives = vm.registryEntry?.externalDrives.mapValues({ $0.path }) ?? [:] + let isTakeoverAllowed = data.vmWindows[vm] is VMRemoteSessionState && (vm.state == .started || vm.state == .paused) return M.VirtualMachineInformation(id: vm.id, name: vm.detailsTitleLabel, path: vm.pathUrl.path, isShortcut: vm.isShortcut, isSuspended: vm.registryEntry?.isSuspended ?? false, + isTakeoverAllowed: isTakeoverAllowed, backend: vm.wrapped is UTMQemuVirtualMachine ? .qemu : .unknown, state: vm.wrapped?.state ?? .stopped, mountedDrives: mountedDrives) @@ -850,9 +877,10 @@ extension UTMRemoteServer { } extension UTMRemoteServer { - class Remote { + class Remote: Identifiable { typealias M = UTMRemoteMessageClient fileprivate(set) var peer: Peer! + let id = UUID() func close() { peer.close() @@ -876,8 +904,8 @@ extension UTMRemoteServer { try await _mountedDrivesHasChanged(parameters: .init(id: id, mountedDrives: mountedDrives)) } - func virtualMachine(id: UUID, didTransitionToState state: UTMVirtualMachineState) async throws { - try await _virtualMachineDidTransition(parameters: .init(id: id, state: state)) + func virtualMachine(id: UUID, didTransitionToState state: UTMVirtualMachineState, isTakeoverAllowed: Bool) async throws { + try await _virtualMachineDidTransition(parameters: .init(id: id, state: state, isTakeoverAllowed: isTakeoverAllowed)) } func virtualMachine(id: UUID, didErrorWithMessage message: String) async throws { diff --git a/Remote/UTMRemoteSpiceVirtualMachine.swift b/Remote/UTMRemoteSpiceVirtualMachine.swift index ea18dd78..5acde0a4 100644 --- a/Remote/UTMRemoteSpiceVirtualMachine.swift +++ b/Remote/UTMRemoteSpiceVirtualMachine.swift @@ -197,7 +197,7 @@ extension UTMRemoteSpiceVirtualMachine { } func start(options: UTMVirtualMachineStartOptions) async throws { - try await _state.operation(before: .stopped, during: .starting, after: .started) { + try await _state.operation(before: [.stopped, .started, .paused], during: .starting, after: .started) { let spiceServer = try await server.startVirtualMachine(id: id, options: options) var options = UTMSpiceIOOptions() if await !config.sound.isEmpty { @@ -250,8 +250,12 @@ extension UTMRemoteSpiceVirtualMachine { } func resume() async throws { - try await _state.operation(before: .paused, during: .resuming, after: .started) { - try await server.resumeVirtualMachine(id: id) + if ioService == nil { + return try await start(options: []) + } else { + try await _state.operation(before: .paused, during: .resuming, after: .started) { + try await server.resumeVirtualMachine(id: id) + } } }