General bootstrap env that overrides with dynamicEnv for SSH forwarding

chavezMac · chavezMac · commit c396115203b1 · 2026-03-13T20:56:04.000-07:00
diff --git a/Package.swift b/Package.swift
@@ -354,8 +354,8 @@ let package = Package(
                 .product(name: "Containerization", package: "containerization"),
                 .product(name: "ContainerizationExtras", package: "containerization"),
                 .product(name: "ContainerizationOCI", package: "containerization"),
-                "ContainerAPIService",
                 "ContainerResource",
+                "ContainerAPIService",
             ]
         ),
         .target(
@@ -394,6 +394,7 @@ let package = Package(
             dependencies: [
                 .product(name: "Containerization", package: "containerization"),
                 "ContainerResource",
+                "ContainerSandboxService",
                 "ContainerSandboxServiceClient",
             ]
         ),
diff --git a/Sources/ContainerCommands/Container/ContainerRun.swift b/Sources/ContainerCommands/Container/ContainerRun.swift
@@ -128,8 +128,11 @@ extension Application {
                     try? io.close()
                 }
 
-                let sshAuthSocketPath = ck.0.ssh ? ProcessInfo.processInfo.environment["SSH_AUTH_SOCK"] : nil
-                let process = try await client.bootstrap(id: id, stdio: io.stdio, sshAuthSocketPath: sshAuthSocketPath)
+                var dynamicEnv: [String: String] = [:]
+                if ck.0.ssh, let sshAuthSock = ProcessInfo.processInfo.environment["SSH_AUTH_SOCK"] {
+                    dynamicEnv["SSH_AUTH_SOCK"] = sshAuthSock
+                }
+                let process = try await client.bootstrap(id: id, stdio: io.stdio, dynamicEnv: dynamicEnv)
                 progress.finish()
 
                 if !self.managementFlags.cidfile.isEmpty {
diff --git a/Sources/ContainerCommands/Container/ContainerStart.swift b/Sources/ContainerCommands/Container/ContainerStart.swift
@@ -87,8 +87,11 @@ extension Application {
                     try? io.close()
                 }
 
-                let sshAuthSocketPath = container.configuration.ssh ? ProcessInfo.processInfo.environment["SSH_AUTH_SOCK"] : nil
-                let process = try await client.bootstrap(id: container.id, stdio: io.stdio, sshAuthSocketPath: sshAuthSocketPath)
+                var dynamicEnv: [String: String] = [:]
+                if container.configuration.ssh, let sshAuthSock = ProcessInfo.processInfo.environment["SSH_AUTH_SOCK"] {
+                    dynamicEnv["SSH_AUTH_SOCK"] = sshAuthSock
+                }
+                let process = try await client.bootstrap(id: container.id, stdio: io.stdio, dynamicEnv: dynamicEnv)
                 progress.finish()
 
                 if detach {
diff --git a/Sources/Services/ContainerAPIService/Client/ContainerClient.swift b/Sources/Services/ContainerAPIService/Client/ContainerClient.swift
@@ -113,8 +113,8 @@ public struct ContainerClient: Sendable {
     }
 
     /// Bootstrap the container's init process.
-    /// - Parameter sshAuthSocketPath: Optional path to the current shell's SSH agent socket, supplied at bootstrap time when SSH forwarding is enabled.
-    public func bootstrap(id: String, stdio: [FileHandle?], sshAuthSocketPath: String? = nil) async throws -> ClientProcess {
+    /// - Parameter dynamicEnv: Optional start-time environment overrides passed through bootstrap.
+    public func bootstrap(id: String, stdio: [FileHandle?], dynamicEnv: [String: String] = [:]) async throws -> ClientProcess {
         let request = XPCMessage(route: .containerBootstrap)
 
         for (i, h) in stdio.enumerated() {
@@ -135,8 +135,9 @@ public struct ContainerClient: Sendable {
 
         do {
             request.set(key: .id, value: id)
-            if let sshAuthSocketPath {
-                request.set(key: .sshAuthSocketPath, value: sshAuthSocketPath)
+            if !dynamicEnv.isEmpty {
+                let encodedDynamicEnv = try JSONEncoder().encode(dynamicEnv)
+                request.set(key: .dynamicEnv, value: encodedDynamicEnv)
             }
             try await xpcClient.send(request)
             return ClientProcessImpl(containerId: id, xpcClient: xpcClient)
diff --git a/Sources/Services/ContainerAPIService/Client/XPC+.swift b/Sources/Services/ContainerAPIService/Client/XPC+.swift
@@ -113,7 +113,7 @@ public enum XPCKeys: String {
     case initImage
 
     /// SSH agent socket path supplied at bootstrap time (current client shell).
-    case sshAuthSocketPath
+    case dynamicEnv
 
     /// Volume
     case volume
diff --git a/Sources/Services/ContainerAPIService/Server/Containers/ContainersHarness.swift b/Sources/Services/ContainerAPIService/Server/Containers/ContainersHarness.swift
@@ -55,8 +55,13 @@ public struct ContainersHarness: Sendable {
             )
         }
         let stdio = message.stdio()
-        let sshAuthSocketPath = message.string(key: .sshAuthSocketPath)
-        try await service.bootstrap(id: id, stdio: stdio, sshAuthSocketPath: sshAuthSocketPath)
+        let dynamicEnv: [String: String] =
+            if let data = message.dataNoCopy(key: .dynamicEnv) {
+                try JSONDecoder().decode([String: String].self, from: data)
+            } else {
+                [:]
+            }
+        try await service.bootstrap(id: id, stdio: stdio, dynamicEnv: dynamicEnv)
         return message.reply()
     }
 
diff --git a/Sources/Services/ContainerAPIService/Server/Containers/ContainersService.swift b/Sources/Services/ContainerAPIService/Server/Containers/ContainersService.swift
@@ -398,7 +398,8 @@ public actor ContainersService {
     }
 
     /// Bootstrap the init process of the container.
-    public func bootstrap(id: String, stdio: [FileHandle?]) async throws {
+    /// - Parameter dynamicEnv: Optional start-time environment overrides passed from the client.
+    public func bootstrap(id: String, stdio: [FileHandle?], dynamicEnv: [String: String] = [:]) async throws {
         log.debug(
             "ContainersService: enter",
             metadata: [
@@ -473,7 +474,7 @@ public actor ContainersService {
                     id: id,
                     runtime: runtime
                 )
-                try await sandboxClient.bootstrap(stdio: stdio, allocatedAttachments: allocatedAttachments)
+                try await sandboxClient.bootstrap(stdio: stdio, allocatedAttachments: allocatedAttachments, dynamicEnv: dynamicEnv)
 
                 try await self.exitMonitor.registerProcess(
                     id: id,
diff --git a/Sources/Services/ContainerSandboxService/Client/SandboxClient.swift b/Sources/Services/ContainerSandboxService/Client/SandboxClient.swift
@@ -77,7 +77,8 @@ public struct SandboxClient: Sendable {
 
 // Runtime Methods
 extension SandboxClient {
-    public func bootstrap(stdio: [FileHandle?], allocatedAttachments: [AllocatedAttachment]) async throws {
+    /// - Parameter dynamicEnv: Optional start-time environment overrides passed from the API service.
+    public func bootstrap(stdio: [FileHandle?], allocatedAttachments: [AllocatedAttachment], dynamicEnv: [String: String] = [:]) async throws {
         let request = XPCMessage(route: SandboxRoutes.bootstrap.rawValue)
 
         for (i, h) in stdio.enumerated() {
@@ -96,8 +97,9 @@ extension SandboxClient {
             }
         }
 
-        if let sshAuthSocketPath {
-            request.set(key: SandboxKeys.sshAuthSocketPath.rawValue, value: sshAuthSocketPath)
+        if !dynamicEnv.isEmpty {
+            let encodedDynamicEnv = try JSONEncoder().encode(dynamicEnv)
+            request.set(key: SandboxKeys.dynamicEnv.rawValue, value: encodedDynamicEnv)
         }
 
         do {
diff --git a/Sources/Services/ContainerSandboxService/Client/SandboxKeys.swift b/Sources/Services/ContainerSandboxService/Client/SandboxKeys.swift
@@ -43,6 +43,9 @@ public enum SandboxKeys: String {
     /// Container statistics
     case statistics
 
+    /// SSH agent socket path supplied at bootstrap time (current client shell).
+    case dynamicEnv
+
     /// Network resource keys.
     case allocatedAttachments
     case networkAdditionalData
diff --git a/Sources/Services/ContainerSandboxService/Server/SandboxService.swift b/Sources/Services/ContainerSandboxService/Server/SandboxService.swift
@@ -108,6 +108,38 @@ public actor SandboxService {
         return URL(fileURLWithPath: resolved.path)
     }
 
+    /// Merges start-time environment overrides into a base environment list.
+    /// Existing keys are replaced in-place and new keys are appended.
+    /// Made static for unit testability.
+    public static func mergedEnvironmentVariables(base: [String], overrides: [String: String]) -> [String] {
+        guard !overrides.isEmpty else {
+            return base
+        }
+
+        var env = base
+        var pending = overrides
+
+        for (index, entry) in env.enumerated() {
+            guard let separator = entry.firstIndex(of: "=") else {
+                continue
+            }
+
+            let key = String(entry[..<separator])
+            if let overrideValue = pending.removeValue(forKey: key) {
+                env[index] = "\(key)=\(overrideValue)"
+            }
+        }
+
+        for key in pending.keys.sorted() {
+            guard let value = pending[key] else {
+                continue
+            }
+            env.append("\(key)=\(value)")
+        }
+
+        return env
+    }
+
     /// Create an instance with a bundle that describes the container.
     ///
     /// - Parameters:
@@ -177,10 +209,15 @@ public actor SandboxService {
             try bundle.createLogFile()
 
             var config = try bundle.configuration
-            // Extract sshAuthSocketPath from the XPC request; if present and config.ssh is true,
-            // we mount the socket and set SSH_AUTH_SOCK in the process environment (replacing the
-            // previous behavior that used the sandbox's launch env).
-            let bootstrapSshAuthPath = message.string(key: SandboxKeys.sshAuthSocketPath.rawValue)
+            // Extract dynamic env overrides from the XPC request; when SSH forwarding is enabled,
+            // SSH_AUTH_SOCK from this map provides the host socket path to mount.
+            let dynamicEnv: [String: String] =
+                if let data = message.dataNoCopy(key: SandboxKeys.dynamicEnv.rawValue) {
+                    try JSONDecoder().decode([String: String].self, from: data)
+                } else {
+                    [:]
+                }
+            let bootstrapSshAuthPath = dynamicEnv[Self.sshAuthSocketEnvVar]
 
             if config.ssh {
                 if let resolved = Self.resolveSSHAuthSocketHostPath(config: config, bootstrapOverridePath: bootstrapSshAuthPath) {
@@ -273,7 +310,12 @@ public actor SandboxService {
             let id = config.id
             let rootfs = try bundle.containerRootfs.asMount
             let container = try LinuxContainer(id, rootfs: rootfs, vmm: vmm, logger: self.log) { czConfig in
-                try Self.configureContainer(czConfig: &czConfig, config: config, bootstrapOverridePath: bootstrapSshAuthPath)
+                try Self.configureContainer(
+                    czConfig: &czConfig,
+                    config: config,
+                    bootstrapOverridePath: bootstrapSshAuthPath,
+                    dynamicEnv: dynamicEnv
+                )
                 czConfig.interfaces = interfaces
                 czConfig.process.stdout = stdout
                 czConfig.process.stderr = stderr
@@ -899,7 +941,8 @@ public actor SandboxService {
     private static func configureContainer(
         czConfig: inout LinuxContainer.Configuration,
         config: ContainerConfiguration,
-        bootstrapOverridePath: String? = nil
+        bootstrapOverridePath: String? = nil,
+        dynamicEnv: [String: String] = [:]
     ) throws {
         czConfig.cpus = config.resources.cpus
         czConfig.memoryInBytes = config.resources.memoryInBytes
@@ -958,7 +1001,12 @@ public actor SandboxService {
                 searchDomains: dns.searchDomains, options: dns.options)
         }
 
-        try Self.configureInitialProcess(czConfig: &czConfig, config: config, bootstrapOverridePath: bootstrapOverridePath)
+        try Self.configureInitialProcess(
+            czConfig: &czConfig,
+            config: config,
+            bootstrapOverridePath: bootstrapOverridePath,
+            dynamicEnv: dynamicEnv
+        )
     }
 
     private func getDefaultNameservers(allocatedAttachments: [AllocatedAttachment]) async throws -> [String] {
@@ -976,17 +1024,20 @@ public actor SandboxService {
     private static func configureInitialProcess(
         czConfig: inout LinuxContainer.Configuration,
         config: ContainerConfiguration,
-        bootstrapOverridePath: String? = nil
+        bootstrapOverridePath: String? = nil,
+        dynamicEnv: [String: String] = [:]
     ) throws {
         let process = config.initProcess
 
         czConfig.process.arguments = [process.executable] + process.arguments
-        czConfig.process.environmentVariables = process.environment
+        czConfig.process.environmentVariables = Self.mergedEnvironmentVariables(
+            base: process.environment,
+            overrides: dynamicEnv
+        )
 
         if Self.sshAuthSocketHostUrl(config: config, bootstrapOverridePath: bootstrapOverridePath) != nil {
-            if !czConfig.process.environmentVariables.contains(where: { $0.starts(with: "\(Self.sshAuthSocketEnvVar)=") }) {
-                czConfig.process.environmentVariables.append("\(Self.sshAuthSocketEnvVar)=\(Self.sshAuthSocketGuestPath)")
-            }
+            czConfig.process.environmentVariables.removeAll(where: { $0.starts(with: "\(Self.sshAuthSocketEnvVar)=") })
+            czConfig.process.environmentVariables.append("\(Self.sshAuthSocketEnvVar)=\(Self.sshAuthSocketGuestPath)")
         }
 
         czConfig.process.terminal = process.terminal
diff --git a/Tests/CLITests/Subcommands/Run/TestCLIRunLifecycle.swift b/Tests/CLITests/Subcommands/Run/TestCLIRunLifecycle.swift
@@ -15,6 +15,7 @@
 //===----------------------------------------------------------------------===//
 
 import ContainerizationError
+import Foundation
 import Testing
 
 class TestCLIRunLifecycle: CLITest {
@@ -138,4 +139,43 @@ class TestCLIRunLifecycle: CLITest {
             )
         }
     }
+
+    @Test func testRunStartWithSSHSetsGuestAuthSockPath() throws {
+        guard ProcessInfo.processInfo.environment["SSH_AUTH_SOCK"] != nil else {
+            return
+        }
+
+        let name = getTestName()
+        defer {
+            try? doStop(name: name)
+            try? doRemove(name: name, force: true)
+        }
+
+        try doLongRun(name: name, args: ["--ssh"], autoRemove: false)
+        try waitForContainerRunning(name)
+
+        var output = try doExec(name: name, cmd: ["sh", "-lc", "echo -n ${SSH_AUTH_SOCK:-}"])
+        #expect(output.trimmingCharacters(in: .whitespacesAndNewlines) == "/run/host-services/ssh-auth.sock")
+
+        try doStop(name: name)
+        try doStart(name: name)
+        try waitForContainerRunning(name)
+
+        output = try doExec(name: name, cmd: ["sh", "-lc", "echo -n ${SSH_AUTH_SOCK:-}"])
+        #expect(output.trimmingCharacters(in: .whitespacesAndNewlines) == "/run/host-services/ssh-auth.sock")
+    }
+
+    @Test func testRunWithoutSSHDoesNotSetGuestAuthSockPath() throws {
+        let name = getTestName()
+        defer {
+            try? doStop(name: name)
+            try? doRemove(name: name, force: true)
+        }
+
+        try doLongRun(name: name, autoRemove: false)
+        try waitForContainerRunning(name)
+
+        let output = try doExec(name: name, cmd: ["sh", "-lc", "echo -n ${SSH_AUTH_SOCK:-}"])
+        #expect(output.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty)
+    }
 }
diff --git a/Tests/ContainerSandboxServiceTests/DynamicEnvMergeTests.swift b/Tests/ContainerSandboxServiceTests/DynamicEnvMergeTests.swift
@@ -0,0 +1,50 @@
+//===----------------------------------------------------------------------===//
+// Copyright © 2026 Apple Inc. and the container project authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//===----------------------------------------------------------------------===//
+
+import ContainerSandboxService
+import Testing
+
+struct DynamicEnvMergeTests {
+    @Test
+    func testDynamicEnvMergeAddsMissingKey() {
+        let overrides = ["SSH_AUTH_SOCK": "/run/host-services/ssh-auth.sock"]
+        let base = ["PATH=/usr/local/bin:/usr/bin:/bin:/usr/sbin:/sbin"]
+
+        let mergedEnv = SandboxService.mergedEnvironmentVariables(base: base, overrides: overrides)
+
+        #expect(mergedEnv.contains("PATH=/usr/local/bin:/usr/bin:/bin:/usr/sbin:/sbin"))
+        #expect(mergedEnv.contains("SSH_AUTH_SOCK=/run/host-services/ssh-auth.sock"))
+    }
+
+    @Test
+    func testDynamicEnvMergeOverridesExistingKey() {
+        let overrides = ["FOO": "updated"]
+        let base = ["FOO=original", "PATH=/usr/bin"]
+
+        let mergedEnv = SandboxService.mergedEnvironmentVariables(base: base, overrides: overrides)
+
+        #expect(mergedEnv.contains("FOO=updated"))
+        #expect(!mergedEnv.contains("FOO=original"))
+        #expect(mergedEnv.contains("PATH=/usr/bin"))
+    }
+
+    @Test
+    func testDynamicEnvMergeNoOverridesLeavesBaseUnchanged() {
+        let base = ["FOO=bar", "PATH=/usr/bin"]
+        let mergedEnv = SandboxService.mergedEnvironmentVariables(base: base, overrides: [:])
+        #expect(mergedEnv == base)
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -55,8 +55,13 @@ public struct ContainersHarness: Sendable {`
`55`	`55`	`)`
`56`	`56`	`}`
`57`	`57`	`let stdio = message.stdio()`
`58`		`- let sshAuthSocketPath = message.string(key: .sshAuthSocketPath)`
`59`		`- try await service.bootstrap(id: id, stdio: stdio, sshAuthSocketPath: sshAuthSocketPath)`
	`58`	`+ let dynamicEnv: [String: String] =`
	`59`	`+ if let data = message.dataNoCopy(key: .dynamicEnv) {`
	`60`	`+ try JSONDecoder().decode([String: String].self, from: data)`
	`61`	`+ } else {`
	`62`	`+ [:]`
	`63`	`+ }`
	`64`	`+ try await service.bootstrap(id: id, stdio: stdio, dynamicEnv: dynamicEnv)`
`60`	`65`	`return message.reply()`
`61`	`66`	`}`
`62`	`67`
Original file line number	Diff line number	Diff line change
`@@ -77,7 +77,8 @@ public struct SandboxClient: Sendable {`
`77`	`77`
`78`	`78`	`// Runtime Methods`
`79`	`79`	`extension SandboxClient {`
`80`		`- public func bootstrap(stdio: [FileHandle?], allocatedAttachments: [AllocatedAttachment]) async throws {`
	`80`	`+ /// - Parameter dynamicEnv: Optional start-time environment overrides passed from the API service.`
	`81`	`+ public func bootstrap(stdio: [FileHandle?], allocatedAttachments: [AllocatedAttachment], dynamicEnv: [String: String] = [:]) async throws {`
`81`	`82`	`let request = XPCMessage(route: SandboxRoutes.bootstrap.rawValue)`
`82`	`83`
`83`	`84`	`for (i, h) in stdio.enumerated() {`
`@@ -96,8 +97,9 @@ extension SandboxClient {`
`96`	`97`	`}`
`97`	`98`	`}`
`98`	`99`
`99`		`- if let sshAuthSocketPath {`
`100`		`- request.set(key: SandboxKeys.sshAuthSocketPath.rawValue, value: sshAuthSocketPath)`
	`100`	`+ if !dynamicEnv.isEmpty {`
	`101`	`+ let encodedDynamicEnv = try JSONEncoder().encode(dynamicEnv)`
	`102`	`+ request.set(key: SandboxKeys.dynamicEnv.rawValue, value: encodedDynamicEnv)`
`101`	`103`	`}`
`102`	`104`
`103`	`105`	`do {`