From f171ab2d4d16f28d47e0eded9d60d78756ed778d Mon Sep 17 00:00:00 2001 From: Paul Flynn Date: Mon, 8 Jul 2024 21:14:32 -0400 Subject: [PATCH] Refactor KASWebSocket initialization and enhance error handling 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. --- ...3169F6C6-AD33-4C10-8516-0CD808F6EB48.plist | 43 +++++++++++++++++++ .../OpenTDFKitTests.xcbaseline/Info.plist | 33 ++++++++++++++ OpenTDFKit/KASWebSocket.swift | 14 ++++-- OpenTDFKitTests/KASWebsocketTests.swift | 13 ++++-- 4 files changed, 95 insertions(+), 8 deletions(-) create mode 100644 .swiftpm/xcode/xcshareddata/xcbaselines/OpenTDFKitTests.xcbaseline/3169F6C6-AD33-4C10-8516-0CD808F6EB48.plist create mode 100644 .swiftpm/xcode/xcshareddata/xcbaselines/OpenTDFKitTests.xcbaseline/Info.plist diff --git a/.swiftpm/xcode/xcshareddata/xcbaselines/OpenTDFKitTests.xcbaseline/3169F6C6-AD33-4C10-8516-0CD808F6EB48.plist b/.swiftpm/xcode/xcshareddata/xcbaselines/OpenTDFKitTests.xcbaseline/3169F6C6-AD33-4C10-8516-0CD808F6EB48.plist new file mode 100644 index 0000000..db6a0a6 --- /dev/null +++ b/.swiftpm/xcode/xcshareddata/xcbaselines/OpenTDFKitTests.xcbaseline/3169F6C6-AD33-4C10-8516-0CD808F6EB48.plist @@ -0,0 +1,43 @@ + + + + + classNames + + KASWebsocketTests + + testEncryptDecrypt() + + com.apple.dt.XCTMetric_CPU.cycles + + baselineAverage + 1650206.736600 + baselineIntegrationDisplayName + Local Baseline + + com.apple.dt.XCTMetric_CPU.instructions_retired + + baselineAverage + 17757462.366800 + baselineIntegrationDisplayName + Local Baseline + + com.apple.dt.XCTMetric_CPU.time + + baselineAverage + 1.325406 + baselineIntegrationDisplayName + Local Baseline + + com.apple.dt.XCTMetric_Memory.physical + + baselineAverage + 3257.190400 + baselineIntegrationDisplayName + Local Baseline + + + + + + diff --git a/.swiftpm/xcode/xcshareddata/xcbaselines/OpenTDFKitTests.xcbaseline/Info.plist b/.swiftpm/xcode/xcshareddata/xcbaselines/OpenTDFKitTests.xcbaseline/Info.plist new file mode 100644 index 0000000..aa8d1c5 --- /dev/null +++ b/.swiftpm/xcode/xcshareddata/xcbaselines/OpenTDFKitTests.xcbaseline/Info.plist @@ -0,0 +1,33 @@ + + + + + runDestinationsByUUID + + 3169F6C6-AD33-4C10-8516-0CD808F6EB48 + + localComputer + + busSpeedInMHz + 0 + cpuCount + 1 + cpuKind + Apple M1 Max + cpuSpeedInMHz + 0 + logicalCPUCoresPerPackage + 10 + modelCode + Mac13,1 + physicalCPUCoresPerPackage + 10 + platformIdentifier + com.apple.platform.macosx + + targetArchitecture + arm64e + + + + diff --git a/OpenTDFKit/KASWebSocket.swift b/OpenTDFKit/KASWebSocket.swift index 4ea7656..3078497 100644 --- a/OpenTDFKit/KASWebSocket.swift +++ b/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) diff --git a/OpenTDFKitTests/KASWebsocketTests.swift b/OpenTDFKitTests/KASWebsocketTests.swift index 5d0aca5..ec80d0c 100644 --- a/OpenTDFKitTests/KASWebsocketTests.swift +++ b/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() }