From 5eccdf74120276e1005407a1fbbf804daa70da20 Mon Sep 17 00:00:00 2001
From: Nikolay Edigaryev <edigaryev@gmail.com>
Date: Tue, 2 Jul 2024 18:12:35 +0400
Subject: [PATCH] Support customizing VM disks and mounting remote VMs in `tart
 run` (#847)

* Support remote VM names in --disk command-line argument

* tart set: introduce "--disk" to support replacing VM's disk contents

* Complete the code comment
---
 Sources/tart/Commands/Run.swift | 31 +++++++++++++++++++++++++++----
 Sources/tart/Commands/Set.swift | 11 +++++++++++
 2 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/Sources/tart/Commands/Run.swift b/Sources/tart/Commands/Run.swift
index f37c76dc..2a85da4e 100644
--- a/Sources/tart/Commands/Run.swift
+++ b/Sources/tart/Commands/Run.swift
@@ -67,12 +67,16 @@ struct Run: AsyncParsableCommand {
   var vncExperimental: Bool = false
 
   @Option(help: ArgumentHelp("""
-  Additional disk attachments with an optional read-only specifier\n(e.g. --disk=\"disk.bin\" --disk=\"ubuntu.iso:ro\" --disk=\"/dev/disk0\" --disk=\"nbd://localhost:10809/myDisk\")
+  Additional disk attachments with an optional read-only specifier\n(e.g. --disk=\"disk.bin\" --disk=\"ubuntu.iso:ro\" --disk=\"/dev/disk0\" --disk "ghcr.io/cirruslabs/xcode:16.0:ro" --disk=\"nbd://localhost:10809/myDisk\")
   """, discussion: """
-  Can be either a disk image file, a block device like a local SSD on AWS EC2 Mac instances or a Network Block Device (NBD).
+  The disk attachment can be a:
 
-  Learn how to create a disk image using Disk Utility here:
-  https://support.apple.com/en-gb/guide/disk-utility/dskutl11888/mac
+  * path to a disk image file
+  * path to a block device (for example, a local SSD on AWS EC2 Mac instances)
+  * remote VM name whose disk will be mounted
+  * Network Block Device (NBD) URL
+
+  Learn how to create a disk image using Disk Utility here: https://support.apple.com/en-gb/guide/disk-utility/dskutl11888/mac
 
   To work with block devices, the easiest way is to modify their permissions (e.g. by using "sudo chown $USER /dev/diskX") or to run the Tart binary as root, which affects locating Tart VMs.
 
@@ -496,6 +500,25 @@ struct Run: AsyncParsableCommand {
         continue
       }
 
+      // Support remote VM names in --disk command-line argument
+      if let remoteName = try? RemoteName(diskPath) {
+        let vmDir = try VMStorageOCI().open(remoteName)
+
+        // Unfortunately, VZDiskImageStorageDeviceAttachment does not support
+        // FileHandle, so we can't easily clone the disk, open it and unlink(2)
+        // to simplify the garbage collection, so use an intermediate directory.
+        let clonedDiskURL = try Config().tartTmpDir.appendingPathComponent("run-disk-\(UUID().uuidString)")
+
+        try FileManager.default.copyItem(at: vmDir.diskURL, to: clonedDiskURL)
+
+        let lock = try FileLock(lockURL: clonedDiskURL)
+        try lock.lock()
+
+        let diskImageAttachment = try VZDiskImageStorageDeviceAttachment(url: clonedDiskURL, readOnly: diskReadOnly)
+        result.append(VZVirtioBlockDeviceConfiguration(attachment: diskImageAttachment))
+        continue
+      }
+
       // Error out if the disk is locked by the host (e.g. it was mounted in Finder),
       // see https://github.com/cirruslabs/tart/issues/323 for more details.
       if try !diskReadOnly && !FileLock(lockURL: diskFileURL).trylock() {
diff --git a/Sources/tart/Commands/Set.swift b/Sources/tart/Commands/Set.swift
index 33bd31af..d4b1c9c9 100644
--- a/Sources/tart/Commands/Set.swift
+++ b/Sources/tart/Commands/Set.swift
@@ -25,6 +25,9 @@ struct Set: AsyncParsableCommand {
   #endif
   var randomSerial: Bool = false
 
+  @Option(help: ArgumentHelp("Replace the VM's disk contents with the disk contents at path.", valueName: "path"))
+  var disk: String?
+
   @Option(help: ArgumentHelp("Resize the VMs disk to the specified size in GB (note that the disk size can only be increased to avoid losing data)",
                              discussion: """
                              Disk resizing works on most cloud-ready Linux distributions out-of-the box (e.g. Ubuntu Cloud Images
@@ -73,6 +76,14 @@ struct Set: AsyncParsableCommand {
 
     try vmConfig.save(toURL: vmDir.configURL)
 
+    if let disk = disk {
+      let temporaryDiskURL = try Config().tartTmpDir.appendingPathComponent("set-disk-\(UUID().uuidString)")
+
+      try FileManager.default.copyItem(atPath: disk, toPath: temporaryDiskURL.path())
+
+      _ = try FileManager.default.replaceItemAt(vmDir.diskURL, withItemAt: temporaryDiskURL)
+    }
+
     if diskSize != nil {
       try vmDir.resizeDisk(diskSize!)
     }