Refactor KASWebSocket initialization and enhance error handling (#6)

This commit refactors the KASWebSocket's initializer to accept a URL parameter, enabling flexibility in the WebSocket endpoint. It addresses previous hard-coded "kasUrl" in the connect() method. Also, it introduces better error handling to manage scenarios where the received data isn't equal to the expected 93 bytes. Tests have been updated to comply with these changes. Additionally, new performance baselines for KASWebsocketTests have been established.

.swiftpm/xcode/xcshareddata/xcbaselines/OpenTDFKitTests.xcbaseline/3169F6C6-AD33-4C10-8516-0CD808F6EB48.plist

@@ -0,0 +1,43 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>classNames</key>
+	<dict>
+		<key>KASWebsocketTests</key>
+		<dict>
+			<key>testEncryptDecrypt()</key>
+			<dict>
+				<key>com.apple.dt.XCTMetric_CPU.cycles</key>
+				<dict>
+					<key>baselineAverage</key>
+					<real>1650206.736600</real>
+					<key>baselineIntegrationDisplayName</key>
+					<string>Local Baseline</string>
+				</dict>
+				<key>com.apple.dt.XCTMetric_CPU.instructions_retired</key>
+				<dict>
+					<key>baselineAverage</key>
+					<real>17757462.366800</real>
+					<key>baselineIntegrationDisplayName</key>
+					<string>Local Baseline</string>
+				</dict>
+				<key>com.apple.dt.XCTMetric_CPU.time</key>
+				<dict>
+					<key>baselineAverage</key>
+					<real>1.325406</real>
+					<key>baselineIntegrationDisplayName</key>
+					<string>Local Baseline</string>
+				</dict>
+				<key>com.apple.dt.XCTMetric_Memory.physical</key>
+				<dict>
+					<key>baselineAverage</key>
+					<real>3257.190400</real>
+					<key>baselineIntegrationDisplayName</key>
+					<string>Local Baseline</string>
+				</dict>
+			</dict>
+		</dict>
+	</dict>
+</dict>
+</plist>

.swiftpm/xcode/xcshareddata/xcbaselines/OpenTDFKitTests.xcbaseline/Info.plist

@@ -0,0 +1,33 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>runDestinationsByUUID</key>
+	<dict>
+		<key>3169F6C6-AD33-4C10-8516-0CD808F6EB48</key>
+		<dict>
+			<key>localComputer</key>
+			<dict>
+				<key>busSpeedInMHz</key>
+				<integer>0</integer>
+				<key>cpuCount</key>
+				<integer>1</integer>
+				<key>cpuKind</key>
+				<string>Apple M1 Max</string>
+				<key>cpuSpeedInMHz</key>
+				<integer>0</integer>
+				<key>logicalCPUCoresPerPackage</key>
+				<integer>10</integer>
+				<key>modelCode</key>
+				<string>Mac13,1</string>
+				<key>physicalCPUCoresPerPackage</key>
+				<integer>10</integer>
+				<key>platformIdentifier</key>
+				<string>com.apple.platform.macosx</string>
+			</dict>
+			<key>targetArchitecture</key>
+			<string>arm64e</string>
+		</dict>
+	</dict>
+</dict>
+</plist>

OpenTDFKit/KASWebSocket.swift

@@ -53,12 +53,14 @@ public class KASWebSocket {
     private var salt: Data?
     private var rewrapCallback: ((Data, SymmetricKey?) -> Void)?
     private var kasPublicKeyCallback: ((P256.KeyAgreement.PublicKey) -> Void)?
+    private let kasUrl: URL
 
-    public init() {
+    public init(kasUrl: URL) {
         // create key
         myPrivateKey = P256.KeyAgreement.PrivateKey()
         // Initialize a URLSession with a default configuration
         urlSession = URLSession(configuration: .default)
+        self.kasUrl = kasUrl
     }
 
     public func setRewrapCallback(_ callback: @escaping (Data, SymmetricKey?) -> Void) {
@@ -71,8 +73,7 @@ public class KASWebSocket {
 
     public func connect() {
         // Create the WebSocket task with the specified URL
-        let url = URL(string: "wss://kas.arkavo.net")!
-        webSocketTask = urlSession.webSocketTask(with: url)
+        webSocketTask = urlSession.webSocketTask(with: kasUrl)
         webSocketTask?.resume()
         // Start receiving messages
         receiveMessage()
@@ -170,7 +171,12 @@ public class KASWebSocket {
 //        print("BEGIN handleRewrappedKeyMessage")
         // print("wrapped_dek_shared_secret \(data.hexEncodedString())")
         guard data.count == 93 else {
-            print("Received data is not the expected 93 bytes (33 for identifier + 60 for key)")
+            if data.count == 33 {
+                // DENY -- Notify the app with the identifier
+                rewrapCallback?(data, nil)
+                return
+            }
+            print("RewrappedKeyMessage not the expected 93 bytes (33 for identifier + 60 for key): \(data.count)")
             return
         }
         let identifier = data.prefix(33)

OpenTDFKitTests/KASWebsocketTests.swift

@@ -6,7 +6,8 @@ final class KASWebsocketTests: XCTestCase {
     func testEncryptDecrypt() throws {
         measure(metrics: [XCTCPUMetric()]) {
             let nanoTDFManager = NanoTDFManager()
-            let webSocket = KASWebSocket()
+            let webSocket = KASWebSocket(kasUrl: URL(string: "wss://kas.arkavo.net")!)
+//            let webSocket = KASWebSocket(kasUrl: URL(string: "ws://localhost:8080")!)
             let plaintext = "Keep this message secret".data(using: .utf8)!
             webSocket.setRewrapCallback { identifier, symmetricKey in
 //                defer {
@@ -17,6 +18,10 @@ final class KASWebsocketTests: XCTestCase {
 //                print("Received Rewrapped Symmetric key: \(String(describing: symmetricKey))")
                 let nanoTDF = nanoTDFManager.getNanoTDF(withIdentifier: identifier)
                 nanoTDFManager.removeNanoTDF(withIdentifier: identifier)
+                if symmetricKey == nil {
+                    // DENY
+                    return
+                }
                 let payload = nanoTDF?.payload
                 let rawIV = payload?.iv
                 // Pad the IV
@@ -45,7 +50,7 @@ final class KASWebsocketTests: XCTestCase {
             webSocket.setKASPublicKeyCallback { publicKey in
                 let kasRL = ResourceLocator(protocolEnum: .http, body: "localhost:8080")
                 let kasMetadata = KasMetadata(resourceLocator: kasRL!, publicKey: publicKey, curve: .secp256r1)
-                let remotePolicy = ResourceLocator(protocolEnum: .sharedResourceDirectory, body: "localhost/123")
+                let remotePolicy = ResourceLocator(protocolEnum: .sharedResourceDirectory, body: "5Cqk3ERPToSMuY8UoKJtcmo4fs1iVyQpq6ndzWzpzWezAF1W")
                 var policy = Policy(type: .remote, body: nil, remote: remotePolicy, binding: nil)
 
                 do {
@@ -81,7 +86,7 @@ final class KASWebsocketTests: XCTestCase {
     }
 
     func testWebsocket() throws {
-        let webSocket = KASWebSocket()
+        let webSocket = KASWebSocket(kasUrl: URL(string: "ws://localhost:8080")!)
         let expectation = XCTestExpectation(description: "Receive rewrapped key")
         // Create a 33-byte identifier
         let testIdentifier = Data((0 ..< 33).map { _ in UInt8.random(in: 0 ... 255) })
@@ -111,7 +116,7 @@ final class KASWebsocketTests: XCTestCase {
         // Send a request for KAS key for encrypt
 //        webSocket.sendKASKeyMessage()
         // wait
-        Thread.sleep(forTimeInterval: 2.0)
+        Thread.sleep(forTimeInterval: 1.0)
         // Optionally, disconnect when done or needed
         webSocket.disconnect()
     }