Sesame-Server/Sources/App/DeviceManager.swift

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<Bool>

    private let messagesToDeviceMetric: Metric<Int>
    
    /// 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<Data>?
    
    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<Data> {
        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
    }
}