Sesame-Server/Sources/App/DeviceManager.swift

import Foundation
import WebSocketKit
import Vapor
import Clairvoyant

enum DeviceState: UInt8 {
    
    case disconnected = 0
    case connected = 1
    case authenticated = 2
}

extension DeviceState: MetricValue {
    
    static let valueType: MetricType = .customType(named: "DeviceState")
}

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 let deviceTimeout: Int64

    private let deviceStateMetric: Metric<DeviceState>

    private let messagesToDeviceMetric: Metric<Int>
    
    private let scheduler: AsyncScheduler
    
    var deviceState: DeviceState {
        guard let connection, !connection.isClosed else {
            return .disconnected
        }
        guard deviceIsAuthenticated else {
            return .connected
        }
        return .authenticated
    }
    
    /// 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: CheckedContinuation<Data, Error>?
    
    init(deviceKey: Data, remoteKey: Data, deviceTimeout: Int64, scheduler: AsyncScheduler) {
        self.deviceKey = deviceKey
        self.remoteKey = remoteKey
        self.deviceTimeout = deviceTimeout
        self.deviceStateMetric = .init(
            "sesame.device",
            name: "Device connected",
            description: "Shows if the device is connected via WebSocket")
        self.messagesToDeviceMetric = .init(
            "sesame.messages",
            name: "Forwarded Messages",
            description: "The number of messages transmitted to the device")
        self.scheduler = scheduler
    }

    private func updateDeviceConnectionMetric() async {
        _ = try? await deviceStateMetric.update(deviceState)
    }

    private func updateMessageCountMetric() async {
        let lastValue = await messagesToDeviceMetric.lastValue()?.value ?? 0
        _ = try? await messagesToDeviceMetric.update(lastValue + 1)
    }
    
    // MARK: API
    
    private var deviceStatus: String {
        "\(deviceState.rawValue)"
    }
    
    func sendMessageToDevice(_ message: Data, on eventLoop: EventLoop) async throws -> Data {
        guard let socket = connection, !socket.isClosed else {
            connection = nil
            throw MessageResult.deviceNotConnected
        }
        guard requestInProgress == nil else {
            throw MessageResult.operationInProgress
        }
        do {
            try await socket.send(Array(message))
            await updateMessageCountMetric()
        } catch {
            throw MessageResult.deviceNotConnected
        }
        startTimeoutForDeviceRequest(on: eventLoop)
        
        let result: Data = try await withCheckedThrowingContinuation { continuation in
            self.requestInProgress = continuation
        }
        return result
    }
    
    private func startTimeoutForDeviceRequest(on eventLoop: EventLoop) {
        eventLoop.scheduleTask(in: .seconds(deviceTimeout)) { [weak self] in
            self?.resumeDeviceRequest(with: .deviceTimedOut)
        }
    }
    
    private func resumeDeviceRequest(with data: Data) {
        requestInProgress?.resume(returning: data)
        requestInProgress = nil
    }
    
    private func resumeDeviceRequest(with result: MessageResult) {
        requestInProgress?.resume(throwing: result)
        requestInProgress = nil
    }
    
    func authenticateDevice(hash: String) async {
        guard let key = Data(fromHexEncodedString: hash),
              SHA256.hash(data: key) == self.deviceKey else {
            log("Invalid device key")
            await removeDeviceConnection()
            return
        }
        guard let connection, !connection.isClosed else {
            await updateDeviceConnectionMetric()
            return
        }
        deviceIsAuthenticated = true
        await updateDeviceConnectionMetric()
    }

    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")
            self.resumeDeviceRequest(with: .invalidDeviceResponse)
            return
        }
        self.resumeDeviceRequest(with: data)
    }
    
    func didCloseDeviceSocket() {
        deviceIsAuthenticated = false
        connection = nil
    }
    
    func removeDeviceConnection() async {
        try? await connection?.close()
        connection = nil
        deviceIsAuthenticated = false
        await updateDeviceConnectionMetric()
    }
    
    func createNewDeviceConnection(_ socket: WebSocket) async {
        await removeDeviceConnection()

        socket.onBinary { [weak self] _, data in
            self?.processDeviceResponse(data)
        }
        socket.onText { [weak self] _, text async in
            await self?.authenticateDevice(hash: text)
        }

        _ = socket.onClose.always { [weak self] _ in
            self?.didCloseDeviceSocket()
        }

        connection = socket
        await updateDeviceConnectionMetric()
    }
}