wehub-resource-sync
495 lines
15 KiB
Swift
495 lines
15 KiB
Swift
//===----------------------------------------------------------------------===//
|
|
// 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"
|
|
}
|
|
}
|
|
}
|