diff --git a/FlyingSocks/Sources/AsyncSocket.swift b/FlyingSocks/Sources/AsyncSocket.swift index 045f070b..72c516f1 100644 --- a/FlyingSocks/Sources/AsyncSocket.swift +++ b/FlyingSocks/Sources/AsyncSocket.swift @@ -129,6 +129,20 @@ public struct AsyncSocket: Sendable { return buffer } + public func receive(atMost length: Int = 4096) async throws -> (Socket.Address, [UInt8]) { + try Task.checkCancellation() + + repeat { + do { + return try socket.receive(length: length) + } catch SocketError.blocked { + try await pool.suspendSocket(socket, untilReadyFor: .read) + } catch { + throw error + } + } while true + } + /// Reads bytes from the socket up to by not over/ /// - Parameter bytes: The max number of bytes to read /// - Returns: an array of the read bytes capped to the number of bytes provided. @@ -163,6 +177,15 @@ public struct AsyncSocket: Sendable { } } + public func send(_ data: Data, to address: some SocketAddress) async throws { + let sent = try await pool.loopUntilReady(for: .write, on: socket) { + try socket.send(Array(data), to: address) + } + guard sent == data.count else { + throw SocketError.disconnected + } + } + public func close() throws { try socket.close() } diff --git a/FlyingSocks/Sources/Socket+Darwin.swift b/FlyingSocks/Sources/Socket+Darwin.swift index efa3a643..d0197350 100644 --- a/FlyingSocks/Sources/Socket+Darwin.swift +++ b/FlyingSocks/Sources/Socket+Darwin.swift @@ -42,6 +42,7 @@ extension Socket.FileDescriptor { extension Socket { static let stream = Int32(SOCK_STREAM) + static let datagram = Int32(SOCK_DGRAM) static let in_addr_any = Darwin.in_addr(s_addr: Darwin.in_addr_t(0)) static func makeAddressINET(port: UInt16) -> Darwin.sockaddr_in { @@ -157,14 +158,22 @@ extension Socket { Darwin.read(fd, buffer, nbyte) } + static func recvfrom(_ fd: FileDescriptorType, _ buffer: UnsafeMutableRawPointer!, _ nbyte: Int, _ flags: Int32, _ addr: UnsafeMutablePointer!, _ len: UnsafeMutablePointer!) -> Int { + Darwin.recvfrom(fd, buffer, nbyte, flags, addr, len) + } + static func write(_ fd: FileDescriptorType, _ buffer: UnsafeRawPointer!, _ nbyte: Int) -> Int { Darwin.write(fd, buffer, nbyte) } + static func sendto(_ fd: FileDescriptorType, _ buffer: UnsafeRawPointer!, _ nbyte: Int, _ flags: Int32, _ destaddr: UnsafePointer!, _ destlen: socklen_t) -> Int { + Darwin.sendto(fd, buffer, nbyte, flags, destaddr, destlen) + } + static func close(_ fd: FileDescriptorType) -> Int32 { Darwin.close(fd) } - + static func unlink(_ addr: UnsafePointer!) -> Int32 { Darwin.unlink(addr) } diff --git a/FlyingSocks/Sources/Socket.swift b/FlyingSocks/Sources/Socket.swift index 7de00526..8f1043e2 100644 --- a/FlyingSocks/Sources/Socket.swift +++ b/FlyingSocks/Sources/Socket.swift @@ -214,6 +214,35 @@ public struct Socket: Sendable, Hashable { return count } + public func receive(length: Int) throws -> (Socket.Address, [UInt8]) { + var address: sockaddr_storage? + let bytes = try [UInt8](unsafeUninitializedCapacity: length) { buffer, count in + (address, count) = try receive(into: buffer.baseAddress!, length: length) + } + + return try (Self.makeAddress(from: address!), bytes) + } + + private func receive(into buffer: UnsafeMutablePointer, length: Int) throws -> (sockaddr_storage, Int) { + var addr = sockaddr_storage() + var size = socklen_t(MemoryLayout.size) + let count = withUnsafeMutablePointer(to: &addr) { + $0.withMemoryRebound(to: sockaddr.self, capacity: 1) { + Socket.recvfrom(file.rawValue, buffer, length, 0, $0, &size) + } + } + guard count > 0 else { + if errno == EWOULDBLOCK { + throw SocketError.blocked + } else if errno == EBADF || count == 0 { + throw SocketError.disconnected + } else { + throw SocketError.makeFailed("RecvFrom") + } + } + return (addr, count) + } + public func write(_ data: Data, from index: Data.Index = 0) throws -> Data.Index { precondition(index >= 0) guard index < data.endIndex else { return data.endIndex } @@ -237,6 +266,29 @@ public struct Socket: Sendable, Hashable { return sent } + public func send(_ bytes: [UInt8], to address: some SocketAddress) throws -> Int { + try bytes.withUnsafeBytes { buffer in + try send(buffer.baseAddress!, length: bytes.count, to: address) + } + } + + private func send(_ pointer: UnsafeRawPointer, length: Int, to address: A) throws -> Int { + var addr = address + let sent = withUnsafePointer(to: &addr) { + $0.withMemoryRebound(to: sockaddr.self, capacity: 1) { + Socket.sendto(file.rawValue, pointer, length, 0, $0, socklen_t(MemoryLayout.size)) + } + } + guard sent >= 0 || errno == EISCONN else { + if errno == EINPROGRESS { + throw SocketError.blocked + } else { + throw SocketError.makeFailed("SendTo") + } + } + return sent + } + public func close() throws { if Socket.close(file.rawValue) == -1 { throw SocketError.makeFailed("Close") diff --git a/FlyingSocks/Sources/SocketAddress.swift b/FlyingSocks/Sources/SocketAddress.swift index 68839421..90389e8c 100644 --- a/FlyingSocks/Sources/SocketAddress.swift +++ b/FlyingSocks/Sources/SocketAddress.swift @@ -186,3 +186,17 @@ extension Socket { } } } + +public extension Socket.Address { + + func makeSocketAddress() throws -> any SocketAddress { + switch self { + case let .unix(path): + return .unix(path: path) + case let .ip4(ip, port: port): + return try .inet(ip4: ip, port: port) + case let .ip6(ip, port: port): + return try .inet6(ip6: ip, port: port) + } + } +}