diff --git a/Sources/App/API/MessageResult.swift b/Sources/App/API/MessageResult.swift index 1038dfa..ba41746 100644 --- a/Sources/App/API/MessageResult.swift +++ b/Sources/App/API/MessageResult.swift @@ -3,7 +3,7 @@ import Foundation /** A result from sending a key to the device. */ -enum MessageResult: UInt8 { +enum MessageResult: UInt8, Error { /// Text content was received, although binary data was expected case textReceived = 1 @@ -48,6 +48,8 @@ enum MessageResult: UInt8 { case invalidUrlParameter = 20 case invalidResponseAuthentication = 21 + + case invalidDeviceResponse = 22 } extension MessageResult: CustomStringConvertible { @@ -84,6 +86,8 @@ extension MessageResult: CustomStringConvertible { return "The url parameter could not be found" case .invalidResponseAuthentication: return "The response could not be authenticated" + case .invalidDeviceResponse: + return "The device responded with invalid data" } } } diff --git a/Sources/App/DeviceManager.swift b/Sources/App/DeviceManager.swift index 03667ce..4de97bb 100644 --- a/Sources/App/DeviceManager.swift +++ b/Sources/App/DeviceManager.swift @@ -3,6 +3,18 @@ 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 @@ -17,30 +29,38 @@ final class DeviceManager { /// 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 deviceStateMetric: Metric private let messagesToDeviceMetric: Metric 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: EventLoopPromise? + private var requestInProgress: CheckedContinuation? init(deviceKey: Data, remoteKey: Data, deviceTimeout: Int64, scheduler: AsyncScheduler) { self.deviceKey = deviceKey self.remoteKey = remoteKey self.deviceTimeout = deviceTimeout - self.deviceConnectedMetric = .init( - "sesame.connected", + self.deviceStateMetric = .init( + "sesame.device", name: "Device connected", description: "Shows if the device is connected via WebSocket") self.messagesToDeviceMetric = .init( @@ -50,61 +70,72 @@ final class DeviceManager { self.scheduler = scheduler } - private func updateDeviceConnectionMetric() { - scheduler.schedule { [weak self] in - guard let self else { return } - _ = try? await deviceConnectedMetric.update(deviceIsConnected) - } + private func updateDeviceConnectionMetric() async { + _ = try? await deviceStateMetric.update(deviceState) } - private func updateMessageCountMetric() { - scheduler.schedule { [weak self] in - guard let self else { return } - let lastValue = await self.messagesToDeviceMetric.lastValue()?.value ?? 0 - _ = try? await messagesToDeviceMetric.update(lastValue + 1) - } + private func updateMessageCountMetric() async { + let lastValue = await messagesToDeviceMetric.lastValue()?.value ?? 0 + _ = try? await messagesToDeviceMetric.update(lastValue + 1) } // MARK: API private var deviceStatus: String { - deviceIsConnected ? "1" : "0" + "\(deviceState.rawValue)" } - func sendMessageToDevice(_ message: Data, on eventLoop: EventLoop) -> EventLoopFuture { + func sendMessageToDevice(_ message: Data, on eventLoop: EventLoop) async throws -> Data { guard let socket = connection, !socket.isClosed else { connection = nil - return eventLoop.makeSucceededFuture(MessageResult.deviceNotConnected.encoded) + throw MessageResult.deviceNotConnected } guard requestInProgress == nil else { - return eventLoop.makeSucceededFuture(MessageResult.operationInProgress.encoded) + throw MessageResult.operationInProgress } - 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) + do { + try await socket.send(Array(message)) + await updateMessageCountMetric() + } catch { + throw MessageResult.deviceNotConnected } - return result.futureResult + startTimeoutForDeviceRequest(on: eventLoop) + + let result: Data = try await withCheckedThrowingContinuation { continuation in + self.requestInProgress = continuation + } + return result } - func authenticateDevice(hash: String) { - defer { updateDeviceConnectionMetric() } + 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") - _ = connection?.close() - deviceIsAuthenticated = false + await removeDeviceConnection() + return + } + guard let connection, !connection.isClosed else { + await updateDeviceConnectionMetric() return } - log("Device authenticated") deviceIsAuthenticated = true + await updateDeviceConnectionMetric() } func authenticateRemote(_ token: Data) -> Bool { @@ -115,60 +146,39 @@ final class DeviceManager { 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 } - 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) + self.resumeDeviceRequest(with: 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 - } - socket.close().whenSuccess { log("Socket closed") } + func removeDeviceConnection() async { + try? await connection?.close() connection = nil - log("Removed device connection") + deviceIsAuthenticated = false + await updateDeviceConnectionMetric() } - func createNewDeviceConnection(_ socket: WebSocket) { - defer { updateDeviceConnectionMetric() } + func createNewDeviceConnection(_ socket: WebSocket) async { + await removeDeviceConnection() - socket.onBinary { _, data in - self.processDeviceResponse(data) + socket.onBinary { [weak self] _, data in + self?.processDeviceResponse(data) } - socket.onText { _, text in - self.authenticateDevice(hash: text) + socket.onText { [weak self] _, text async in + await self?.authenticateDevice(hash: text) } - _ = socket.onClose.always { _ in - self.didCloseDeviceSocket() + _ = socket.onClose.always { [weak self] _ in + self?.didCloseDeviceSocket() } - isOpeningNewConnection = true - removeDeviceConnection() connection = socket - log("Socket connected") - isOpeningNewConnection = false + await updateDeviceConnectionMetric() } } diff --git a/Sources/App/configure.swift b/Sources/App/configure.swift index a4e8afc..218dddb 100755 --- a/Sources/App/configure.swift +++ b/Sources/App/configure.swift @@ -47,13 +47,6 @@ public func configure(_ app: Application) throws { provider.asyncScheduler = asyncScheduler provider.registerRoutes(app) - // Gracefully shut down by closing potentially open socket - DispatchQueue.global(qos: .utility).asyncAfter(deadline: .now() + .seconds(5)) { - _ = app.server.onShutdown.always { _ in - deviceManager.removeDeviceConnection() - } - } - asyncScheduler.schedule { _ = try await status.update(.nominal) } @@ -62,6 +55,14 @@ public func configure(_ app: Application) throws { df.dateStyle = .short df.timeStyle = .short print("[\(df.string(from: Date()))] Server started") + + // Gracefully shut down by closing potentially open socket + DispatchQueue.global(qos: .utility).asyncAfter(deadline: .now() + .seconds(5)) { + _ = app.server.onShutdown.always { _ in + print("[\(df.string(from: Date()))] Server shutdown") + //await deviceManager.removeDeviceConnection() + } + } } private func loadKeys(at url: URL) throws -> (deviceKey: Data, remoteKey: Data) { diff --git a/Sources/App/routes.swift b/Sources/App/routes.swift index c9e6f96..ea8d4c1 100755 --- a/Sources/App/routes.swift +++ b/Sources/App/routes.swift @@ -11,33 +11,33 @@ extension RouteAPI { } } -private func messageTransmission(_ req: Request) -> EventLoopFuture { +private func messageTransmission(_ req: Request) async throws -> Data { guard let body = req.body.data else { - return req.eventLoop.makeSucceededFuture(MessageResult.noBodyData.encoded) + throw MessageResult.noBodyData } guard let message = ServerMessage(decodeFrom: body) else { - return req.eventLoop.makeSucceededFuture(MessageResult.invalidMessageSize.encoded) + throw MessageResult.invalidMessageSize } guard deviceManager.authenticateRemote(message.authToken) else { - return req.eventLoop.makeSucceededFuture(MessageResult.messageAuthenticationFailed.encoded) + throw MessageResult.messageAuthenticationFailed } - return deviceManager.sendMessageToDevice(message.message, on: req.eventLoop) + return try await deviceManager.sendMessageToDevice(message.message, on: req.eventLoop) } -private func deviceStatus(_ req: Request) -> EventLoopFuture { +private func deviceStatus(_ req: Request) -> MessageResult { guard let body = req.body.data else { - return req.eventLoop.makeSucceededFuture(.noBodyData) + return .noBodyData } guard let authToken = ServerMessage.token(from: body) else { - return req.eventLoop.makeSucceededFuture(.invalidMessageSize) + return .invalidMessageSize } guard deviceManager.authenticateRemote(authToken) else { - return req.eventLoop.makeSucceededFuture(.messageAuthenticationFailed) + return .messageAuthenticationFailed } guard deviceManager.deviceIsConnected else { - return req.eventLoop.makeSucceededFuture(.deviceNotConnected) + return .deviceNotConnected } - return req.eventLoop.makeSucceededFuture(.deviceConnected) + return .deviceConnected } func routes(_ app: Application) throws { @@ -50,10 +50,9 @@ func routes(_ app: Application) throws { The request returns one byte of data, which is the raw value of a `MessageResult`. Possible results are `noBodyData`, `invalidMessageSize`, `deviceNotConnected`, `deviceConnected`. */ - app.post(RouteAPI.getDeviceStatus.path) { req in - deviceStatus(req).map { - Response(status: .ok, body: .init(data: $0.encoded)) - } + app.post(RouteAPI.getDeviceStatus.path) { request in + let result = deviceStatus(request) + return Response(status: .ok, body: .init(data: result.encoded)) } /** @@ -64,9 +63,12 @@ func routes(_ app: Application) throws { The request returns one or `Message.length+1` bytes of data, where the first byte is the raw value of a `MessageResult`, and the optional following bytes contain the response message of the device. This request does not complete until either the device responds or the request times out. The timeout is specified by `KeyManagement.deviceTimeout`. */ - app.post(RouteAPI.postMessage.path) { req in - messageTransmission(req).map { - Response(status: .ok, body: .init(data: $0)) + app.post(RouteAPI.postMessage.path) { request async throws in + do { + let result = try await messageTransmission(request) + return Response(status: .ok, body: .init(data: result)) + } catch let error as MessageResult { + return Response(status: .ok, body: .init(data: error.encoded)) } } @@ -75,7 +77,7 @@ func routes(_ app: Application) throws { - Returns: Nothing - Note: The first message from the device over the connection must be a valid auth token. */ - app.webSocket(RouteAPI.socket.path) { req, socket in - deviceManager.createNewDeviceConnection(socket) + app.webSocket(RouteAPI.socket.path) { req, socket async in + await deviceManager.createNewDeviceConnection(socket) } }