From ab37e574a08a37a3a111d803e1d3951b4f0432b2 Mon Sep 17 00:00:00 2001
From: Danny Canter <danny_canter@apple.com>
Date: Tue, 29 Jul 2025 12:41:54 -0700
Subject: [PATCH] ContainerManager: Make root/imagestore suppliable (#228)

We should be able to supply an ImageStore if the user wants a different
ContentStore, or a different path.
---
 .../Containerization/ContainerManager.swift   | 72 ++++++++++++++++++-
 1 file changed, 69 insertions(+), 3 deletions(-)

diff --git a/Sources/Containerization/ContainerManager.swift b/Sources/Containerization/ContainerManager.swift
index b012d1fb..be222472 100644
--- a/Sources/Containerization/ContainerManager.swift
+++ b/Sources/Containerization/ContainerManager.swift
@@ -192,13 +192,15 @@ public struct ContainerManager: Sendable {
         }
     }
 
-    /// Create a new manager with the provided kernel and initfs mount.
+    /// Create a new manager with the provided kernel, initfs mount, image store
+    /// and optional network implementation.
     public init(
         kernel: Kernel,
         initfs: Mount,
+        imageStore: ImageStore,
         network: Network? = nil
     ) throws {
-        self.imageStore = ImageStore.default
+        self.imageStore = imageStore
         self.network = network
         try Self.createRootDirectory(path: self.imageStore.path)
         self.vmm = VZVirtualMachineManager(
@@ -208,13 +210,77 @@ public struct ContainerManager: Sendable {
         )
     }
 
+    /// Create a new manager with the provided kernel, initfs mount, root state
+    /// directory and optional network implementation.
+    public init(
+        kernel: Kernel,
+        initfs: Mount,
+        root: URL? = nil,
+        network: Network? = nil
+    ) throws {
+        if let root {
+            self.imageStore = try ImageStore(path: root)
+        } else {
+            self.imageStore = ImageStore.default
+        }
+        self.network = network
+        try Self.createRootDirectory(path: self.imageStore.path)
+        self.vmm = VZVirtualMachineManager(
+            kernel: kernel,
+            initialFilesystem: initfs,
+            bootlog: self.imageStore.path.appendingPathComponent("bootlog.log").absolutePath()
+        )
+    }
+
+    /// Create a new manager with the provided kernel, initfs reference, image store
+    /// and optional network implementation.
+    public init(
+        kernel: Kernel,
+        initfsReference: String,
+        imageStore: ImageStore,
+        network: Network? = nil
+    ) async throws {
+        self.imageStore = imageStore
+        self.network = network
+        try Self.createRootDirectory(path: self.imageStore.path)
+
+        let initPath = self.imageStore.path.appendingPathComponent("initfs.ext4")
+        let initImage = try await self.imageStore.getInitImage(reference: initfsReference)
+        let initfs = try await {
+            do {
+                return try await initImage.initBlock(at: initPath, for: .linuxArm)
+            } catch let err as ContainerizationError {
+                guard err.code == .exists else {
+                    throw err
+                }
+                return .block(
+                    format: "ext4",
+                    source: initPath.absolutePath(),
+                    destination: "/",
+                    options: ["ro"]
+                )
+            }
+        }()
+
+        self.vmm = VZVirtualMachineManager(
+            kernel: kernel,
+            initialFilesystem: initfs,
+            bootlog: self.imageStore.path.appendingPathComponent("bootlog.log").absolutePath()
+        )
+    }
+
     /// Create a new manager with the provided kernel and image reference for the initfs.
     public init(
         kernel: Kernel,
         initfsReference: String,
+        root: URL? = nil,
         network: Network? = nil
     ) async throws {
-        self.imageStore = ImageStore.default
+        if let root {
+            self.imageStore = try ImageStore(path: root)
+        } else {
+            self.imageStore = ImageStore.default
+        }
         self.network = network
         try Self.createRootDirectory(path: self.imageStore.path)