From 2dfabeb4f9f5a85c90b09532c228fcdd46e0fac9 Mon Sep 17 00:00:00 2001
From: Danny Canter <danny_canter@apple.com>
Date: Wed, 25 Jun 2025 11:07:31 -0700
Subject: [PATCH] Vminitd: Rework ManagedContainer (#168)

This gets rid of some design decisions that I was not fond of. This
changes the process start for both execs and the init process to happen
internally in the object, as well as makes it so that all exec
functionalities can happen on the container object itself. This allows
us to get rid of the `if request.id == containerID` branches in the
rpcs, and handle this entirely in the object itself using its internal
state.

The existing integration tests should stress this just fine I believe.
---
 .../Sources/vminitd/ManagedContainer.swift    | 48 +++++++++++--------
 vminitd/Sources/vminitd/Server+GRPC.swift     | 40 ++++------------
 2 files changed, 36 insertions(+), 52 deletions(-)

diff --git a/vminitd/Sources/vminitd/ManagedContainer.swift b/vminitd/Sources/vminitd/ManagedContainer.swift
index 9907ea7f..484ea510 100644
--- a/vminitd/Sources/vminitd/ManagedContainer.swift
+++ b/vminitd/Sources/vminitd/ManagedContainer.swift
@@ -91,34 +91,29 @@ extension ManagedContainer {
         self._execs[id] = process
     }
 
-    func getExec(id: String) throws -> ManagedProcess {
-        guard let exec = self._execs[id] else {
-            throw ContainerizationError(
-                .invalidState,
-                message: "exec \(id) does not exist in container \(self.id)"
-            )
-        }
-        return exec
+    func start(execID: String) async throws -> Int32 {
+        let proc = try self.getExecOrInit(execID: execID)
+        return try await ProcessSupervisor.default.start(process: proc)
     }
 
-    func start() throws -> Int32 {
-        try self.initProcess.start()
+    func wait(execID: String) async throws -> Int32 {
+        let proc = try self.getExecOrInit(execID: execID)
+        return await proc.wait()
     }
 
-    func wait() async -> Int32 {
-        await self.initProcess.wait()
+    func kill(execID: String, _ signal: Int32) throws {
+        let proc = try self.getExecOrInit(execID: execID)
+        try proc.kill(signal)
     }
 
-    func kill(_ signal: Int32) throws {
-        try self.initProcess.kill(signal)
+    func resize(execID: String, size: Terminal.Size) throws {
+        let proc = try self.getExecOrInit(execID: execID)
+        try proc.resize(size: size)
     }
 
-    func resize(size: Terminal.Size) throws {
-        try self.initProcess.resize(size: size)
-    }
-
-    func close() throws {
-        try self.initProcess.close()
+    func close(execID: String) throws {
+        let proc = try self.getExecOrInit(execID: execID)
+        try proc.close()
     }
 
     func deleteExec(id: String) throws {
@@ -134,6 +129,19 @@ extension ManagedContainer {
     func delete() throws {
         try self._bundle.delete()
     }
+
+    func getExecOrInit(execID: String) throws -> ManagedProcess {
+        if execID == self.id {
+            return self.initProcess
+        }
+        guard let proc = self._execs[execID] else {
+            throw ContainerizationError(
+                .invalidState,
+                message: "exec \(execID) does not exist in container \(self.id)"
+            )
+        }
+        return proc
+    }
 }
 
 extension ContainerizationOCI.Bundle {
diff --git a/vminitd/Sources/vminitd/Server+GRPC.swift b/vminitd/Sources/vminitd/Server+GRPC.swift
index 038db9ee..ba99ec80 100644
--- a/vminitd/Sources/vminitd/Server+GRPC.swift
+++ b/vminitd/Sources/vminitd/Server+GRPC.swift
@@ -465,13 +465,7 @@ extension Initd: Com_Apple_Containerization_Sandbox_V3_SandboxContextAsyncProvid
         }
 
         let ctr = try await self.state.get(container: request.containerID)
-
-        if request.id == request.containerID {
-            try await ctr.kill(request.signal)
-        } else {
-            let proc = try await ctr.getExec(id: request.id)
-            try proc.kill(request.signal)
-        }
+        try await ctr.kill(execID: request.id, request.signal)
 
         return .init()
     }
@@ -520,16 +514,7 @@ extension Initd: Com_Apple_Containerization_Sandbox_V3_SandboxContextAsyncProvid
 
         do {
             let ctr = try await self.state.get(container: request.containerID)
-
-            // FIXME: This should just happen inside of ManagedContainer.
-            let pid: Int32
-            if request.id == request.containerID {
-                let process = ctr.initProcess
-                pid = try await ProcessSupervisor.default.start(process: process)
-            } else {
-                let process = try await ctr.getExec(id: request.id)
-                pid = try await ProcessSupervisor.default.start(process: process)
-            }
+            let pid = try await ctr.start(execID: request.id)
 
             return .with {
                 $0.pid = pid
@@ -565,14 +550,11 @@ extension Initd: Com_Apple_Containerization_Sandbox_V3_SandboxContextAsyncProvid
 
         do {
             let ctr = try await self.state.get(container: request.containerID)
-
-            let size = Terminal.Size(width: UInt16(request.columns), height: UInt16(request.rows))
-            if request.id == request.containerID {
-                try await ctr.resize(size: size)
-            } else {
-                let proc = try await ctr.getExec(id: request.id)
-                try proc.resize(size: size)
-            }
+            let size = Terminal.Size(
+                width: UInt16(request.columns),
+                height: UInt16(request.rows)
+            )
+            try await ctr.resize(execID: request.id, size: size)
         } catch {
             log.error(
                 "resizeProcess",
@@ -607,13 +589,7 @@ extension Initd: Com_Apple_Containerization_Sandbox_V3_SandboxContextAsyncProvid
         do {
             let ctr = try await self.state.get(container: request.containerID)
 
-            let exitCode: Int32
-            if request.id == request.containerID {
-                exitCode = await ctr.wait()
-            } else {
-                let proc = try await ctr.getExec(id: request.id)
-                exitCode = await proc.wait()
-            }
+            let exitCode = try await ctr.wait(execID: request.id)
 
             return .with {
                 $0.exitCode = exitCode