From 48cd4b47e45460288fa3b9aa5de1d8e0e13cbce2 Mon Sep 17 00:00:00 2001
From: Nikolay Edigaryev <edigaryev@gmail.com>
Date: Thu, 21 Jul 2022 12:37:43 +0300
Subject: [PATCH] Move full-fledged VNC support to --experimental-vnc (#154)

---
 Sources/tart/Commands/Run.swift               | 48 ++++++++++++-------
 .../FullFledgedVNC.swift}                     | 28 +++++------
 Sources/tart/VNC/ScreenSharingVNC.swift       | 25 ++++++++++
 Sources/tart/VNC/VNC.swift                    |  6 +++
 4 files changed, 77 insertions(+), 30 deletions(-)
 rename Sources/tart/{VNCWrapper.swift => VNC/FullFledgedVNC.swift} (63%)
 create mode 100644 Sources/tart/VNC/ScreenSharingVNC.swift
 create mode 100644 Sources/tart/VNC/VNC.swift

diff --git a/Sources/tart/Commands/Run.swift b/Sources/tart/Commands/Run.swift
index 645164f..4e26c62 100644
--- a/Sources/tart/Commands/Run.swift
+++ b/Sources/tart/Commands/Run.swift
@@ -24,41 +24,57 @@ struct Run: AsyncParsableCommand {
   
   @Flag(help: ArgumentHelp(
           "Use screen sharing instead of the built-in UI.",
-          discussion: "Useful since VNC supports copy/paste, drag and drop, etc.\nNote that Remote Login option should be enabled inside the VM.")) 
+          discussion: "Useful since Screen Sharing supports copy/paste, drag and drop, etc.\n"
+            + "Note that Remote Login option should be enabled inside the VM."))
   var vnc: Bool = false
 
+  @Flag(help: ArgumentHelp(
+    "Use Virtualization.Framework's VNC server instead of the build-in UI.",
+    discussion: "Useful since this type of VNC is available in recovery mode and in macOS installation.\n"
+      + "Note that this feature is experimental and there may be bugs present when using VNC."))
+  var vncExperimental: Bool = false
+
   @Flag var withSoftnet: Bool = false
 
+  func validate() throws {
+    if vnc && vncExperimental {
+      throw ValidationError("--vnc and --vnc-experimental are mutually exclusive")
+    }
+  }
+
   @MainActor
   func run() async throws {    
     let vmDir = try VMStorageLocal().open(name)
     vm = try VM(vmDir: vmDir, withSoftnet: withSoftnet)
 
-    var vncWrapper: VNCWrapper?
-
-    if vnc {
-        vncWrapper = VNCWrapper(virtualMachine: vm!.virtualMachine)
-    }
+    let vncImpl: VNC? = try {
+      if vnc {
+        let vmConfig = try VMConfig.init(fromURL: vmDir.configURL)
+        return ScreenSharingVNC(vmConfig: vmConfig)
+      } else if vncExperimental {
+        return FullFledgedVNC(virtualMachine: vm!.virtualMachine)
+      } else {
+        return nil
+      }
+    }()
 
     let task = Task {
       do {
-        if let vncWrapper = vncWrapper {
-          let port = try await vncWrapper.waitForPort()
-
-          let url = URL(string: "vnc://:\(vncWrapper.password)@127.0.0.1:\(port)")!
+        if let vncImpl = vncImpl {
+          let vncURL = try await vncImpl.waitForURL()
 
           if noGraphics || ProcessInfo.processInfo.environment["CI"] != nil {
-            print("VNC server is running at \(url)")
+            print("VNC server is running at \(vncURL)")
           } else {
-            print("Opening \(url)...")
-            NSWorkspace.shared.open(url)
+            print("Opening \(vncURL)...")
+            NSWorkspace.shared.open(vncURL)
           }
         }
 
         try await vm!.run(recovery)
 
-        if let vncWrapper = vncWrapper {
-          try vncWrapper.stop()
+        if let vncImpl = vncImpl {
+          try vncImpl.stop()
         }
 
         Foundation.exit(0)
@@ -79,7 +95,7 @@ struct Run: AsyncParsableCommand {
     }
     sigintSrc.activate()
 
-    if noGraphics || vnc {
+    if noGraphics || vnc || vncExperimental {
       dispatchMain()
     } else {
       runUI()
diff --git a/Sources/tart/VNCWrapper.swift b/Sources/tart/VNC/FullFledgedVNC.swift
similarity index 63%
rename from Sources/tart/VNCWrapper.swift
rename to Sources/tart/VNC/FullFledgedVNC.swift
index d92ef41..62e94b8 100644
--- a/Sources/tart/VNCWrapper.swift
+++ b/Sources/tart/VNC/FullFledgedVNC.swift
@@ -2,7 +2,7 @@ import Foundation
 import Dynamic
 import Virtualization
 
-class VNCWrapper {
+class FullFledgedVNC: VNC {
     let password: String
     private let vnc: Dynamic
 
@@ -15,6 +15,19 @@ class VNCWrapper {
         vnc.start()
     }
 
+    func waitForURL() async throws -> URL {
+      while true {
+        // Port is 0 shortly after start(),
+        // but will be initialized later
+        if let port = vnc.port.asUInt16, port != 0 {
+          return URL(string: "vnc://:\(password)@127.0.0.1:\(port)")!
+        }
+
+        // Wait 50 ms.
+        try await Task.sleep(nanoseconds: 50_000_000)
+      }
+    }
+
     func stop() throws {
         vnc.stop()
     }
@@ -22,17 +35,4 @@ class VNCWrapper {
     deinit {
         try? stop()
     }
-
-    func waitForPort() async throws -> UInt16 {
-        while true {
-            // Port is 0 shortly after start(),
-            // but will be initialized later
-            if let port = vnc.port.asUInt16, port != 0 {
-                return port
-            }
-
-            // Wait 50 ms.
-            try await Task.sleep(nanoseconds: 50_000_000)
-        }
-    }
 }
diff --git a/Sources/tart/VNC/ScreenSharingVNC.swift b/Sources/tart/VNC/ScreenSharingVNC.swift
new file mode 100644
index 0000000..982295f
--- /dev/null
+++ b/Sources/tart/VNC/ScreenSharingVNC.swift
@@ -0,0 +1,25 @@
+import Foundation
+import Dynamic
+import Virtualization
+
+class ScreenSharingVNC: VNC {
+  let vmConfig: VMConfig
+
+  init(vmConfig: VMConfig) {
+    self.vmConfig = vmConfig
+  }
+
+  func waitForURL() async throws -> URL {
+    let ip = try await IP.resolveIP(vmConfig, secondsToWait: 60)
+
+    if let ip = ip {
+      return URL(string: "vnc://\(ip)")!
+    }
+
+    throw IPNotFound()
+  }
+
+  func stop() throws {
+    // nothing to do
+  }
+}
diff --git a/Sources/tart/VNC/VNC.swift b/Sources/tart/VNC/VNC.swift
new file mode 100644
index 0000000..76afe53
--- /dev/null
+++ b/Sources/tart/VNC/VNC.swift
@@ -0,0 +1,6 @@
+import Foundation
+
+protocol VNC {
+  func waitForURL() async throws -> URL
+  func stop() throws
+}