Swift Foundation WebSocket

WebSocket.swift

import Foundation

public final class WebSocket {
    public typealias OnOpen = (URLSessionWebSocketTask, _ protocol: String?) -> Void
    public typealias OnClose = (URLSessionWebSocketTask, _ closeCode: URLSessionWebSocketTask.CloseCode, _ reason: Data?) -> Void
    public typealias OnMessage = (URLSessionWebSocketTask.Message) -> Void
    public typealias OnError = (Error) -> Void

    public enum ReadyState: Int {
        case initial = -1, connecting = 0, `open` = 1, closing = 2, closed = 3
    }

    private let object = WebSocketObject()

    @ObjectProxy(\.onOpen)    public var onOpen
    @ObjectProxy(\.onClose)   public var onClose
    @ObjectProxy(\.onMessage) public var onMessage
    @ObjectProxy(\.onError)   public var onError

    public var readyState: ReadyState { object.readyState }

    public func connect(url: URL) {
        object.connect(url: url)
    }

    public func close() {
        object.disconnect()
    }

    public func send(_ message: URLSessionWebSocketTask.Message, completionHandler: @escaping ((any Error)?) -> Void) {
        guard let task = object.task else { preconditionFailure("WebSocket.connect(url:) should be called before send") }
        task.send(message, completionHandler: completionHandler)
    }

    public func send(_ message: URLSessionWebSocketTask.Message) async throws {
        guard let task = object.task else { preconditionFailure("WebSocket.connect(url:) should be called before send") }
        try await task.send(message)
    }

    private final class WebSocketObject: NSObject, URLSessionWebSocketDelegate {
        var readyState: ReadyState = .initial
        var onOpen: OnOpen?
        var onClose: OnClose?
        var onMessage: OnMessage?
        var onError: OnError?
        var task: URLSessionWebSocketTask?

        lazy var session = URLSession(configuration: .default,
                                      delegate: self,
                                      delegateQueue: nil)

        deinit { disconnect() }

        func connect(url: URL) {
            precondition(readyState == .initial, "WebSocket.connect(url:) should be called once only")
            readyState = .connecting
            let task = session.webSocketTask(with: url)
            self.task = task
            task.resume()
        }

        func disconnect() {
            readyState = .closing
            task?.cancel(with: .normalClosure, reason: nil)
        }

        func setupReceiveMessage() {
            guard let task, readyState == .open else { return }
            task.receive { [weak self] result in
                guard let self else { return }
                switch result {
                case .success(let message):
                    self.onMessage?(message)
                    self.setupReceiveMessage()
                case .failure(let error):
                    self.onError?(error)
                }
            }
        }

        func urlSession(_ session: URLSession,
                        webSocketTask: URLSessionWebSocketTask,
                        didOpenWithProtocol protocol: String?) {
            readyState = .open
            setupReceiveMessage()
            onOpen?(webSocketTask, `protocol`)
        }

        func urlSession(_ session: URLSession,
                        webSocketTask: URLSessionWebSocketTask,
                        didCloseWith closeCode: URLSessionWebSocketTask.CloseCode,
                        reason: Data?) {
            readyState = .closed
            task = nil
            onClose?(webSocketTask, closeCode, reason)
        }
    }

    @propertyWrapper
    private struct ObjectProxy<Value> {
        let keyPath: ReferenceWritableKeyPath<WebSocketObject, Value>

        init(_ keyPath: ReferenceWritableKeyPath<WebSocketObject, Value>) {
            self.keyPath = keyPath
        }

        @available(*, unavailable)
        var wrappedValue: Value {
            get { fatalError() } set { fatalError() }
        }

        static subscript(_enclosingInstance instance: WebSocket,
                         wrapped wrappedKeyPath: ReferenceWritableKeyPath<WebSocket, Value>,
                         storage storageKeyPath: ReferenceWritableKeyPath<WebSocket, Self>) -> Value {
            get {
                let `self` = instance[keyPath: storageKeyPath]
                return instance.object[keyPath: self.keyPath]
            }
            set {
                let `self` = instance[keyPath: storageKeyPath]
                instance.object[keyPath: self.keyPath] = newValue
            }
        }
    }
}