.gitignore

@@ -3,5 +3,3 @@ Package.resolved
 .swiftpm
 .build
 Resources/config.json
-Resources/logs
-Resources/keys

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
@@ -14,17 +26,28 @@ final class DeviceManager {
     /// 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 deviceConnectedMetric: Metric<Bool>
+    private let deviceStateMetric: Metric<DeviceState>
 
     private let messagesToDeviceMetric: Metric<Int>
     
-    var deviceIsConnected: Bool {
+    var deviceState: DeviceState {
         guard let connection, !connection.isClosed else {
-            return false
+            return .disconnected
         }
-        return true
+        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
@@ -34,10 +57,10 @@ final class DeviceManager {
         self.deviceKey = deviceKey
         self.remoteKey = remoteKey
         self.deviceTimeout = deviceTimeout
-        self.deviceConnectedMetric = .init(
+        self.deviceStateMetric = .init(
-            "sesame.connected",
+            "sesame.device",
-            name: "Device connection",
+            name: "Device status",
-            description: "Shows if the device is connected via WebSocket")
+            description: "Shows if the device is connected and authenticated via WebSocket")
         self.messagesToDeviceMetric = .init(
             "sesame.messages",
             name: "Forwarded Messages",
@@ -45,7 +68,7 @@ final class DeviceManager {
     }
 
     func updateDeviceConnectionMetric() async {
-        _ = try? await deviceConnectedMetric.update(deviceIsConnected)
+        _ = try? await deviceStateMetric.update(deviceState)
     }
 
     private func updateMessageCountMetric() async {
@@ -55,6 +78,10 @@ final class DeviceManager {
     
     // 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
@@ -92,6 +119,21 @@ final class DeviceManager {
         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)
@@ -108,34 +150,32 @@ final class DeviceManager {
     }
     
     func didCloseDeviceSocket() {
+        deviceIsAuthenticated = false
         connection = nil
     }
     
     func removeDeviceConnection() async {
         try? await connection?.close()
         connection = nil
+        deviceIsAuthenticated = false
         await updateDeviceConnectionMetric()
     }
     
-    func createNewDeviceConnection(socket: WebSocket, auth: String) async {
+    func createNewDeviceConnection(_ socket: WebSocket) async {
-        guard let key = Data(fromHexEncodedString: auth),
-              SHA256.hash(data: key) == self.deviceKey else {
-            log("Invalid device key")
-            return
-        }
         await removeDeviceConnection()
-        
-        connection = socket
         socket.eventLoop.execute {
-            socket.pingInterval = .seconds(10)
-            
             socket.onBinary { [weak self] _, data in
                 self?.processDeviceResponse(data)
             }
-            socket.onClose.whenComplete { [weak self] _ in
+            socket.onText { [weak self] _, text async in
+                await self?.authenticateDevice(hash: text)
+            }
+
+            _ = socket.onClose.always { [weak self] _ in
                 self?.didCloseDeviceSocket()
             }
         }
+        connection = socket
         await updateDeviceConnectionMetric()
     }
 }

Sources/App/routes.swift

@@ -77,11 +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) { request, socket async in
+    app.webSocket(RouteAPI.socket.path) { req, socket async in
-        guard let authToken = request.headers.first(name: "Authorization") else {
+        await deviceManager.createNewDeviceConnection(socket)
-            try? await socket.close()
-            return
-        }
-        await deviceManager.createNewDeviceConnection(socket: socket, auth: authToken)
     }
 }

Sources/Run/main.swift

@@ -5,8 +5,8 @@ var env = Environment.production //.detect()
 try LoggingSystem.bootstrap(from: &env)
 let app = Application(env)
 defer {
-    shutdown()
     app.shutdown()
+    shutdown()
 }
 
 try configure(app)