spb/os-vault Public

1import Foundation23public protocol SolanaSocket {4    /// Connection status of the socket5    var isConnected: Bool { get }67    /// Delegation8    var delegate: SolanaSocketEventsDelegate? { get set }910    /// Connect to socket11    func connect()1213    /// Disconnect from socket14    func disconnect()1516    /// Subscribe to an entity ('account', 'program', 'signature', for example)17    /// - Parameters:18    ///   - type: type of entity, '.account', '.program',...19    ///   - params: params to be sent20    /// - Returns: id of the request21    @discardableResult func subscribe<T: Encodable>(22        type: SocketEntity,23        params: T,24        commitment: String,25        encoding: String26    ) async throws27        -> String2829    /// Unsubscribe to an entity ('account', 'program', 'signature', for example)30    /// - Parameters:31    ///   - type: type of entity, '.account', '.program',...32    ///   - socketId: id of the subscription33    /// - Returns: id of the request34    @discardableResult func unsubscribe(type: SocketEntity, socketId: UInt64) async throws -> String35}3637public extension SolanaSocket {38    /// Subscribe to `accountNotification`39    /// - Parameter publickey: account to be subscribed40    /// - Returns: id of the request41    @discardableResult func accountSubscribe(publickey: String, commitment: String = "recent") async throws -> String {42        try await subscribe(type: .account, params: publickey, commitment: commitment, encoding: "jsonParsed")43    }4445    /// Subscribe to `signatureNotification`46    /// - Parameter signature: signature to be subscribed47    /// - Returns: id of the request48    @discardableResult func signatureSubscribe(signature: String,49                                               commitment: String = "confirmed") async throws -> String50    {51        try await subscribe(type: .signature, params: signature, commitment: commitment, encoding: "base64")52    }5354    /// Subscribe to `logsNotification`55    /// - Parameter mentions: accounts to be subscribed56    /// - Returns: id of the request57    @discardableResult func logsSubscribe(mentions: [String], commitment: String = "confirmed") async throws -> String {58        try await subscribe(type: .logs, params: ["mentions": mentions], commitment: commitment, encoding: "base64")59    }6061    /// Subscribe to all events62    /// - Returns: id of the request63    @discardableResult func logsSubscribeAll(commitment: String = "confirmed") async throws -> String {64        try await subscribe(type: .logs, params: "all", commitment: commitment, encoding: "base64")65    }6667    /// Subscribe to `programNotification`68    /// - Parameter publickey: program to be subscribed69    /// - Returns: id of the request70    @discardableResult func programSubscribe(publickey: String,71                                             commitment: String = "confirmed") async throws -> String72    {73        try await subscribe(type: .program, params: publickey, commitment: commitment, encoding: "base64")74    }75}7677public class Socket: NSObject, SolanaSocket {78    // MARK: - Properties7980    /// Connection status of the socket81    public var isConnected: Bool = false8283    /// Socket task to handle socket event84    private var task: WebSocketTask!8586    /// Timer to send pings to prevent idie time out87    private var wsHeartBeat: Timer!8889    /// Async task to keep track of asynchronous receiving task90    private var asyncTask: Task<Void, Error>?9192    /// Delegation93    public weak var delegate: SolanaSocketEventsDelegate?9495    // MARK: - Initializers9697    /// Initializer for Socket98    /// - Parameters:99    ///   - url: url of the socket100    ///   - enableDebugLogs: enable/disable logging101    ///   - socketTaskProviderType: type of task provider, default is `URLSession.self`102    public init<T: WebSocketTaskProvider>(103        url: URL,104        socketTaskProviderType _: T.Type105    ) {106        super.init()107        let urlSession = T(configuration: .default, delegate: self, delegateQueue: .current!)108        task = urlSession.createWebSocketTask(with: url)109    }110111    /// Convenience initializer for socket using `URLSession` as `WebSocketTaskProvider`112    /// - Parameters:113    ///   - url: url of the socket114    ///   - enableDebugLogs: enable/disable logging115    public convenience init(116        url: URL117    ) {118        self.init(url: url, socketTaskProviderType: URLSession.self)119    }120121    deinit {122        disconnect()123    }124125    // MARK: - Methods126127    /// Connect to socket128    public func connect() {129        task.resume()130    }131132    /// Disconnect from socket133    public func disconnect() {134        delegate?.disconnected(reason: "", code: 0)135136        asyncTask?.cancel()137        task.cancel()138        delegate = nil139        wsHeartBeat?.invalidate()140        wsHeartBeat = nil141    }142143    /// Subscribe to an entity ('account', 'program', 'signature', for example)144    /// - Parameters:145    ///   - type: type of entity, '.account', '.program',...146    ///   - params: params to be sent147    /// - Returns: id of the request148    @discardableResult public func subscribe<T: Encodable>(149        type entity: SocketEntity, params: T,150        commitment: String,151        encoding: String152    ) async throws -> String {153        let method: SocketMethod = .init(entity, .subscribe)154        let params: [Encodable] = [params, ["commitment": commitment, "encoding": encoding]]155        let request = RequestAPI(method: method.rawValue, params: params)156        return try await writeToSocket(request: request)157    }158159    /// Unsubscribe to an entity ('account', 'program', 'signature', for example)160    /// - Parameters:161    ///   - type: type of entity, '.account', '.program',...162    ///   - socketId: id of the subscription163    /// - Returns: id of the request164    @discardableResult public func unsubscribe(type entity: SocketEntity, socketId: UInt64) async throws -> String {165        let method: SocketMethod = .init(entity, .unsubscribe)166        let params: [Encodable] = [socketId]167        let request = RequestAPI(method: method.rawValue, params: params)168        return try await writeToSocket(request: request)169    }170171    /// Emit message to socket172    /// - Parameter request: request to be sent173    /// - Returns: request id174    @discardableResult private func writeToSocket(request: RequestAPI) async throws -> String {175        let jsonData = try JSONEncoder().encode(request)176        Logger.log(event: "request", message: "\(String(data: jsonData, encoding: .utf8) ?? "")", logLevel: .info)177        try await task.send(.data(jsonData))178        return request.id179    }180181    /// Read message from socket one at a time182    private func readMessage() async throws {183        try Task.checkCancellation()184        let message = try await task.receive()185        switch message {186        case let .string(text):187            Logger.log(event: "event", message: "Receive string from socket: \(text)", logLevel: .debug)188            guard let data = text.data(using: .utf8) else { return }189            do {190                // TODO: Fix this mess code191                let jsonResponse = try JSONSerialization.jsonObject(with: data, options: []) as! [String: Any]192                if let jsonType = jsonResponse["method"] as? String,193                   let type = SocketMethod(rawValue: jsonType),194                   type.action == .notification195                {196                    switch type.entity {197                    case .account:198                        if let notification = try? JSONDecoder()199                            .decode(SocketNativeAccountNotification.self, from: data)200                        {201                            delegate?.nativeAccountNotification(notification: notification)202                        } else {203                            let notification = try JSONDecoder().decode(SocketTokenAccountNotification.self, from: data)204                            delegate?.tokenAccountNotification(notification: notification)205                        }206207                    case .signature:208                        let notification = try JSONDecoder().decode(SocketSignatureNotification.self, from: data)209                        delegate?.signatureNotification(notification: notification)210                    case .logs:211                        let notification = try JSONDecoder().decode(SocketLogsNotification.self, from: data)212                        delegate?.logsNotification(notification: notification)213                    case .program:214                        let notification = try JSONDecoder().decode(SocketProgramAccountNotification.self, from: data)215                        delegate?.programNotification(notification: notification)216                    default:217                        break218                    }219220                } else {221                    if let subscription = try? JSONDecoder().decode(SocketSubscriptionResponse.self, from: data),222                       let socketId = subscription.result,223                       let id = subscription.id224                    {225                        delegate?.subscribed(socketId: socketId, id: id)226                    }227228                    if let subscription = try? JSONDecoder().decode(SocketUnsubscriptionResponse.self, from: data),229                       subscription.result == true,230                       let id = subscription.id231                    {232                        delegate?.unsubscribed(id: id)233                    }234                }235            } catch {236                delegate?.error(error: error)237            }238        case let .data(data):239            print("Received binary message: \(data)")240        @unknown default:241            fatalError()242        }243    }244245    private func ping() {246        Logger.log(event: "request", message: "Ping socket", logLevel: .debug)247        task.sendPing { error in248            if let error = error {249                print("Ping failed: \(error)")250            }251        }252    }253}254255extension Socket: URLSessionWebSocketDelegate {256    public func urlSession(_: URLSession, webSocketTask _: URLSessionWebSocketTask, didOpenWithProtocol _: String?) {257        isConnected = true258        wsHeartBeat?.invalidate()259        wsHeartBeat = Timer.scheduledTimer(withTimeInterval: 5, repeats: true) { [weak self] _ in260            // Ping server every 5s to prevent idle timeouts261            self?.ping()262        }263        delegate?.connected()264265        Logger.log(event: "urlSession", message: "Socket disconnected", logLevel: .debug)266267        asyncTask = Task.detached { [weak self] in268            while true {269                guard let self = self else { break }270                try await self.readMessage()271            }272        }273    }274275    public func urlSession(276        _: URLSession,277        webSocketTask _: URLSessionWebSocketTask,278        didCloseWith closeCode: URLSessionWebSocketTask.CloseCode,279        reason: Data?280    ) {281        isConnected = false282        wsHeartBeat?.invalidate()283        task.resume()284        delegate?.disconnected(reason: reason?.jsonString ?? "", code: closeCode.rawValue)285286        Logger.log(event: "urlSession", message: "Socket disconnected", logLevel: .debug)287288        asyncTask?.cancel()289    }290}291