chore: import upstream snapshot with attribution

Sources/ContainerizationOS/Socket/Socket.swift

@@ -0,0 +1,494 @@
+//===----------------------------------------------------------------------===//
+// 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 CShim
+import Foundation
+import Synchronization
+
+#if canImport(Musl)
+import Musl
+#elseif canImport(Glibc)
+import Glibc
+#elseif canImport(Darwin)
+import Darwin
+#else
+#error("Socket not supported on this platform.")
+#endif
+
+#if !os(Windows)
+let sysFchmod = fchmod
+let sysRead = read
+let sysUnlink = unlink
+let sysSend = send
+let sysClose = close
+let sysShutdown = shutdown
+let sysBind = bind
+let sysSocket = socket
+let sysSetsockopt = setsockopt
+let sysGetsockopt = getsockopt
+let sysListen = listen
+let sysAccept = accept
+let sysConnect = connect
+let sysIoctl: @convention(c) (CInt, CUnsignedLong, UnsafeMutableRawPointer) -> CInt = ioctl
+let sysRecvmsg = recvmsg
+#endif
+
+/// Thread-safe socket wrapper.
+public final class Socket: Sendable {
+    public enum TimeoutOption {
+        case send
+        case receive
+    }
+
+    public enum ShutdownOption {
+        case read
+        case write
+        case readWrite
+    }
+
+    private enum SocketState {
+        case created
+        case connected
+        case listening
+    }
+
+    private struct State {
+        let socketState: SocketState
+        let handle: FileHandle?
+        let type: SocketType
+        let acceptSource: DispatchSourceRead?
+    }
+
+    private let _closeOnDeinit: Bool
+    private let _queue: DispatchQueue
+
+    private let state: Mutex<State>
+
+    public var fileDescriptor: Int32 {
+        guard let handle = state.withLock({ $0.handle }) else {
+            return -1
+        }
+        return handle.fileDescriptor
+    }
+
+    public convenience init(type: SocketType, closeOnDeinit: Bool = true) throws {
+        let sockFD = sysSocket(type.domain, type.type, 0)
+        if sockFD < 0 {
+            throw SocketError.withErrno("failed to create socket: \(sockFD)", errno: errno)
+        }
+        self.init(fd: sockFD, type: type, closeOnDeinit: closeOnDeinit)
+    }
+
+    init(fd: Int32, type: SocketType, closeOnDeinit: Bool) {
+        _queue = DispatchQueue(label: "com.apple.containerization.socket")
+        _closeOnDeinit = closeOnDeinit
+        let state = State(
+            socketState: .created,
+            handle: FileHandle(fileDescriptor: fd, closeOnDealloc: false),
+            type: type,
+            acceptSource: nil
+        )
+        self.state = Mutex(state)
+    }
+
+    /// Internal initializer for wrapping already-connected file descriptors (e.g., from socketpair)
+    /// Ideally we just get rid of the state machine in this class. Not sure how much value it provides..
+    init(fd: Int32, type: SocketType, closeOnDeinit: Bool, connected: Bool) {
+        _queue = DispatchQueue(label: "com.apple.containerization.socket")
+        _closeOnDeinit = closeOnDeinit
+        let state = State(
+            socketState: connected ? .connected : .created,
+            handle: FileHandle(fileDescriptor: fd, closeOnDealloc: false),
+            type: type,
+            acceptSource: nil
+        )
+        self.state = Mutex(state)
+    }
+
+    deinit {
+        if _closeOnDeinit {
+            try? close()
+        }
+    }
+}
+
+extension Socket {
+    static func errnoToError(msg: String) -> SocketError {
+        SocketError.withErrno("\(msg) (\(_errnoString(errno)))", errno: errno)
+    }
+
+    public func connect() throws {
+        try state.withLock { currentState in
+            guard currentState.socketState == .created else {
+                throw SocketError.invalidOperationOnSocket("connect")
+            }
+            guard let handle = currentState.handle else {
+                throw SocketError.closed
+            }
+
+            var res: Int32 = 0
+            try currentState.type.withSockAddr { (ptr, length) in
+                res = Syscall.retrying {
+                    sysConnect(handle.fileDescriptor, ptr, length)
+                }
+            }
+
+            if res == -1 {
+                throw Socket.errnoToError(msg: "could not connect to socket \(currentState.type)")
+            }
+
+            currentState = State(
+                socketState: .connected,
+                handle: handle,
+                type: currentState.type,
+                acceptSource: currentState.acceptSource
+            )
+        }
+    }
+
+    public func listen() throws {
+        try state.withLock { currentState in
+            guard currentState.socketState == .created else {
+                throw SocketError.invalidOperationOnSocket("listen")
+            }
+            guard let handle = currentState.handle else {
+                throw SocketError.closed
+            }
+
+            try currentState.type.beforeBind(fd: handle.fileDescriptor)
+
+            var rc: Int32 = 0
+            try currentState.type.withSockAddr { (ptr, length) in
+                rc = sysBind(handle.fileDescriptor, ptr, length)
+            }
+
+            if rc < 0 {
+                throw Socket.errnoToError(msg: "could not bind to \(currentState.type)")
+            }
+
+            try currentState.type.beforeListen(fd: handle.fileDescriptor)
+
+            if sysListen(handle.fileDescriptor, SOMAXCONN) < 0 {
+                throw Socket.errnoToError(msg: "listen failed on \(currentState.type)")
+            }
+
+            currentState = State(
+                socketState: .listening,
+                handle: handle,
+                type: currentState.type,
+                acceptSource: currentState.acceptSource
+            )
+        }
+    }
+
+    public func close() throws {
+        try state.withLock { currentState in
+            guard let handle = currentState.handle else {
+                // Already closed.
+                return
+            }
+
+            let acceptSource = currentState.acceptSource
+
+            acceptSource?.cancel()
+            try handle.close()
+
+            currentState = State(
+                socketState: currentState.socketState,
+                handle: nil,
+                type: currentState.type,
+                acceptSource: nil
+            )
+        }
+    }
+
+    public func write(data: any DataProtocol) throws -> Int {
+        let handle = try state.withLock { currentState in
+            guard currentState.socketState == .connected else {
+                throw SocketError.invalidOperationOnSocket("write")
+            }
+            guard let handle = currentState.handle else {
+                throw SocketError.closed
+            }
+            return handle
+        }
+
+        if data.isEmpty {
+            return 0
+        }
+
+        try handle.write(contentsOf: data)
+        return data.count
+    }
+
+    public func acceptStream(closeOnDeinit: Bool = true) throws -> AsyncThrowingStream<Socket, Swift.Error> {
+        let source = try state.withLock { currentState -> DispatchSourceRead in
+            guard currentState.socketState == .listening else {
+                throw SocketError.invalidOperationOnSocket("accept")
+            }
+            guard let handle = currentState.handle else {
+                throw SocketError.closed
+            }
+            guard currentState.acceptSource == nil else {
+                throw SocketError.acceptStreamExists
+            }
+
+            let source = DispatchSource.makeReadSource(
+                fileDescriptor: handle.fileDescriptor,
+                queue: _queue
+            )
+
+            currentState = State(
+                socketState: currentState.socketState,
+                handle: handle,
+                type: currentState.type,
+                acceptSource: source
+            )
+
+            return source
+        }
+
+        return AsyncThrowingStream { cont in
+            source.setCancelHandler {
+                cont.finish()
+            }
+            source.setEventHandler(handler: {
+                if source.data == 0 {
+                    source.cancel()
+                    return
+                }
+
+                do {
+                    let connection = try self.accept(closeOnDeinit: closeOnDeinit)
+                    cont.yield(connection)
+                } catch SocketError.closed {
+                    source.cancel()
+                } catch {
+                    cont.yield(with: .failure(error))
+                    source.cancel()
+                }
+            })
+            source.activate()
+        }
+    }
+
+    public func accept(closeOnDeinit: Bool = true) throws -> Socket {
+        let (handle, socketType) = try state.withLock { currentState in
+            guard currentState.socketState == .listening else {
+                throw SocketError.invalidOperationOnSocket("accept")
+            }
+            guard let handle = currentState.handle else {
+                throw SocketError.closed
+            }
+            return (handle, currentState.type)
+        }
+
+        let (clientFD, newSocketType) = try socketType.accept(fd: handle.fileDescriptor)
+        return Socket(
+            fd: clientFD,
+            type: newSocketType,
+            closeOnDeinit: closeOnDeinit,
+            connected: true
+        )
+    }
+
+    /// Receive a file descriptor via SCM_RIGHTS control message.
+    /// This is commonly used for passing file descriptors between processes via Unix domain sockets.
+    public func receiveFileDescriptor() throws -> Int32 {
+        let handle = try state.withLock { currentState in
+            guard currentState.socketState == .connected else {
+                throw SocketError.invalidOperationOnSocket("receiveFileDescriptor")
+            }
+            guard let handle = currentState.handle else {
+                throw SocketError.closed
+            }
+            return handle
+        }
+
+        var msg = msghdr()
+        var iov = iovec()
+        var buf: UInt8 = 0
+
+        iov.iov_base = withUnsafeMutablePointer(to: &buf) { UnsafeMutableRawPointer($0) }
+        iov.iov_len = 1
+
+        msg.msg_iov = withUnsafeMutablePointer(to: &iov) { $0 }
+        msg.msg_iovlen = 1
+
+        var cmsgBuf = [UInt8](repeating: 0, count: Int(CZ_CMSG_SPACE(Int(MemoryLayout<Int32>.size))))
+        msg.msg_control = withUnsafeMutablePointer(to: &cmsgBuf[0]) { UnsafeMutableRawPointer($0) }
+        msg.msg_controllen = numericCast(cmsgBuf.count)
+
+        let recvResult = withUnsafeMutablePointer(to: &msg) { msgPtr in
+            sysRecvmsg(handle.fileDescriptor, msgPtr, 0)
+        }
+
+        guard recvResult >= 0 else {
+            throw Socket.errnoToError(msg: "recvmsg failed")
+        }
+
+        // Extract file descriptor from control message
+        let cmsgPtr = withUnsafeMutablePointer(to: &msg) { CZ_CMSG_FIRSTHDR($0) }
+        guard let cmsg = cmsgPtr else {
+            throw SocketError.invalidFileDescriptor
+        }
+
+        guard cmsg.pointee.cmsg_level == SOL_SOCKET,
+            cmsg.pointee.cmsg_type == SCM_RIGHTS
+        else {
+            throw SocketError.invalidFileDescriptor
+        }
+
+        guard let dataPtr = CZ_CMSG_DATA(cmsg) else {
+            throw SocketError.invalidFileDescriptor
+        }
+
+        let fdPtr = dataPtr.assumingMemoryBound(to: Int32.self)
+        let fd = fdPtr.pointee
+        guard fd >= 0 else {
+            throw SocketError.invalidFileDescriptor
+        }
+
+        return fd
+    }
+
+    public func read(buffer: inout Data) throws -> Int {
+        let handle = try state.withLock { currentState in
+            guard currentState.socketState == .connected else {
+                throw SocketError.invalidOperationOnSocket("read")
+            }
+            guard let handle = currentState.handle else {
+                throw SocketError.closed
+            }
+            return handle
+        }
+
+        var bytesRead = 0
+        let bufferSize = buffer.count
+        try buffer.withUnsafeMutableBytes { pointer in
+            guard let baseAddress = pointer.baseAddress else {
+                throw SocketError.missingBaseAddress
+            }
+
+            bytesRead = Syscall.retrying {
+                sysRead(handle.fileDescriptor, baseAddress, bufferSize)
+            }
+            if bytesRead < 0 {
+                throw Socket.errnoToError(msg: "error reading from connection")
+            } else if bytesRead == 0 {
+                throw SocketError.closed
+            }
+        }
+        return bytesRead
+    }
+
+    public func shutdown(how: ShutdownOption) throws {
+        let handle = try state.withLock { currentState in
+            guard let handle = currentState.handle else {
+                throw SocketError.closed
+            }
+            return handle
+        }
+
+        var howOpt: Int32 = 0
+        switch how {
+        case .read:
+            howOpt = Int32(SHUT_RD)
+        case .write:
+            howOpt = Int32(SHUT_WR)
+        case .readWrite:
+            howOpt = Int32(SHUT_RDWR)
+        }
+
+        if sysShutdown(handle.fileDescriptor, howOpt) < 0 {
+            throw Socket.errnoToError(msg: "shutdown failed")
+        }
+    }
+
+    public func setSockOpt(sockOpt: Int32 = 0, ptr: UnsafeRawPointer, stride: UInt32) throws {
+        let handle = try state.withLock { currentState in
+            guard let handle = currentState.handle else {
+                throw SocketError.closed
+            }
+            return handle
+        }
+
+        if setsockopt(handle.fileDescriptor, SOL_SOCKET, sockOpt, ptr, stride) < 0 {
+            throw Socket.errnoToError(msg: "failed to set sockopt")
+        }
+    }
+
+    public func setTimeout(option: TimeoutOption, seconds: Int) throws {
+        let handle = try state.withLock { currentState in
+            guard let handle = currentState.handle else {
+                throw SocketError.closed
+            }
+            return handle
+        }
+
+        var sockOpt: Int32 = 0
+        switch option {
+        case .receive:
+            sockOpt = SO_RCVTIMEO
+        case .send:
+            sockOpt = SO_SNDTIMEO
+        }
+
+        var timer = timeval()
+        timer.tv_sec = seconds
+        timer.tv_usec = 0
+
+        if setsockopt(
+            handle.fileDescriptor,
+            SOL_SOCKET,
+            sockOpt,
+            &timer,
+            socklen_t(MemoryLayout<timeval>.size)
+        ) < 0 {
+            throw Socket.errnoToError(msg: "failed to set read timeout")
+        }
+    }
+
+    static func _errnoString(_ err: Int32?) -> String {
+        String(validatingCString: strerror(errno)) ?? "error: \(errno)"
+    }
+}
+
+public enum SocketError: Error, Equatable, CustomStringConvertible {
+    case closed
+    case acceptStreamExists
+    case invalidOperationOnSocket(String)
+    case missingBaseAddress
+    case withErrno(_ msg: String, errno: Int32)
+    case invalidFileDescriptor
+
+    public var description: String {
+        switch self {
+        case .closed:
+            return "socket: closed"
+        case .acceptStreamExists:
+            return "accept stream already exists"
+        case .invalidOperationOnSocket(let operation):
+            return "socket: invalid operation on socket '\(operation)'"
+        case .missingBaseAddress:
+            return "socket: missing base address"
+        case .withErrno(let msg, _):
+            return "socket: error \(msg)"
+        case .invalidFileDescriptor:
+            return "socket: invalid file descriptor received"
+        }
+    }
+}