From a65641226c3bcf6f27470a98e6d5a3cc9f528e0e Mon Sep 17 00:00:00 2001
From: Steve <stephenlclarke@mac.com>
Date: Sat, 5 Sep 2026 22:17:28 +0100
Subject: [PATCH] fix(vsock): retain connection owner after handoff

---
 .../VZVirtualMachine+Helpers.swift            |  5 ++-
 Sources/Containerization/VsockListener.swift  | 19 +++++++-
 .../VsockListenerTests.swift                  | 43 +++++++++++++++++++
 3 files changed, 64 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/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..e3563a0eb
--- /dev/null
+++ b/Tests/ContainerizationTests/VsockListenerTests.swift
@@ -0,0 +1,43 @@
+//===----------------------------------------------------------------------===//
+// 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 Foundation
+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)
+    }
+}
+
+private final class ConnectionOwner {}
+#endif