Update device authentication

Sources/App/DeviceManager.swift

@@ -3,18 +3,6 @@ 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
@@ -26,28 +14,17 @@ 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 deviceStateMetric: Metric<DeviceState>
+    private let deviceConnectedMetric: Metric<Bool>
 
     private let messagesToDeviceMetric: Metric<Int>
     
-    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)
+        guard let connection, !connection.isClosed else {
+            return false
+        }
+        return true
     }
     
     /// A promise to finish the request once the device responds or times out
@@ -57,10 +34,10 @@ final class DeviceManager {
         self.deviceKey = deviceKey
         self.remoteKey = remoteKey
         self.deviceTimeout = deviceTimeout
-        self.deviceStateMetric = .init(
-            "sesame.device",
-            name: "Device status",
-            description: "Shows if the device is connected and authenticated via WebSocket")
+        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",
@@ -68,7 +45,7 @@ final class DeviceManager {
     }
 
     func updateDeviceConnectionMetric() async {
-        _ = try? await deviceStateMetric.update(deviceState)
+        _ = try? await deviceConnectedMetric.update(deviceIsConnected)
     }
 
     private func updateMessageCountMetric() async {
@@ -78,10 +55,6 @@ 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
@@ -120,21 +93,6 @@ final class DeviceManager {
         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
@@ -150,32 +108,34 @@ 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) async {
+    func createNewDeviceConnection(socket: WebSocket, auth: String) 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.onText { [weak self] _, text async in
-                await self?.authenticateDevice(hash: text)
-            }
-
-            _ = socket.onClose.always { [weak self] _ in
+            socket.onClose.whenComplete { [weak self] _ in
                 self?.didCloseDeviceSocket()
             }
         }
-        connection = socket
         await updateDeviceConnectionMetric()
     }
 }

Sources/App/routes.swift

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