macOS 27: DiskImageKit Example

DiskImageKitArticleDemoView.swift

import SwiftUI
import Observation
import DiskImageKit
import Virtualization

@available(macOS 27.0, *)
@MainActor
@Observable
final class VMDiskSnapshotModel {

    struct LayerRow: Identifiable, Hashable {
        let id = UUID()
        let url: URL
        let format: String
        let layerType: String
        let size: Int
        let blockCount: Int
        let blockSize: Int
    }

    private(set) var image: DiskImage?
    private(set) var stack: (any StackedImage)?
    private(set) var layers: [LayerRow] = []
    private(set) var errorMessage: String?
    private(set) var status = "Idle"

    let workingDirectory: URL

    init() {
        let base = FileManager.default.urls(
            for: .applicationSupportDirectory,
            in: .userDomainMask
        ).first!

        workingDirectory = base
            .appendingPathComponent("DiskImageKitDemo", isDirectory: true)
    }

    var baseURL: URL {
        workingDirectory
            .appendingPathComponent("base.asif")
    }

    var cacheURL: URL {
        workingDirectory
            .appendingPathComponent("cache.asif")
    }

    var overlayURL: URL {
        workingDirectory
            .appendingPathComponent("overlay.asif")
    }

    func createBaseImage() {
        do {
            try prepareDirectory()
            removeIfExists(baseURL)

            let gigabytes = 8
            let byteCount = gigabytes * 1_000_000_000
            let blockSize = DiskImage.BlockSize.bytes512
            let blockCount = byteCount / 512

            let image = try DiskImage(
                creating: .asif(
                    url: baseURL,
                    blockCount: blockCount,
                    blockSize: blockSize
                )
            )

            self.image = image
            self.stack = nil
            self.status = "Created \(gigabytes) GB ASIF base image"
            self.errorMessage = nil

            refreshLayers(from: image)
        } catch {
            fail(error)
        }
    }

    func openBaseImage() {
        do {
            let image = try DiskImage(
                opening: .open(url: baseURL)
            )

            self.image = image
            self.stack = nil
            self.status = "Opened base image"
            self.errorMessage = nil

            refreshLayers(from: image)
        } catch {
            fail(error)
        }
    }

    func addCacheLayer() {
        guard let image else {
            status = "Create or open a base image first"
            return
        }

        do {
            try prepareDirectory()
            removeIfExists(cacheURL)

            let stacked = try image.appending(
                .asifLayer(
                    url: cacheURL,
                    type: .cache
                )
            )

            self.stack = stacked
            self.image = stacked
            self.status = "Added cache layer"
            self.errorMessage = nil

            refreshLayers(from: stacked)
        } catch {
            fail(error)
        }
    }

    func addOverlayLayer() {
        guard let currentImage = image else {
            status = "Create or open a base image first"
            return
        }

        do {
            try prepareDirectory()
            removeIfExists(overlayURL)

            let stacked = try currentImage.appending(
                .asifLayer(
                    url: overlayURL,
                    type: .overlay
                )
            )

            self.stack = stacked
            self.image = stacked
            self.status = "Added overlay layer"
            self.errorMessage = nil

            refreshLayers(from: stacked)
        } catch {
            fail(error)
        }
    }

    func extendLogicalSize(toGigabytes gigabytes: Int) {
        guard let image else {
            status = "Create or open an image first"
            return
        }

        do {
            let bytes = gigabytes * 1_000_000_000
            let divisor = image.blockSize == .bytes4096 ? 4096 : 512
            let blockCount = bytes / divisor

            try image.truncate(
                blockCount: blockCount
            )

            status = "Resized logical disk to \(gigabytes) GB"
            errorMessage = nil

            refreshLayers(from: image)
        } catch {
            fail(error)
        }
    }

    func makeVirtualizationAttachment() throws -> VZDiskImageStorageDeviceAttachment {
        guard let image else {
            throw DemoError.noImage
        }

        return try VZDiskImageStorageDeviceAttachment(
            diskImage: image
        )
    }

    func reset() {
        image = nil
        stack = nil
        layers = []
        errorMessage = nil
        status = "Idle"

        removeIfExists(baseURL)
        removeIfExists(cacheURL)
        removeIfExists(overlayURL)
    }

    private func refreshLayers(from image: DiskImage) {
        if let stacked = image as? any StackedImage {
            layers = stacked.layers.map(Self.makeRow)
        } else {
            layers = [Self.makeRow(image)]
        }
    }

    private static func makeRow(_ image: DiskImage) -> LayerRow {
        LayerRow(
            url: image.url,
            format: String(describing: image.format),
            layerType: image.layerType.map(String.init(describing:)) ?? "base",
            size: image.size,
            blockCount: image.blockCount,
            blockSize: image.blockSize.rawValue
        )
    }

    private func prepareDirectory() throws {
        try FileManager.default.createDirectory(
            at: workingDirectory,
            withIntermediateDirectories: true
        )
    }

    private func removeIfExists(_ url: URL) {
        try? FileManager.default.removeItem(
            at: url
        )
    }

    private func fail(_ error: Error) {
        errorMessage = error.localizedDescription
        status = "Operation failed"
    }

    enum DemoError: LocalizedError {
        case noImage

        var errorDescription: String? {
            "No disk image is currently open."
        }
    }
}

@available(macOS 27.0, *)
struct VMDiskSnapshotView: View {

    @State private var model = VMDiskSnapshotModel()

    var body: some View {
        NavigationSplitView {
            List(model.layers) { layer in
                VStack(alignment: .leading, spacing: 4) {
                    Text(layer.url.lastPathComponent)
                        .font(.headline)

                    Text(layer.layerType)
                        .font(.caption)
                        .foregroundStyle(.secondary)

                    Text(layer.format)
                        .font(.caption.monospaced())
                        .foregroundStyle(.secondary)
                }
                .padding(.vertical, 4)
            }
            .navigationTitle("Disk Layers")
        } detail: {
            VStack(alignment: .leading, spacing: 18) {
                Text("VM Disk Snapshot Manager")
                    .font(.largeTitle.bold())

                Text(model.status)
                    .foregroundStyle(.secondary)

                if let errorMessage = model.errorMessage {
                    Text(errorMessage)
                        .foregroundStyle(.red)
                }

                HStack {
                    Button("Create Base") {
                        model.createBaseImage()
                    }

                    Button("Open Base") {
                        model.openBaseImage()
                    }

                    Button("Add Cache") {
                        model.addCacheLayer()
                    }

                    Button("Add Overlay") {
                        model.addOverlayLayer()
                    }

                    Button("Extend to 12 GB") {
                        model.extendLogicalSize(
                            toGigabytes: 12
                        )
                    }

                    Button("Reset", role: .destructive) {
                        model.reset()
                    }
                }

                if let image = model.image {
                    GroupBox("Current Image") {
                        Grid(alignment: .leading) {
                            GridRow {
                                Text("Format")
                                Text(String(describing: image.format))
                            }

                            GridRow {
                                Text("Logical size")
                                Text(ByteCountFormatter.string(
                                    fromByteCount: Int64(image.size),
                                    countStyle: .file
                                ))
                            }

                            GridRow {
                                Text("Blocks")
                                Text(image.blockCount.formatted())
                            }

                            GridRow {
                                Text("Block size")
                                Text("\(image.blockSize.rawValue) bytes")
                            }

                            GridRow {
                                Text("Open mode")
                                Text(String(describing: image.openMode))
                            }

                            GridRow {
                                Text("Top URL")
                                Text(image.url.path)
                                    .textSelection(.enabled)
                            }
                        }
                        .padding(8)
                    }
                }

                Spacer()
            }
            .padding()
        }
    }
}

struct DiskImageKitArticleDemoView: View {

    var body: some View {
        if #available(macOS 27.0, *) {
            VMDiskSnapshotView()
        } else {
            ContentUnavailableView(
                "macOS 27 Required",
                systemImage: "externaldrive.badge.exclamationmark",
                description: Text(
                    "This sample uses the DiskImageKit framework."
                )
            )
        }
    }
}

#Preview {
    DiskImageKitArticleDemoView()
}