Support stacked VM archive import and export

Author: yzhuang-oai

Sources/tart/Commands/Import.swift

@@ -21,6 +21,9 @@ struct Import: AsyncParsableCommand {
 
     // Create a temporary VM directory to which we will load the export file
     let tmpVMDir = try VMDirectory.temporary()
+    defer {
+      try? tmpVMDir.removeFromDisk()
+    }
 
     // Lock the temporary VM directory to prevent it's garbage collection
     // while we're running
@@ -30,10 +33,8 @@ struct Import: AsyncParsableCommand {
     // Populate the temporary VM directory with the export file contents
     print("importing...")
     try tmpVMDir.importFromArchive(path: path)
-
-    if tmpVMDir.isStackedVM || tmpVMDir.isStackedCachedImage {
-      try? FileManager.default.removeItem(at: tmpVMDir.baseURL)
-      throw RuntimeError.ImportFailed("importing stacked VMs is not supported yet")
+    guard tmpVMDir.initialized else {
+      throw RuntimeError.ImportFailed("archive does not contain a runnable VM")
     }
 
     try await withTaskCancellationHandler(operation: {
@@ -50,7 +51,7 @@ struct Import: AsyncParsableCommand {
 
       try lock.unlock()
     }, onCancel: {
-      try? FileManager.default.removeItem(at: tmpVMDir.baseURL)
+      try? tmpVMDir.removeFromDisk()
     })
   }
 }

Sources/tart/DiskImageStack.swift

@@ -45,6 +45,22 @@ struct DiskImageStack {
   let blockSize: UInt64
   let blockCount: UInt64
 
+  static var isSupported: Bool {
+    #if canImport(DiskImageKit)
+      if #available(macOS 27.0, *) {
+        return true
+      }
+    #endif
+
+    return false
+  }
+
+  static func requireSupport() throws {
+    guard isSupported else {
+      throw DiskImageStackError.unavailable
+    }
+  }
+
   /// Reads a disk image's current block layout without resolving or validating a
   /// whole stack. This is used for the VM's private writable overlay, whose
   /// size may be newer than the pinned immutable parent manifest.
@@ -69,21 +85,7 @@ struct DiskImageStack {
     #if canImport(DiskImageKit)
       if #available(macOS 27.0, *) {
         let image = try DiskImage(opening: .open(url: url, mode: .readOnly))
-        let matchesFormat = switch expectedFormat {
-        case .raw:
-          image.format == .raw
-        case .asif:
-          image.format == .asif
-        }
-        guard matchesFormat else {
-          throw DiskImageStackError.invalidDiskImage(url, "base disk format does not match")
-        }
-        guard image.layerType == nil, image.parentUUID == nil else {
-          throw DiskImageStackError.invalidDiskImage(url, "base disk must not be an overlay")
-        }
-        if expectedFormat == .asif && image.layerUUID == nil {
-          throw DiskImageStackError.invalidDiskImage(url, "ASIF base disk is missing a UUID")
-        }
+        try validateBase(image, at: url, expectedFormat: expectedFormat)
 
         return DiskImageBlockLayout(
           blockSize: UInt64(image.blockSize.rawValue),
@@ -214,7 +216,7 @@ struct DiskImageStack {
       }
 
       let baseImage = try DiskImage(opening: .open(url: baseURL, mode: .readOnly))
-      try validateBase(baseImage, at: baseURL, expectedFormat: baseFormat)
+      try Self.validateBase(baseImage, at: baseURL, expectedFormat: baseFormat)
 
       var image = baseImage
 
@@ -239,7 +241,7 @@ struct DiskImageStack {
     }
 
     @available(macOS 27.0, *)
-    private func validateBase(
+    private static func validateBase(
       _ image: DiskImage,
       at url: URL,
       expectedFormat: DiskImageFormat

Sources/tart/VMDirectory+Archive.swift

@@ -1,3 +1,4 @@
+import Foundation
 import System
 import AppleArchive
 
@@ -10,8 +11,14 @@ fileprivate let permissions = FilePermissions(rawValue: 0o644)
 // [2]: https://developer.apple.com/documentation/compression/algorithm/lzfse
 extension VMDirectory {
   func exportToArchive(path: String) throws {
-    guard !isStackedVM && !isStackedCachedImage else {
-      throw RuntimeError.ExportFailed("exporting stacked VMs is not supported yet")
+    let temporaryArchive = try stackedArchiveDirectoryIfNeeded()
+    let archiveSourceURL = temporaryArchive?.vmDirectory.baseURL ?? baseURL
+
+    defer {
+      if let temporaryArchive {
+        try? temporaryArchive.lock.unlock()
+        try? temporaryArchive.vmDirectory.removeFromDisk()
+      }
     }
 
     guard let fileStream = ArchiveByteStream.fileStream(
@@ -53,7 +60,7 @@ extension VMDirectory {
       return
     }
 
-    try encodeStream.writeDirectoryContents(archiveFrom: FilePath(baseURL.path), keySet: keySet)
+    try encodeStream.writeDirectoryContents(archiveFrom: FilePath(archiveSourceURL.path), keySet: keySet)
   }
 
   func importFromArchive(path: String) throws {
@@ -96,5 +103,145 @@ extension VMDirectory {
     }
 
     _ = try ArchiveStream.process(readingFrom: decodeStream, writingTo: extractStream)
+
+    if isStackedVM {
+      try restoreStackedArchive()
+    }
+  }
+
+  /// Builds a self-contained staging directory for a stacked archive, if this
+  /// directory currently resolves to a stacked VM or cached image.
+  private func stackedArchiveDirectoryIfNeeded() throws -> (vmDirectory: VMDirectory, lock: FileLock)? {
+    guard isStackedVM || isStackedCachedImage else {
+      return nil
+    }
+    try DiskImageStack.requireSupport()
+
+    let contentStore = try ContentStore()
+    let archiveVMDir = try VMDirectory.temporary()
+    let archiveVMDirLock = try FileLock(lockURL: archiveVMDir.baseURL)
+    try archiveVMDirLock.lock()
+
+    do {
+      let stagedSource: (isStackedVM: Bool, contentDigests: [String])? = try contentStore.withPruneLock {
+        () -> (isStackedVM: Bool, contentDigests: [String])? in
+        // OCI tags are mutable symlinks. Resolve one digest record while tag
+        // replacement and cached-image deletion are blocked, then copy every
+        // source-owned file before releasing the lock.
+        let sourceVMDir = VMDirectory(baseURL: baseURL.resolvingSymlinksInPath())
+        guard sourceVMDir.isStackedVM || sourceVMDir.isStackedCachedImage else {
+          return nil
+        }
+
+        let sourceIsStackedVM = sourceVMDir.isStackedVM
+        let sourceVMLock: PIDLock?
+        if sourceIsStackedVM {
+          let lock = try sourceVMDir.lock()
+          guard try lock.trylock() else {
+            throw RuntimeError.ExportFailed("VM \"\(sourceVMDir.name)\" must be stopped before export")
+          }
+          sourceVMLock = lock
+
+          // Holding the PID lock proves that the VM is not running. A saved
+          // state file is the remaining suspended state that must reject export.
+          guard !FileManager.default.fileExists(atPath: sourceVMDir.stateURL.path) else {
+            try? lock.unlock()
+            throw RuntimeError.ExportFailed("VM \"\(sourceVMDir.name)\" must be stopped before export")
+          }
+        } else {
+          sourceVMLock = nil
+        }
+        defer { try? sourceVMLock?.unlock() }
+
+        try FileManager.default.copyItem(at: sourceVMDir.configURL, to: archiveVMDir.configURL)
+        try FileManager.default.copyItem(at: sourceVMDir.nvramURL, to: archiveVMDir.nvramURL)
+        try FileManager.default.copyItem(at: sourceVMDir.manifestURL, to: archiveVMDir.manifestURL)
+        if sourceIsStackedVM {
+          try FileManager.default.copyItem(at: sourceVMDir.overlayURL, to: archiveVMDir.overlayURL)
+        }
+
+        return (sourceIsStackedVM, try archiveVMDir.diskContentDigests())
+      }
+
+      guard let stagedSource else {
+        try archiveVMDirLock.unlock()
+        try archiveVMDir.removeFromDisk()
+        return nil
+      }
+
+      if !stagedSource.isStackedVM {
+        try archiveVMDir.diskImageStack().createWritableOverlay()
+      }
+
+      // The staged manifest is now an in-progress reference, so immutable
+      // content remains protected while these potentially large copies run
+      // without holding the global prune lock.
+      for contentDigest in stagedSource.contentDigests {
+        guard let sourceURL = try contentStore.existingContentURL(for: contentDigest) else {
+          throw RuntimeError.ExportFailed("VM is missing cached disk content \(contentDigest)")
+        }
+
+        let destinationURL = try contentStore.contentURL(
+          for: contentDigest,
+          under: archiveContentStoreURL(in: archiveVMDir)
+        )
+        try FileManager.default.createDirectory(
+          at: destinationURL.deletingLastPathComponent(),
+          withIntermediateDirectories: true
+        )
+        try FileManager.default.copyItem(at: sourceURL, to: destinationURL)
+      }
+
+      return (archiveVMDir, archiveVMDirLock)
+    } catch {
+      try? archiveVMDirLock.unlock()
+      try? archiveVMDir.removeFromDisk()
+      throw error
+    }
+  }
+
+  /// Restores immutable files from an archive into the shared content store,
+  /// removes the archive-only payload, then validates the resulting stack.
+  private func restoreStackedArchive() throws {
+    try DiskImageStack.requireSupport()
+
+    let contentStore = try ContentStore()
+    // The extracted manifest is already a reference; synchronize publication
+    // with a concurrent prune before installing its immutable content.
+    try contentStore.synchronizePublishedReferences()
+    for contentDigest in try diskContentDigests() {
+      if try contentStore.existingContentURL(for: contentDigest) != nil {
+        continue
+      }
+
+      let archivedContentURL = try contentStore.contentURL(
+        for: contentDigest,
+        under: archiveContentStoreURL(in: self)
+      )
+      guard FileManager.default.fileExists(atPath: archivedContentURL.path) else {
+        throw RuntimeError.ImportFailed("archive is missing disk content \(contentDigest)")
+      }
+
+      let temporaryURL = try contentStore.temporaryContentURL(for: contentDigest)
+      do {
+        try FileManager.default.copyItem(at: archivedContentURL, to: temporaryURL)
+        _ = try contentStore.install(temporaryURL, contentDigest: contentDigest)
+      } catch {
+        try? FileManager.default.removeItem(at: temporaryURL)
+        throw error
+      }
+    }
+
+    try FileManager.default.removeItem(at: archiveContentStoreURL(in: self))
+    try? FileManager.default.removeItem(at: stateURL)
+
+    // Opening the attachment validates the reconstructed immutable stack and
+    // imported writable overlay before the VM enters local storage.
+    _ = try diskImageStack().makeAttachment()
   }
+
+  private func archiveContentStoreURL(in vmDir: VMDirectory) -> URL {
+    vmDir.baseURL.appendingPathComponent("content", isDirectory: true)
+  }
+
 }