Fix socket descriptor ownership and lifetime

Sources/CShim/include/socket_helpers.h

@@ -21,9 +21,21 @@
 #include <stdint.h>
 
 // Helper functions to access CMSG macros from Swift
-struct cmsghdr* CZ_CMSG_FIRSTHDR(struct msghdr *msg);
 void* CZ_CMSG_DATA(struct cmsghdr *cmsg);
 size_t CZ_CMSG_SPACE(size_t length);
 size_t CZ_CMSG_LEN(size_t length);
 
+// Receives exactly one SCM_RIGHTS descriptor and sets FD_CLOEXEC before returning it.
+// On failure, it returns -1 and closes the descriptors in each well-formed control
+// record; parsing stops at a malformed record. errno is one of:
+// - a recvmsg or fcntl error
+// - EMSGSIZE for truncated control data, usually from descriptor exhaustion
+// - EPROTO for other invalid control data
+// - ENODATA for ordinary EOF
+// The caller owns the result and handles readiness and EINTR.
+int CZ_receive_fd(int socket_fd);
+
+// Linux-only accept4(2) with SOCK_CLOEXEC; glibc declares accept4 only under _GNU_SOURCE.
+int CZ_accept_cloexec(int socket_fd);
+
 #endif /* socket_helpers_h */

Sources/CShim/socket_helpers.c

@@ -14,11 +14,21 @@
  * limitations under the License.
  */
 
+#if defined(__linux__)
+#define _GNU_SOURCE
+#endif
+
 #include "socket_helpers.h"
 
-struct cmsghdr* CZ_CMSG_FIRSTHDR(struct msghdr *msg) {
-    return CMSG_FIRSTHDR(msg);
-}
+#include <errno.h>
+#include <fcntl.h>
+#include <string.h>
+#include <unistd.h>
+
+// Linux UAPI value; older libc headers may not expose it yet.
+#if defined(__linux__) && !defined(SCM_PIDFD)
+#define SCM_PIDFD 0x04
+#endif
 
 void* CZ_CMSG_DATA(struct cmsghdr *cmsg) {
     return CMSG_DATA(cmsg);
@@ -31,3 +41,114 @@ size_t CZ_CMSG_SPACE(size_t length) {
 size_t CZ_CMSG_LEN(size_t length) {
     return CMSG_LEN(length);
 }
+
+int CZ_receive_fd(int socket_fd) {
+    // XNU installs every right before copying out control data, so truncation would
+    // hide installed descriptors. Its rights are bounded by MCLBYTES of internal
+    // pointers (254 rights, 1028 bytes copied out); Linux allows 253 rights plus
+    // credential and pidfd records. 4 KiB holds the largest set either kernel delivers.
+    union {
+        struct cmsghdr alignment;
+        unsigned char bytes[4096];
+    } control;
+    unsigned char byte;
+    struct iovec iov = { .iov_base = &byte, .iov_len = sizeof(byte) };
+    struct msghdr message = {0};
+    message.msg_iov = &iov;
+    message.msg_iovlen = 1;
+    message.msg_control = control.bytes;
+    message.msg_controllen = sizeof(control.bytes);
+
+    int flags = 0;
+#ifdef __linux__
+    flags = MSG_CMSG_CLOEXEC;
+#endif
+    ssize_t count = recvmsg(socket_fd, &message, flags);
+    if (count < 0) {
+        return -1;
+    }
+
+    int received = -1;
+    size_t offset = 0;
+    size_t length = message.msg_controllen;
+    int invalid = (message.msg_flags & MSG_CTRUNC) != 0;
+    if (length > sizeof(control.bytes)) {
+        length = sizeof(control.bytes);
+        invalid = 1;
+    }
+    while (length - offset >= sizeof(struct cmsghdr)) {
+        struct cmsghdr header;
+        memcpy(&header, control.bytes + offset, sizeof(header));
+        size_t header_length = CMSG_LEN(0);
+        if (header.cmsg_len < header_length || header.cmsg_len > length - offset) {
+            // A malformed length cannot establish which bytes are descriptors.
+            invalid = 1;
+            break;
+        }
+        size_t payload_length = header.cmsg_len - header_length;
+        if (header.cmsg_level == SOL_SOCKET && header.cmsg_type == SCM_RIGHTS) {
+            if (payload_length == 0 || payload_length % sizeof(int) != 0) {
+                invalid = 1;
+            }
+            for (size_t i = 0; i + sizeof(int) <= payload_length; i += sizeof(int)) {
+                int fd;
+                memcpy(&fd, control.bytes + offset + header_length + i, sizeof(fd));
+                if (fd < 0) {
+                    invalid = 1;
+                } else if (received < 0) {
+                    received = fd;
+                } else {
+                    // Never retry close: the descriptor number may already be reused.
+                    close(fd);
+                    invalid = 1;
+                }
+            }
+#ifdef __linux__
+        } else if (header.cmsg_level == SOL_SOCKET && header.cmsg_type == SCM_PIDFD) {
+            // SO_PASSPIDFD installs a descriptor even though this record is rejected.
+            invalid = 1;
+            if (payload_length >= sizeof(int)) {
+                int fd;
+                memcpy(&fd, control.bytes + offset + header_length, sizeof(fd));
+                if (fd >= 0) {
+                    close(fd);
+                }
+            }
+#endif
+        } else {
+            invalid = 1;
+        }
+        size_t stride = CMSG_SPACE(payload_length);
+        if (stride > length - offset) {
+            // Native trailing alignment padding may be omitted.
+            break;
+        }
+        offset += stride;
+    }
+
+    int error = (message.msg_flags & MSG_CTRUNC) != 0 ? EMSGSIZE : EPROTO;
+    if (!invalid && received >= 0) {
+#ifndef __linux__
+        int descriptor_flags = fcntl(received, F_GETFD);
+        if (descriptor_flags >= 0 && fcntl(received, F_SETFD, descriptor_flags | FD_CLOEXEC) >= 0) {
+            return received;
+        }
+        error = errno;
+#else
+        return received;
+#endif
+    } else if (!invalid && received < 0 && count == 0) {
+        error = ENODATA;
+    }
+    if (received >= 0) {
+        close(received);
+    }
+    errno = error;
+    return -1;
+}
+
+#if defined(__linux__)
+int CZ_accept_cloexec(int socket_fd) {
+    return accept4(socket_fd, NULL, NULL, SOCK_CLOEXEC);
+}
+#endif

Sources/Containerization/CHStdioPortSlot.swift

@@ -15,7 +15,9 @@
 //===----------------------------------------------------------------------===//
 
 #if os(Linux)
+import CShim
 import ContainerizationError
+import ContainerizationOS
 import Foundation
 import Logging
 import Synchronization
@@ -48,14 +50,17 @@ final class CHStdioPortSlot: Sendable {
     let port: UInt32
     let path: URL
     let listenFd: Int32
+    private let descriptor: OwnedDescriptor
 
     private struct State {
         /// The listener entitled to accepted connections right now.
-        var owner: VsockListener?
+        weak var owner: VsockListener?
+        var claim: UUID?
         /// Whether the accept loop is running. Once started it runs until
-        /// `shutdown()`, and it owns closing `listenFd`.
+        /// `shutdown()`, holding a descriptor use until it exits.
         var accepting: Bool
-        /// Set by `shutdown()`. Blocks further claims and stops the loop.
+        /// Set by `shutdown()` or a failed accept loop. Blocks further claims
+        /// and stops the loop.
         var closed: Bool
     }
 
@@ -70,11 +75,12 @@ final class CHStdioPortSlot: Sendable {
     /// fds per slot plus a lock to keep them from being written after reuse.
     private static let pollTimeoutMilliseconds: Int32 = 1000
 
-    init(port: UInt32, path: URL, listenFd: Int32) {
+    init(port: UInt32, path: URL, listenDescriptor: consuming DescriptorTransfer) {
         self.port = port
         self.path = path
-        self.listenFd = listenFd
-        self.state = Mutex(State(owner: nil, accepting: false, closed: false))
+        self.descriptor = listenDescriptor.takeOwner()
+        self.listenFd = descriptor.fileDescriptor
+        self.state = Mutex(State(owner: nil, claim: nil, accepting: false, closed: false))
     }
 
     /// Lend the slot to `listener`.
@@ -88,26 +94,37 @@ final class CHStdioPortSlot: Sendable {
             guard !state.closed else {
                 throw ContainerizationError(.invalidState, message: "vsock port \(port) is closed")
             }
-            guard state.owner == nil else {
+            // Loading the weak owner here can run its deinit under this mutex,
+            // reentering relinquish when the temporary strong reference drops.
+            guard state.claim == nil else {
                 throw ContainerizationError(
                     .invalidState,
                     message: "vsock port \(port) is already being listened on"
                 )
             }
             state.owner = listener
+            state.claim = listener.identity
         }
     }
 
-    /// Give up ownership without disturbing the socket. The accept loop keeps
-    /// running for the next tenant; connections that arrive in between are
+    /// Give up the matching claim without disturbing the socket. The accept
+    /// loop keeps running for the next tenant; connections that arrive in between are
     /// dropped by the loop.
-    func relinquish() {
-        state.withLock { $0.owner = nil }
+    func relinquish(_ identity: UUID) {
+        state.withLock {
+            guard $0.claim == identity else { return }
+            $0.owner = nil
+            $0.claim = nil
+        }
+    }
+
+    deinit {
+        try? descriptor.close()
     }
 
     /// Start the accept loop, if this is the slot's first tenant. The loop then
     /// runs until `shutdown()`, so ownership handoff never has to stop and
-    /// restart it — which is what makes `relinquish()`/`claim(by:)` safe
+    /// restart it — which is what makes `relinquish(_:)`/`claim(by:)` safe
     /// back-to-back with no settling period.
     func startAcceptingIfNeeded(logger: Logger?) {
         let shouldStart = state.withLock { state -> Bool in
@@ -117,7 +134,7 @@ final class CHStdioPortSlot: Sendable {
         }
         guard shouldStart else { return }
 
-        // The loop blocks in poll(2)/accept(2), which is inappropriate for
+        // The loop blocks in poll(2), which is inappropriate for
         // Swift's cooperative thread pool: a pool thread parked in a syscall
         // can't service other tasks until it returns. With even a few of these,
         // detached tasks queue behind the parked threads and never run, which
@@ -131,33 +148,41 @@ final class CHStdioPortSlot: Sendable {
 
     /// Stop the accept loop and release the socket. Idempotent, and safe to
     /// call whether or not a loop was ever started.
+    /// The descriptor closes after the loop releases its active use.
     func shutdown() {
-        // A running loop owns `listenFd` and closes it on its way out; with no
-        // loop there is nobody else to do it.
-        let closeHere = state.withLock { state -> Bool in
-            guard !state.closed else { return false }
+        // A running loop keeps `listenFd` alive through its descriptor use;
+        // with no active loop, closing the descriptor releases it immediately.
+        let listener = state.withLock { state -> VsockListener? in
+            guard !state.closed else { return nil }
             state.closed = true
+            let listener = state.owner
             state.owner = nil
-            return !state.accepting
-        }
-        if closeHere {
-            _ = close(listenFd)
+            state.claim = nil
+            return listener
         }
+        try? listener?.cancelAndDiscard()
+        try? descriptor.close()
     }
 
     private func acceptLoop(logger: Logger?) {
         logger?.debug("vsock acceptLoop starting port=\(port) listenFd=\(listenFd)")
+        guard let use = try? descriptor.acquire() else { return }
         defer {
-            _ = close(listenFd)
+            use.end()
+            try? descriptor.close()
             // Mark the slot closed on the way out, not just on the `shutdown()`
             // path. A slot whose loop has stopped can never serve another
             // tenant, and because it stays in the pool the allocator will hand
             // its port out again; without this, every later claim would succeed
             // and then silently wait out its whole dial-back window.
-            state.withLock {
+            let listener = state.withLock {
                 $0.closed = true
+                let listener = $0.owner
                 $0.owner = nil
+                $0.claim = nil
+                return listener
             }
+            try? listener?.cancelAndDiscard()
             logger?.debug("vsock acceptLoop exited port=\(port)")
         }
 
@@ -184,9 +209,18 @@ final class CHStdioPortSlot: Sendable {
                 return
             }
 
-            let connFd = accept(listenFd, nil, nil)
+            // Accept is nonblocking. Pair it with the current claim so a handover
+            // before delivery cannot send this connection to the next listener.
+            let accepted = state.withLock { state -> (fd: Int32, error: Int32, owner: VsockListener?)? in
+                guard !state.closed else { return nil }
+                let owner = state.owner
+                let fd = CZ_accept_cloexec(listenFd)
+                return (fd, errno, owner)
+            }
+            guard let accepted else { return }
+            let connFd = accepted.fd
             if connFd < 0 {
-                let savedErrno = errno
+                let savedErrno = accepted.error
                 if savedErrno == EINTR || savedErrno == EAGAIN || savedErrno == EWOULDBLOCK || savedErrno == ECONNABORTED {
                     continue
                 }
@@ -203,21 +237,19 @@ final class CHStdioPortSlot: Sendable {
                 return
             }
 
-            let handle = FileHandle(fileDescriptor: connFd, closeOnDealloc: true)
-            guard let owner = state.withLock({ $0.owner }) else {
+            let connection = DescriptorTransfer(adopting: connFd)
+            guard let owner = accepted.owner else {
                 // A dial for a stream that already gave up the port — e.g. a
                 // guest that got to its connect(2) after the host stopped
                 // waiting for it. Dropping it here is what keeps it from being
                 // delivered to the port's next tenant.
                 logger?.warning("vsock dial on port \(port) with no listener; dropping")
-                try? handle.close()
                 continue
             }
-            if case .terminated = owner.yield(handle) {
-                try? handle.close()
+            if !owner.yieldOwned(connection, closeOnDealloc: true) {
                 // That listener can never take another connection, so free the
                 // slot rather than wedging it, and keep accepting.
-                relinquish()
+                relinquish(owner.identity)
             }
         }
     }