From bec1008f4de1e416ea0e112b57b2357875ec7d25 Mon Sep 17 00:00:00 2001
From: Danny Canter <danny_canter@apple.com>
Date: Mon, 3 Nov 2025 14:09:56 -0800
Subject: [PATCH] VsockConnectionStream: Conform to AsyncSequence (#380)

I didn't like how we expose the asyncstream via a public connections
field. We should have the type conform to AsyncSequence and then hide
the underlying stream. This also stops listening on the stdio ports
after we get the initial connection.
---
 Sources/Containerization/LinuxProcess.swift   |  8 ++++---
 .../Containerization/UnixSocketRelay.swift    |  2 +-
 .../VsockConnectionStream.swift               | 21 ++++++++++++++++---
 3 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/Sources/Containerization/LinuxProcess.swift b/Sources/Containerization/LinuxProcess.swift
index 69a498e9..bacc8d40 100644
--- a/Sources/Containerization/LinuxProcess.swift
+++ b/Sources/Containerization/LinuxProcess.swift
@@ -119,19 +119,21 @@ public final class LinuxProcess: Sendable {
 extension LinuxProcess {
     func setupIO(streams: [VsockConnectionStream?]) async throws -> [FileHandle?] {
         let handles = try await Timeout.run(seconds: 3) {
-            await withTaskGroup(of: (Int, FileHandle?).self) { group in
+            try await withThrowingTaskGroup(of: (Int, FileHandle?).self) { group in
                 var results = [FileHandle?](repeating: nil, count: 3)
 
                 for (index, stream) in streams.enumerated() {
                     guard let stream = stream else { continue }
 
                     group.addTask {
-                        let first = await stream.connections.first(where: { _ in true })
+                        let first = await stream.first(where: { _ in true })
+                        stream.finish()
+                        try self.vm.stopListen(stream.port)
                         return (index, first)
                     }
                 }
 
-                for await (index, fileHandle) in group {
+                for try await (index, fileHandle) in group {
                     results[index] = fileHandle
                 }
                 return results
diff --git a/Sources/Containerization/UnixSocketRelay.swift b/Sources/Containerization/UnixSocketRelay.swift
index 3f3184b3..2388ec7c 100644
--- a/Sources/Containerization/UnixSocketRelay.swift
+++ b/Sources/Containerization/UnixSocketRelay.swift
@@ -201,7 +201,7 @@ extension SocketRelay {
             $0.t = Task {
                 do {
                     defer { connectionStream.finish() }
-                    for await connection in connectionStream.connections {
+                    for await connection in connectionStream {
                         try await self.handleGuestVsockConn(
                             vsockConn: connection,
                             hostConnectionPath: hostPath,
diff --git a/Sources/Containerization/VsockConnectionStream.swift b/Sources/Containerization/VsockConnectionStream.swift
index 3bf43b4f..c985410a 100644
--- a/Sources/Containerization/VsockConnectionStream.swift
+++ b/Sources/Containerization/VsockConnectionStream.swift
@@ -21,9 +21,11 @@ import Virtualization
 #endif
 
 /// A stream of vsock connections.
-public final class VsockConnectionStream: NSObject, Sendable {
+public final class VsockConnectionStream: NSObject, Sendable, AsyncSequence {
+    public typealias Element = FileHandle
+
     /// A stream of connections dialed from the remote.
-    public let connections: AsyncStream<FileHandle>
+    private let connections: AsyncStream<FileHandle>
     /// The port the connections are for.
     public let port: UInt32
 
@@ -39,6 +41,10 @@ public final class VsockConnectionStream: NSObject, Sendable {
     public func finish() {
         self.cont.finish()
     }
+
+    public func makeAsyncIterator() -> AsyncStream<FileHandle>.AsyncIterator {
+        connections.makeAsyncIterator()
+    }
 }
 
 #if os(macOS)
@@ -49,9 +55,18 @@ extension VsockConnectionStream: VZVirtioSocketListenerDelegate {
         from _: VZVirtioSocketDevice
     ) -> Bool {
         let fd = dup(conn.fileDescriptor)
+        guard fd != -1 else {
+            return false
+        }
         conn.close()
 
-        cont.yield(FileHandle(fileDescriptor: fd, closeOnDealloc: false))
+        let fh = FileHandle(fileDescriptor: fd, closeOnDealloc: false)
+        let result = cont.yield(fh)
+        if case .terminated = result {
+            try? fh.close()
+            return false
+        }
+
         return true
     }
 }