From 85dfa961c9d29f772553f6314eae38da8af74f43 Mon Sep 17 00:00:00 2001
From: Nikolay Edigaryev <edigaryev@gmail.com>
Date: Tue, 26 Jul 2022 16:25:23 +0300
Subject: [PATCH] tart run: introduce --disk option (#156)

* tart run: introduce --disk option

* Document how to create disks using Disk Utility
---
 Sources/tart/Commands/Run.swift | 37 +++++++++++++++++++++++++++++++--
 Sources/tart/VM.swift           | 26 ++++++++++++++++-------
 2 files changed, 53 insertions(+), 10 deletions(-)

diff --git a/Sources/tart/Commands/Run.swift b/Sources/tart/Commands/Run.swift
index b71c39a..fc879ad 100644
--- a/Sources/tart/Commands/Run.swift
+++ b/Sources/tart/Commands/Run.swift
@@ -36,6 +36,14 @@ struct Run: AsyncParsableCommand {
 
   @Flag var withSoftnet: Bool = false
 
+  @Option(help: ArgumentHelp("""
+    Additional disk attachments with an optional read-only specifier\n(e.g. --disk=\"disk.bin\" --disk=\"disk.bin:ro\")
+    """, discussion: """
+    Learn how to create a disk image using Disk Utility here:
+    https://support.apple.com/en-gb/guide/disk-utility/dskutl11888/mac
+    """))
+  var disk: [String] = []
+
   func validate() throws {
     if vnc && vncExperimental {
       throw ValidationError("--vnc and --vnc-experimental are mutually exclusive")
@@ -43,9 +51,13 @@ struct Run: AsyncParsableCommand {
   }
 
   @MainActor
-  func run() async throws {    
+  func run() async throws {
     let vmDir = try VMStorageLocal().open(name)
-    vm = try VM(vmDir: vmDir, withSoftnet: withSoftnet)
+    vm = try VM(
+      vmDir: vmDir,
+      withSoftnet: withSoftnet,
+      additionalDiskAttachments: additionalDiskAttachments()
+    )
 
     let vncImpl: VNC? = try {
       if vnc {
@@ -102,6 +114,27 @@ struct Run: AsyncParsableCommand {
     }
   }
 
+  func additionalDiskAttachments() throws -> [VZDiskImageStorageDeviceAttachment] {
+    var result: [VZDiskImageStorageDeviceAttachment] = []
+    let readOnlySuffix = ":ro"
+
+    for rawDisk in disk {
+      if rawDisk.hasSuffix(readOnlySuffix) {
+        result.append(try VZDiskImageStorageDeviceAttachment(
+          url: URL(fileURLWithPath: String(rawDisk.prefix(rawDisk.count - readOnlySuffix.count))),
+          readOnly: true
+        ))
+      } else {
+        result.append(try VZDiskImageStorageDeviceAttachment(
+          url: URL(fileURLWithPath: rawDisk),
+          readOnly: false
+        ))
+      }
+    }
+
+    return result
+  }
+
   private func runUI() {
     let nsApp = NSApplication.shared
     nsApp.setActivationPolicy(.regular)
diff --git a/Sources/tart/VM.swift b/Sources/tart/VM.swift
index 76c3bc0..66dde17 100644
--- a/Sources/tart/VM.swift
+++ b/Sources/tart/VM.swift
@@ -25,7 +25,10 @@ class VM: NSObject, VZVirtualMachineDelegate, ObservableObject {
 
   var softnet: Softnet? = nil
 
-  init(vmDir: VMDirectory, withSoftnet: Bool = false) throws {
+  init(vmDir: VMDirectory,
+       withSoftnet: Bool = false,
+       additionalDiskAttachments: [VZDiskImageStorageDeviceAttachment] = []
+  ) throws {
     let auxStorage = VZMacAuxiliaryStorage(contentsOf: vmDir.nvramURL)
 
     name = vmDir.name
@@ -37,7 +40,7 @@ class VM: NSObject, VZVirtualMachineDelegate, ObservableObject {
     }
 
     let configuration = try Self.craftConfiguration(diskURL: vmDir.diskURL, auxStorage: auxStorage, vmConfig: config,
-      softnet: softnet)
+      softnet: softnet, additionalDiskAttachments: additionalDiskAttachments)
     virtualMachine = VZVirtualMachine(configuration: configuration)
 
     super.init()
@@ -94,7 +97,13 @@ class VM: NSObject, VZVirtualMachineDelegate, ObservableObject {
     }
   }
 
-  init(vmDir: VMDirectory, ipswURL: URL?, diskSizeGB: UInt16, withSoftnet: Bool = false) async throws {
+  init(
+    vmDir: VMDirectory,
+    ipswURL: URL?,
+    diskSizeGB: UInt16,
+    withSoftnet: Bool = false,
+    additionalDiskAttachments: [VZDiskImageStorageDeviceAttachment] = []
+  ) async throws {
     let ipswURL = ipswURL != nil ? ipswURL! : try await VM.retrieveLatestIPSW();
 
     // Load the restore image and try to get the requirements
@@ -132,7 +141,7 @@ class VM: NSObject, VZVirtualMachineDelegate, ObservableObject {
     }
 
     let configuration = try Self.craftConfiguration(diskURL: vmDir.diskURL, auxStorage: auxStorage, vmConfig: config,
-      softnet: softnet)
+      softnet: softnet, additionalDiskAttachments: additionalDiskAttachments)
     virtualMachine = VZVirtualMachine(configuration: configuration)
 
     super.init()
@@ -183,7 +192,8 @@ class VM: NSObject, VZVirtualMachineDelegate, ObservableObject {
     diskURL: URL,
     auxStorage: VZMacAuxiliaryStorage,
     vmConfig: VMConfig,
-    softnet: Softnet? = nil
+    softnet: Softnet? = nil,
+    additionalDiskAttachments: [VZDiskImageStorageDeviceAttachment]
   ) throws -> VZVirtualMachineConfiguration {
     let configuration = VZVirtualMachineConfiguration()
 
@@ -247,9 +257,9 @@ class VM: NSObject, VZVirtualMachineDelegate, ObservableObject {
     configuration.networkDevices = [vio]
 
     // Storage
-    let attachment = try VZDiskImageStorageDeviceAttachment(url: diskURL, readOnly: false)
-    let storage = VZVirtioBlockDeviceConfiguration(attachment: attachment)
-    configuration.storageDevices = [storage]
+    var attachments = [try VZDiskImageStorageDeviceAttachment(url: diskURL, readOnly: false)]
+    attachments.append(contentsOf: additionalDiskAttachments)
+    configuration.storageDevices = attachments.map { VZVirtioBlockDeviceConfiguration(attachment: $0) }
 
     // Entropy
     configuration.entropyDevices = [VZVirtioEntropyDeviceConfiguration()]