From a35fc88e689bdc069884dc0bfdc36f181228eab1 Mon Sep 17 00:00:00 2001 From: Steve Date: Sat, 5 Sep 2026 22:17:28 +0100 Subject: [PATCH] fix(vsock): retain connection owner after handoff Keep the Virtualization.framework VSOCK connection object alive while duplicated descriptors remain in use by FileHandle and the gRPC transport. Retain accepted connection owners through the FileHandle lifetime, retain the vminitd handoff handle for the client lifetime, and cover both ownership boundaries with focused regressions. Related: apple/containerization#678 --- .../VZVirtualMachine+Helpers.swift | 5 +- Sources/Containerization/Vminitd.swift | 3 + Sources/Containerization/VsockListener.swift | 19 ++++- .../VsockListenerTests.swift | 76 +++++++++++++++++++ 4 files changed, 100 insertions(+), 3 deletions(-) create mode 100644 Tests/ContainerizationTests/VsockListenerTests.swift diff --git a/Sources/Containerization/VZVirtualMachine+Helpers.swift b/Sources/Containerization/VZVirtualMachine+Helpers.swift index 2cbadb1c8..6231ca9c3 100644 --- a/Sources/Containerization/VZVirtualMachine+Helpers.swift +++ b/Sources/Containerization/VZVirtualMachine+Helpers.swift @@ -144,8 +144,9 @@ extension VZVirtioSocketConnection { if fd == -1 { throw POSIXError.fromErrno() } - self.close() - return FileHandle(fileDescriptor: fd, closeOnDealloc: false) + let handle = FileHandle(fileDescriptor: fd, closeOnDealloc: false) + retainConnectionOwner(self, for: handle) + return handle } } diff --git a/Sources/Containerization/Vminitd.swift b/Sources/Containerization/Vminitd.swift index 7fded4432..bc94a2150 100644 --- a/Sources/Containerization/Vminitd.swift +++ b/Sources/Containerization/Vminitd.swift @@ -31,6 +31,8 @@ public struct Vminitd: Sendable { let client: Com_Apple_Containerization_Sandbox_V3_SandboxContext.Client public let grpcClient: GRPCClient + // The FileHandle retains the Virtualization connection that owns this descriptor. + private let retainedConnection: FileHandle private let connectionTask: Task public init(connection: FileHandle, group: any EventLoopGroup) async throws { @@ -54,6 +56,7 @@ public struct Vminitd: Sendable { let grpcClient = GRPCClient(transport: transport) self.grpcClient = grpcClient self.client = Com_Apple_Containerization_Sandbox_V3_SandboxContext.Client(wrapping: self.grpcClient) + self.retainedConnection = connection // Not very structured concurrency friendly, but we'd need to expose a way on the protocol to "run" the // agent otherwise, which some agents might not even need. self.connectionTask = Task { diff --git a/Sources/Containerization/VsockListener.swift b/Sources/Containerization/VsockListener.swift index 0a20d81cf..ac4e5a9db 100644 --- a/Sources/Containerization/VsockListener.swift +++ b/Sources/Containerization/VsockListener.swift @@ -18,6 +18,7 @@ import Foundation import Synchronization #if os(macOS) +import ObjectiveC import Virtualization #endif @@ -76,6 +77,22 @@ public final class VsockListener: NSObject, Sendable, AsyncSequence { #if os(macOS) +private final class ConnectionOwnerAssociationKey: @unchecked Sendable {} + +private let connectionOwnerAssociationKey = ConnectionOwnerAssociationKey() + +/// Keeps an owner of a duplicated descriptor alive for as long as its file +/// handle can use that descriptor. Virtualization owns the original VSOCK +/// descriptor and requires its connection object to stay alive after accept. +func retainConnectionOwner(_ owner: AnyObject, for handle: FileHandle) { + objc_setAssociatedObject( + handle, + Unmanaged.passUnretained(connectionOwnerAssociationKey).toOpaque(), + owner, + .OBJC_ASSOCIATION_RETAIN_NONATOMIC + ) +} + extension VsockListener: VZVirtioSocketListenerDelegate { public func listener( _: VZVirtioSocketListener, shouldAcceptNewConnection conn: VZVirtioSocketConnection, @@ -85,9 +102,9 @@ extension VsockListener: VZVirtioSocketListenerDelegate { guard fd != -1 else { return false } - conn.close() let fh = FileHandle(fileDescriptor: fd, closeOnDealloc: false) + retainConnectionOwner(conn, for: fh) let result = cont.yield(fh) if case .terminated = result { try? fh.close() diff --git a/Tests/ContainerizationTests/VsockListenerTests.swift b/Tests/ContainerizationTests/VsockListenerTests.swift new file mode 100644 index 000000000..5917316a1 --- /dev/null +++ b/Tests/ContainerizationTests/VsockListenerTests.swift @@ -0,0 +1,76 @@ +//===----------------------------------------------------------------------===// +// 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 +// +// http://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. +//===----------------------------------------------------------------------===// + +#if os(macOS) +import Darwin +import Foundation +import NIO +import Testing + +@testable import Containerization + +struct VsockListenerTests { + @Test func connectionOwnerOutlivesTheDescriptorHandoff() throws { + var handle: FileHandle? = FileHandle(forReadingAtPath: "/dev/null") + try #require(handle != nil) + + var owner: ConnectionOwner? = ConnectionOwner() + weak let weakOwner = owner + retainConnectionOwner(owner!, for: handle!) + owner = nil + + #expect(weakOwner != nil) + + try handle?.close() + handle = nil + + #expect(weakOwner == nil) + } + + @Test func vminitdRetainsConnectionOwnerWhileUsingDescriptor() async throws { + var descriptors = [Int32](repeating: -1, count: 2) + guard socketpair(AF_UNIX, SOCK_STREAM, 0, &descriptors) == 0 else { + throw POSIXError(.init(rawValue: errno) ?? .EIO) + } + + let peer = FileHandle(fileDescriptor: descriptors[1], closeOnDealloc: true) + let group = MultiThreadedEventLoopGroup.singleton + defer { try? peer.close() } + + weak var weakOwner: ConnectionOwner? + var client: Vminitd? + do { + let handle = FileHandle(fileDescriptor: descriptors[0], closeOnDealloc: false) + var owner: ConnectionOwner? = ConnectionOwner() + weakOwner = owner + retainConnectionOwner(owner!, for: handle) + + client = try await Vminitd(connection: handle, group: group) + owner = nil + } + + #expect(weakOwner != nil) + + try peer.close() + try? await client?.close() + client = nil + + #expect(weakOwner == nil) + } +} + +private final class ConnectionOwner {} +#endif