SOCK_DGRAM from [email protected]

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()
     }

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<sockaddr>!, _ len: UnsafeMutablePointer<socklen_t>!) -> 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<sockaddr>!, _ 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<CChar>!) -> Int32 {
         Darwin.unlink(addr)
     }

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<UInt8>, length: Int) throws -> (sockaddr_storage, Int) {
+        var addr = sockaddr_storage()
+        var size = socklen_t(MemoryLayout<sockaddr_storage>.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<A: SocketAddress>(_ 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<A>.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")

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)
+        }
+    }
+}