import Foundation import WebSocketKit import Vapor import Clairvoyant final class DeviceManager { /// The connection to the device private var connection: WebSocket? /// The authentication token of the device for the socket connection private let deviceKey: Data /// The authentication token of the remote private let remoteKey: Data /// Indicate that the socket is fully initialized with an authorized device private var deviceIsAuthenticated = false private var isOpeningNewConnection = false private let deviceTimeout: Int64 private let deviceConnectedMetric: Metric private let messagesToDeviceMetric: Metric /// Indicator for device availability var deviceIsConnected: Bool { deviceIsAuthenticated && !(connection?.isClosed ?? true) } /// A promise to finish the request once the device responds or times out private var requestInProgress: EventLoopPromise? init(deviceKey: Data, remoteKey: Data, deviceTimeout: Int64) async { self.deviceKey = deviceKey self.remoteKey = remoteKey self.deviceTimeout = deviceTimeout self.deviceConnectedMetric = try! await .init( "sesame.connected", name: "Device connected", description: "Shows if the device is connected via WebSocket") self.messagesToDeviceMetric = try! await .init( "sesame.messages", name: "Forwarded Messages", description: "The number of messages transmitted to the device") } private func updateDeviceConnectionMetric() { Task { try? await deviceConnectedMetric.update(deviceIsConnected) } } private func updateMessageCountMetric() { Task { let lastValue = await messagesToDeviceMetric.lastValue()?.value ?? 0 try? await messagesToDeviceMetric.update(lastValue + 1) } } // MARK: API private var deviceStatus: String { deviceIsConnected ? "1" : "0" } func sendMessageToDevice(_ message: Data, on eventLoop: EventLoop) -> EventLoopFuture { guard let socket = connection, !socket.isClosed else { connection = nil return eventLoop.makeSucceededFuture(MessageResult.deviceNotConnected.encoded) } guard requestInProgress == nil else { return eventLoop.makeSucceededFuture(MessageResult.operationInProgress.encoded) } let result = eventLoop.makePromise(of: Data.self) self.requestInProgress = result socket.send(Array(message), promise: nil) updateMessageCountMetric() eventLoop.scheduleTask(in: .seconds(deviceTimeout)) { [weak self] in guard let promise = self?.requestInProgress else { return } self?.requestInProgress = nil log("Timed out waiting for device response") promise.succeed(MessageResult.deviceTimedOut.encoded) } return result.futureResult } func authenticateDevice(hash: String) { defer { updateDeviceConnectionMetric() } guard let key = Data(fromHexEncodedString: hash), SHA256.hash(data: key) == self.deviceKey else { log("Invalid device key") _ = connection?.close() deviceIsAuthenticated = false return } log("Device authenticated") deviceIsAuthenticated = true } func authenticateRemote(_ token: Data) -> Bool { let hash = SHA256.hash(data: token) return hash == remoteKey } func processDeviceResponse(_ buffer: ByteBuffer) { guard let data = buffer.getData(at: 0, length: buffer.readableBytes) else { log("Failed to get data buffer received from device") return } guard let promise = requestInProgress else { log("Received device response \(data) without an active request") return } defer { requestInProgress = nil } log("Device response received") promise.succeed(data) } func didCloseDeviceSocket() { defer { updateDeviceConnectionMetric() } guard !isOpeningNewConnection else { return } deviceIsAuthenticated = false guard connection != nil else { log("Socket closed, but no connection anyway") return } connection = nil log("Socket closed") } func removeDeviceConnection() { defer { updateDeviceConnectionMetric() } deviceIsAuthenticated = false guard let socket = connection else { return } try? socket.close().wait() connection = nil log("Removed device connection") } func createNewDeviceConnection(_ socket: WebSocket) { defer { updateDeviceConnectionMetric() } socket.onBinary { _, data in self.processDeviceResponse(data) } socket.onText { _, text in self.authenticateDevice(hash: text) } _ = socket.onClose.always { _ in self.didCloseDeviceSocket() } isOpeningNewConnection = true removeDeviceConnection() connection = socket log("Socket connected") isOpeningNewConnection = false } }