Refactor KASWebSocket.swift for managing connection state

The code changes involve refactoring the KASWebSocket.swift class to manage the WebSocket connection state with Combine. This reorganization includes the creation of a connectionStateSubject and the introduction of connection and disconnection states, as well as a publisher for these states. The code now also performs periodic "pinging" to check the connection status.

OpenTDFKit/KASWebSocket.swift

@@ -1,48 +1,11 @@
 import CryptoKit
 import Foundation
+import Combine
 
-struct KASKeyMessage {
-    let messageType: Data = .init([0x02])
-
-    func toData() -> Data {
-        messageType
-    }
-}
-
-struct PublicKeyMessage {
-    let messageType: Data = .init([0x01])
-    let publicKey: Data
-
-    func toData() -> Data {
-        var data = Data()
-        data.append(messageType)
-        data.append(publicKey)
-        return data
-    }
-}
-
-struct RewrapMessage {
-    let messageType: Data = .init([0x03])
-    let header: Header
-
-    func toData() -> Data {
-        var data = Data()
-        data.append(messageType)
-        data.append(header.toData())
-        return data
-    }
-}
-
-struct RewrappedKeyMessage {
-    let messageType: Data = .init([0x04])
-    let rewrappedKey: Data
-
-    func toData() -> Data {
-        var data = Data()
-        data.append(messageType)
-        data.append(rewrappedKey)
-        return data
-    }
+public enum WebSocketConnectionState {
+    case disconnected
+    case connecting
+    case connected
 }
 
 public class KASWebSocket {
@@ -54,6 +17,11 @@ public class KASWebSocket {
     private var rewrapCallback: ((Data, SymmetricKey?) -> Void)?
     private var kasPublicKeyCallback: ((P256.KeyAgreement.PublicKey) -> Void)?
     private let kasUrl: URL
+
+    private let connectionStateSubject = CurrentValueSubject<WebSocketConnectionState, Never>(.disconnected)
+    public var connectionStatePublisher: AnyPublisher<WebSocketConnectionState, Never> {
+        connectionStateSubject.eraseToAnyPublisher()
+    }
 
     public init(kasUrl: URL) {
         // create key
@@ -72,19 +40,37 @@ public class KASWebSocket {
     }
 
     public func connect() {
-        // Create the WebSocket task with the specified URL
+        connectionStateSubject.send(.connecting)
         webSocketTask = urlSession.webSocketTask(with: kasUrl)
         webSocketTask?.resume()
         // Start receiving messages
         receiveMessage()
+        pingPeriodically()
+    }
+
+    private func pingPeriodically() {
+        webSocketTask?.sendPing { [weak self] error in
+            if let error = error {
+                print("Error sending ping: \(error)")
+                self?.connectionStateSubject.send(.disconnected)
+            } else {
+                self?.connectionStateSubject.send(.connected)
+            }
+            // Schedule next ping
+            DispatchQueue.main.asyncAfter(deadline: .now() + 5) { [weak self] in
+                self?.pingPeriodically()
+            }
+        }
     }
 
     private func receiveMessage() {
         webSocketTask?.receive { [weak self] result in
             switch result {
             case let .failure(error):
                 print("Failed to receive message: \(error)")
+                self?.connectionStateSubject.send(.disconnected)
             case let .success(message):
+                self?.connectionStateSubject.send(.connected)
                 switch message {
                 case let .string(text):
                     print("Received string: \(text)")
@@ -93,7 +79,6 @@ public class KASWebSocket {
                 @unknown default:
                     fatalError()
                 }
-
                 // Continue receiving messages
                 self?.receiveMessage()
             }
@@ -254,9 +239,19 @@ public class KASWebSocket {
         }
     }
 
+    public func sendPing(completionHandler: @escaping (Error?) -> Void) {
+        webSocketTask?.sendPing { error in
+            if let error = error {
+                print("Error sending ping: \(error)")
+            }
+            completionHandler(error)
+        }
+    }
+
+
     public func disconnect() {
-        // Close the WebSocket connection
         webSocketTask?.cancel(with: .goingAway, reason: nil)
+        connectionStateSubject.send(.disconnected)
     }
 }
 
@@ -266,3 +261,47 @@ extension Data {
         map { String(format: "%02hhx", $0) }.joined()
     }
 }
+
+struct KASKeyMessage {
+    let messageType: Data = .init([0x02])
+
+    func toData() -> Data {
+        messageType
+    }
+}
+
+struct PublicKeyMessage {
+    let messageType: Data = .init([0x01])
+    let publicKey: Data
+
+    func toData() -> Data {
+        var data = Data()
+        data.append(messageType)
+        data.append(publicKey)
+        return data
+    }
+}
+
+struct RewrapMessage {
+    let messageType: Data = .init([0x03])
+    let header: Header
+
+    func toData() -> Data {
+        var data = Data()
+        data.append(messageType)
+        data.append(header.toData())
+        return data
+    }
+}
+
+struct RewrappedKeyMessage {
+    let messageType: Data = .init([0x04])
+    let rewrappedKey: Data
+
+    func toData() -> Data {
+        var data = Data()
+        data.append(messageType)
+        data.append(rewrappedKey)
+        return data
+    }
+}