Add token authentication to KASWebSocket connection

The KASWebSocket class has been updated to include a token for authentication during connection. The token is passed in upon initialization and is added to the header of the WebSocket request. The token is also sent as a WebSocket message on connection. The changes also affect KASWebSocket tests to accommodate the new parameter.

OpenTDFKit/KASWebSocket.swift

@@ -10,25 +10,25 @@ public enum WebSocketConnectionState {
 
 public class KASWebSocket {
     private var webSocketTask: URLSessionWebSocketTask?
-    private let urlSession: URLSession
+    private var urlSession: URLSession?
     private let myPrivateKey: P256.KeyAgreement.PrivateKey!
     private var sharedSecret: SharedSecret?
     private var salt: Data?
     private var rewrapCallback: ((Data, SymmetricKey?) -> Void)?
     private var kasPublicKeyCallback: ((P256.KeyAgreement.PublicKey) -> Void)?
     private let kasUrl: URL
+    private let token: String
 
     private let connectionStateSubject = CurrentValueSubject<WebSocketConnectionState, Never>(.disconnected)
     public var connectionStatePublisher: AnyPublisher<WebSocketConnectionState, Never> {
         connectionStateSubject.eraseToAnyPublisher()
     }
 
-    public init(kasUrl: URL) {
+    public init(kasUrl: URL, token: String) {
         // create key
         myPrivateKey = P256.KeyAgreement.PrivateKey()
-        // Initialize a URLSession with a default configuration
-        urlSession = URLSession(configuration: .default)
         self.kasUrl = kasUrl
+        self.token = token
     }
 
     public func setRewrapCallback(_ callback: @escaping (Data, SymmetricKey?) -> Void) {
@@ -41,8 +41,20 @@ public class KASWebSocket {
 
     public func connect() {
         connectionStateSubject.send(.connecting)
-        webSocketTask = urlSession.webSocketTask(with: kasUrl)
+        // Create a URLRequest object with the WebSocket URL
+        var request = URLRequest(url: kasUrl)
+        // Add the Authorization header to the request
+        request.addValue("Bearer \(token)", forHTTPHeaderField: "Authorization")
+        // Initialize a URLSession with a default configuration
+        urlSession = URLSession(configuration: .default)
+        webSocketTask = urlSession!.webSocketTask(with: request)
         webSocketTask?.resume()
+        let tokenMessage = URLSessionWebSocketTask.Message.string(token)
+        webSocketTask?.send(tokenMessage) { error in
+            if let error {
+                print("token sending error: \(error)")
+            }
+        }
         // Start receiving messages
         receiveMessage()
         pingPeriodically()

OpenTDFKitTests/KASWebsocketTests.swift

@@ -6,8 +6,8 @@ final class KASWebsocketTests: XCTestCase {
     func testEncryptDecrypt() throws {
         measure(metrics: [XCTCPUMetric()]) {
             let nanoTDFManager = NanoTDFManager()
-            let webSocket = KASWebSocket(kasUrl: URL(string: "wss://kas.arkavo.net")!)
-//            let webSocket = KASWebSocket(kasUrl: URL(string: "ws://localhost:8080")!)
+            let webSocket = KASWebSocket(kasUrl: URL(string: "wss://kas.arkavo.net")!, token: "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c")
+//            let webSocket = KASWebSocket(kasUrl: URL(string: "ws://localhost:8080")!, token: "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c")
             let plaintext = "Keep this message secret".data(using: .utf8)!
             webSocket.setRewrapCallback { identifier, symmetricKey in
 //                defer {
@@ -86,7 +86,7 @@ final class KASWebsocketTests: XCTestCase {
     }
 
     func testWebsocket() throws {
-        let webSocket = KASWebSocket(kasUrl: URL(string: "ws://localhost:8080")!)
+        let webSocket = KASWebSocket(kasUrl: URL(string: "ws://localhost:8080")!, token: "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c")
         let expectation = XCTestExpectation(description: "Receive rewrapped key")
         // Create a 33-byte identifier
         let testIdentifier = Data((0 ..< 33).map { _ in UInt8.random(in: 0 ... 255) })