Fix remaining async vsock lifetime bugs

DePasqualeOrg · DePasqualeOrg · commit c4fa3e61cb3c · 2026-03-15T23:17:59.000+01:00
diff --git a/Sources/Containerization/LinuxContainer.swift b/Sources/Containerization/LinuxContainer.swift
@@ -969,7 +969,7 @@ extension LinuxContainer {
     }
 
     /// Dial a vsock port in the container.
-    public func dialVsock(port: UInt32) async throws -> FileHandle {
+    public func dialVsock(port: UInt32) async throws -> VsockConnection {
         try await self.state.withLock {
             let state = try $0.startedState("dialVsock")
             return try await state.vm.dial(port)
@@ -1098,7 +1098,7 @@ extension LinuxContainer {
                     try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, any Error>) in
                         self.copyQueue.async {
                             do {
-                                defer { conn.closeFile() }
+                                defer { try? conn.close() }
 
                                 if isArchive {
                                     let writer = try ArchiveWriter(configuration: .init(format: .pax, filter: .gzip))
@@ -1209,7 +1209,7 @@ extension LinuxContainer {
                     try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, any Error>) in
                         self.copyQueue.async {
                             do {
-                                defer { conn.closeFile() }
+                                defer { try? conn.close() }
 
                                 if metadata.isArchive {
                                     try FileManager.default.createDirectory(at: destination, withIntermediateDirectories: true)
diff --git a/Sources/Containerization/LinuxPod.swift b/Sources/Containerization/LinuxPod.swift
@@ -850,7 +850,7 @@ extension LinuxPod {
     }
 
     /// Dial a vsock port in the pod's VM.
-    public func dialVsock(port: UInt32) async throws -> FileHandle {
+    public func dialVsock(port: UInt32) async throws -> VsockConnection {
         try await self.state.withLock { state in
             let createdState = try state.phase.createdState("dialVsock")
             return try await createdState.vm.dial(port)
diff --git a/Sources/Containerization/LinuxProcess.swift b/Sources/Containerization/LinuxProcess.swift
@@ -48,26 +48,44 @@ public final class LinuxProcess: Sendable {
     }
 
     private struct StdioHandles: Sendable {
-        var stdin: FileHandle?
-        var stdout: FileHandle?
-        var stderr: FileHandle?
+        var stdin: VsockConnection?
+        var stdout: VsockConnection?
+        var stderr: VsockConnection?
 
         mutating func close() throws {
+            var firstError: Error?
+
             if let stdin {
-                try stdin.close()
                 stdin.readabilityHandler = nil
+                do {
+                    try stdin.close()
+                } catch {
+                    firstError = firstError ?? error
+                }
                 self.stdin = nil
             }
             if let stdout {
-                try stdout.close()
                 stdout.readabilityHandler = nil
+                do {
+                    try stdout.close()
+                } catch {
+                    firstError = firstError ?? error
+                }
                 self.stdout = nil
             }
             if let stderr {
-                try stderr.close()
                 stderr.readabilityHandler = nil
+                do {
+                    try stderr.close()
+                } catch {
+                    firstError = firstError ?? error
+                }
                 self.stderr = nil
             }
+
+            if let firstError {
+                throw firstError
+            }
         }
     }
 
@@ -124,10 +142,10 @@ public final class LinuxProcess: Sendable {
 }
 
 extension LinuxProcess {
-    func setupIO(listeners: [VsockListener?]) async throws -> [FileHandle?] {
+    func setupIO(listeners: [VsockListener?]) async throws -> [VsockConnection?] {
         let handles = try await Timeout.run(seconds: 3) {
-            try await withThrowingTaskGroup(of: (Int, FileHandle?).self) { group in
-                var results = [FileHandle?](repeating: nil, count: 3)
+            try await withThrowingTaskGroup(of: (Int, VsockConnection?).self) { group in
+                var results = [VsockConnection?](repeating: nil, count: 3)
 
                 for (index, listener) in listeners.enumerated() {
                     guard let listener else { continue }
@@ -196,7 +214,7 @@ extension LinuxProcess {
         return handles
     }
 
-    func startStdinRelay(handle: FileHandle) {
+    func startStdinRelay(handle: VsockConnection) {
         guard let stdin = self.ioSetup.stdin else { return }
 
         self.state.withLock {
diff --git a/Sources/Containerization/UnixSocketRelay.swift b/Sources/Containerization/UnixSocketRelay.swift
@@ -29,8 +29,13 @@ package final class UnixSocketRelay: Sendable {
     private let log: Logger?
     private let state: Mutex<State>
 
+    private struct ActiveRelay: Sendable {
+        let relay: BidirectionalRelay
+        let guestConnection: VsockConnection
+    }
+
     private struct State {
-        var activeRelays: [String: BidirectionalRelay] = [:]
+        var activeRelays: [String: ActiveRelay] = [:]
         var t: Task<(), Never>? = nil
         var listener: VsockListener? = nil
     }
@@ -75,10 +80,9 @@ extension UnixSocketRelay {
             }
             t.cancel()
             $0.t = nil
-            for (_, relay) in $0.activeRelays {
-                relay.stop()
+            for (_, activeRelay) in $0.activeRelays {
+                activeRelay.relay.stop()
             }
-            $0.activeRelays.removeAll()
 
             switch configuration.direction {
             case .outOf:
@@ -170,12 +174,12 @@ extension UnixSocketRelay {
                 "initiating connection from host to guest",
                 metadata: [
                     "vport": "\(port)",
-                    "hostFd": "\(guestConn.fileDescriptor)",
-                    "guestFd": "\(hostConn.fileDescriptor)",
+                    "hostFd": "\(hostConn.fileDescriptor)",
+                    "guestFd": "\(guestConn.fileDescriptor)",
                 ])
             try await self.relay(
                 hostConn: hostConn,
-                guestFd: guestConn.fileDescriptor
+                guestConn: guestConn
             )
         } catch {
             log?.error("failed to relay between vsock \(port) and \(hostConn)")
@@ -184,7 +188,7 @@ extension UnixSocketRelay {
     }
 
     private func handleGuestVsockConn(
-        vsockConn: FileHandle,
+        vsockConn: VsockConnection,
         hostConnectionPath: URL,
         port: UInt32,
         log: Logger?
@@ -207,7 +211,7 @@ extension UnixSocketRelay {
         do {
             try await self.relay(
                 hostConn: hostSocket,
-                guestFd: vsockConn.fileDescriptor
+                guestConn: vsockConn
             )
         } catch {
             log?.error("failed to relay between vsock \(port) and \(hostPath)")
@@ -216,9 +220,13 @@ extension UnixSocketRelay {
 
     private func relay(
         hostConn: Socket,
-        guestFd: Int32
+        guestConn: VsockConnection
     ) async throws {
         let hostFd = hostConn.fileDescriptor
+        let guestFd = dup(guestConn.fileDescriptor)
+        if guestFd == -1 {
+            throw POSIXError.fromErrno()
+        }
 
         let relayID = UUID().uuidString
         let relay = BidirectionalRelay(
@@ -229,9 +237,21 @@ extension UnixSocketRelay {
         )
 
         state.withLock {
-            $0.activeRelays[relayID] = relay
+            // Retain the original connection until the relay has fully completed.
+            // The relay owns its duplicated fd and will close it itself.
+            $0.activeRelays[relayID] = ActiveRelay(
+                relay: relay,
+                guestConnection: guestConn
+            )
         }
 
         relay.start()
+
+        Task {
+            await relay.waitForCompletion()
+            let _ = self.state.withLock {
+                $0.activeRelays.removeValue(forKey: relayID)
+            }
+        }
     }
 }
diff --git a/Sources/Containerization/VZVirtualMachine+Helpers.swift b/Sources/Containerization/VZVirtualMachine+Helpers.swift
@@ -141,27 +141,22 @@ extension VZVirtualMachine {
 }
 
 extension VZVirtioSocketConnection {
-    /// Duplicates the file descriptor and immediately closes the connection.
+    /// Duplicates the file descriptor and retains the originating vsock connection
+    /// until the returned connection is closed or deallocated.
     ///
-    /// Only safe when the returned fd is used synchronously before any
-    /// suspension point. For deferred use (e.g., gRPC/NIO), use
-    /// ``dupFileDescriptor()`` and keep the connection alive via
-    /// ``VsockTransport``.
-    func dupHandle() throws -> FileHandle {
-        let fd = dup(self.fileDescriptor)
-        if fd == -1 {
-            throw POSIXError.fromErrno()
-        }
-        self.close()
-        return FileHandle(fileDescriptor: fd, closeOnDealloc: false)
+    /// Use this for file descriptors which cross an async boundary or may not be
+    /// consumed immediately by the caller.
+    func retainedConnection() throws -> VsockConnection {
+        try VsockConnection(connection: self)
     }
 
     /// Duplicates the connection's file descriptor without closing the connection.
     ///
     /// The caller must keep the `VZVirtioSocketConnection` alive until the dup'd
     /// descriptor is no longer needed. The Virtualization framework tears down the
     /// vsock endpoint when the connection is closed, which invalidates dup'd
-    /// descriptors.
+    /// descriptors. This is intended for callers which manage lifetime separately,
+    /// such as gRPC transports stored on `Vminitd`.
     func dupFileDescriptor() throws -> FileHandle {
         let fd = dup(self.fileDescriptor)
         if fd == -1 {
diff --git a/Sources/Containerization/VZVirtualMachineInstance.swift b/Sources/Containerization/VZVirtualMachineInstance.swift
@@ -202,14 +202,14 @@ extension VZVirtualMachineInstance: VirtualMachineInstance {
         }
     }
 
-    func dial(_ port: UInt32) async throws -> FileHandle {
+    func dial(_ port: UInt32) async throws -> VsockConnection {
         try await lock.withLock { _ in
             do {
                 let conn = try await vm.connect(
                     queue: queue,
                     port: port
                 )
-                return try conn.dupHandle()
+                return try conn.retainedConnection()
             } catch {
                 if let err = error as? ContainerizationError {
                     throw err
diff --git a/Sources/Containerization/VirtualMachineInstance.swift b/Sources/Containerization/VirtualMachineInstance.swift
@@ -38,7 +38,7 @@ public protocol VirtualMachineInstance: Sendable {
     /// what port the agent is listening on.
     func dialAgent() async throws -> Agent
     /// Dial a vsock port in the guest.
-    func dial(_ port: UInt32) async throws -> FileHandle
+    func dial(_ port: UInt32) async throws -> VsockConnection
     /// Listen on a host vsock port.
     func listen(_ port: UInt32) throws -> VsockListener
     /// Start the virtual machine.
diff --git a/Sources/Containerization/VsockConnection.swift b/Sources/Containerization/VsockConnection.swift
@@ -0,0 +1,92 @@
+//===----------------------------------------------------------------------===//
+// Copyright © 2025-2026 Apple Inc. and the Containerization 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 Foundation
+
+#if os(macOS)
+import Virtualization
+#endif
+
+/// A vsock connection whose duplicated file descriptor keeps the originating
+/// transport alive until the connection is closed.
+///
+/// Uses `@unchecked Sendable` because the mutable close state is protected by
+/// `NSLock`, while the underlying `FileHandle` and `VsockTransport` are shared
+/// across tasks.
+public final class VsockConnection: @unchecked Sendable {
+    private let fileHandle: FileHandle
+    private let transport: VsockTransport
+    private let lock = NSLock()
+    private var isClosed = false
+
+#if os(macOS)
+    init(connection: VZVirtioSocketConnection) throws {
+        let fd = dup(connection.fileDescriptor)
+        if fd == -1 {
+            throw POSIXError.fromErrno()
+        }
+        self.fileHandle = FileHandle(fileDescriptor: fd, closeOnDealloc: false)
+        self.transport = VsockTransport(connection)
+    }
+#endif
+
+    init(fileDescriptor: Int32, transport: VsockTransport) {
+        self.fileHandle = FileHandle(fileDescriptor: fileDescriptor, closeOnDealloc: false)
+        self.transport = transport
+    }
+
+    public var fileDescriptor: Int32 {
+        fileHandle.fileDescriptor
+    }
+
+    public var readabilityHandler: (@Sendable (FileHandle) -> Void)? {
+        get { fileHandle.readabilityHandler }
+        set { fileHandle.readabilityHandler = newValue }
+    }
+
+    public var availableData: Data {
+        fileHandle.availableData
+    }
+
+    public func write(contentsOf data: some DataProtocol) throws {
+        try fileHandle.write(contentsOf: data)
+    }
+
+    public func close() throws {
+        try closeIfNeeded {
+            try fileHandle.close()
+        }
+    }
+
+    private func closeIfNeeded(_ closeUnderlying: () throws -> Void) throws {
+        lock.lock()
+        guard !isClosed else {
+            lock.unlock()
+            return
+        }
+        isClosed = true
+        lock.unlock()
+
+        defer { transport.close() }
+        try closeUnderlying()
+    }
+
+    deinit {
+        try? closeIfNeeded {
+            try fileHandle.close()
+        }
+    }
+}
diff --git a/Sources/Containerization/VsockListener.swift b/Sources/Containerization/VsockListener.swift
diff --git a/Sources/Containerization/VsockTransport.swift b/Sources/Containerization/VsockTransport.swift
diff --git a/Tests/ContainerizationTests/VsockTransportTests.swift b/Tests/ContainerizationTests/VsockTransportTests.swift

Original file line number	Diff line number	Diff line change
`@@ -850,7 +850,7 @@ extension LinuxPod {`
`850`	`850`	`}`
`851`	`851`
`852`	`852`	`/// Dial a vsock port in the pod's VM.`
`853`		`- public func dialVsock(port: UInt32) async throws -> FileHandle {`
	`853`	`+ public func dialVsock(port: UInt32) async throws -> VsockConnection {`
`854`	`854`	`try await self.state.withLock { state in`
`855`	`855`	`let createdState = try state.phase.createdState("dialVsock")`
`856`	`856`	`return try await createdState.vm.dial(port)`
Original file line number	Diff line number	Diff line change
`@@ -202,14 +202,14 @@ extension VZVirtualMachineInstance: VirtualMachineInstance {`
`202`	`202`	`}`
`203`	`203`	`}`
`204`	`204`
`205`		`- func dial(_ port: UInt32) async throws -> FileHandle {`
	`205`	`+ func dial(_ port: UInt32) async throws -> VsockConnection {`
`206`	`206`	`try await lock.withLock { _ in`
`207`	`207`	`do {`
`208`	`208`	`let conn = try await vm.connect(`
`209`	`209`	`queue: queue,`
`210`	`210`	`port: port`
`211`	`211`	`)`
`212`		`- return try conn.dupHandle()`
	`212`	`+ return try conn.retainedConnection()`
`213`	`213`	`} catch {`
`214`	`214`	`if let err = error as? ContainerizationError {`
`215`	`215`	`throw err`