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

    private let deviceTimeout: Int64

    private let deviceConnectedMetric: Metric<Bool>

    private let messagesToDeviceMetric: Metric<Int>
    
    let serverStatus: Metric<ServerStatus>
    
    var deviceIsConnected: Bool {
        guard let connection, !connection.isClosed else {
            return false
        }
        return 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, serverStatus: Metric<ServerStatus>) {
        self.deviceKey = deviceKey
        self.remoteKey = remoteKey
        self.deviceTimeout = deviceTimeout
        self.deviceConnectedMetric = .init(
            "sesame.connected",
            name: "Device connection",
            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.serverStatus = serverStatus
    }
    
    func updateMetricsAfterSystemStart() async {
        _ = try? await serverStatus.update(deviceIsConnected ? .nominal : .reducedFunctionality)
        await updateDeviceConnectionMetric()
    }

    private func updateDeviceConnectionMetric() async {
        _ = try? await deviceConnectedMetric.update(deviceIsConnected)
    }

    private func updateMessageCountMetric() async {
        let lastValue = await messagesToDeviceMetric.lastValue()?.value ?? 0
        _ = try? await messagesToDeviceMetric.update(lastValue + 1)
    }
    
    // MARK: API
    
    func sendMessageToDevice(_ message: Data, authToken: Data, on eventLoop: EventLoop) async throws -> Data {
        guard message.count == SignedMessage.size else {
            throw MessageResult.invalidMessageSizeFromRemote
        }
        guard SHA256.hash(data: authToken) == remoteKey else {
            throw MessageResult.invalidServerAuthenticationFromRemote
        }
        guard let socket = connection, !socket.isClosed else {
            connection = nil
            throw MessageResult.deviceNotConnected
        }
        guard requestInProgress == nil else {
            throw MessageResult.tooManyRequests
        }
        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 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: .invalidMessageSizeFromDevice)
            return
        }
        
        guard data.count == SignedMessage.size else {
            print("Invalid size of device message: \(data.count)")
            self.resumeDeviceRequest(with: .invalidMessageSizeFromDevice)
            return
        }
        self.resumeDeviceRequest(with: data)
    }
    
    func didCloseDeviceSocket() {
        connection = nil
    }
    
    func removeDeviceConnection() async {
        try? await connection?.close()
        connection = nil
        await updateDeviceConnectionMetric()
    }
    
    func createNewDeviceConnection(socket: WebSocket, auth: String) async {
        guard let key = Data(fromHexEncodedString: auth),
              SHA256.hash(data: key) == self.deviceKey else {
            log("Invalid device key while opening socket")
            try? await socket.close()
            return
        }
        await removeDeviceConnection()
        
        connection = socket
        socket.eventLoop.execute {
            socket.pingInterval = .seconds(10)
            
            socket.onText { socket, text in
                print("[WARN] Received text over socket: \(text)")
                // Close connection to prevent spamming the log
                try? await socket.close()
            }
            
            socket.onBinary { [weak self] _, data in
                self?.processDeviceResponse(data)
            }
            socket.onClose.whenComplete { [weak self] _ in
                self?.didCloseDeviceSocket()
            }
        }
        await updateDeviceConnectionMetric()
    }
}