Export sender ratchet key (#9)

* feat: exported sender ratchet public key for the message * feat: exported current ratchet public key for the session
LLCFreedom-Space · Mar 27, 2024 · 955f7a8 · 955f7a8
1 parent 8a454df
commit 955f7a8
Show file tree

Hide file tree

Showing 4 changed files with 72 additions and 0 deletions.
diff --git a/Sources/DXProtocol/Messages/MessageContainer.swift b/Sources/DXProtocol/Messages/MessageContainer.swift
@@ -47,6 +47,20 @@ extension MessageContainer: Codable {
         case secureMessage
         case preKeySecureMessage
     }
+
+    /// The public key of the ratchet key pair from sending chain which was used to encrypt this message
+    public var senderRatchetKey: PublicKey {
+        let result: PublicKey
+
+        switch self {
+        case .secureMessage(let message):
+            result = message.senderRatchetKey
+        case .preKeySecureMessage(let message):
+            result = message.secureMessage.senderRatchetKey
+        }
+
+        return result
+    }
 
     /// Initializes a `MessageContainer` struct from a decoder.
     ///

diff --git a/Sources/DXProtocol/Session/Session.swift b/Sources/DXProtocol/Session/Session.swift
@@ -485,6 +485,11 @@ extension Session {
         self.state.archived = true
         self.previousSessionStates.insert(self.state, at: 0)
     }
+
+    /// Returns public key of the current sender ratchet key pair
+    public func currentSenderRatchetKey() -> PublicKey {
+        self.state.senderChain.ratchetKeyPair.publicKey
+    }
 
     /// Makes the specified state the currently active state of the session.
     /// This method must be used to promote newly created states to existing sessions

diff --git a/Tests/DXProtocolTests/MessagesTests/MessageContainerTests.swift b/Tests/DXProtocolTests/MessagesTests/MessageContainerTests.swift
@@ -27,6 +27,48 @@ import XCTest
 @testable import DXProtocol
 
 final class MessageContainerTests: XCTestCase {
+    func testSenderRatchetKey() throws {
+        let aliceRatchetKey = try KeyPair().publicKey
+        let aliceBaseKey = try KeyPair().publicKey
+        let aliceIdentityKey = IdentityKeyPublic(publicKey: try KeyPair().publicKey)
+        let bobIdentityKey = IdentityKeyPublic(publicKey: try KeyPair().publicKey)
+        let macKeyData = Data([
+            0xce, 0x41, 0x35, 0xc2, 0x19, 0xb1, 0xdb, 0xe2, 0xfa, 0x66, 0x6d, 0xdf, 0xb3, 0x9a, 0x7e,
+            0x46, 0x7b, 0xd8, 0x33, 0xf9, 0xcf, 0x30, 0xd9, 0x9d, 0x96, 0x0e, 0xe7, 0x38, 0x07, 0x25,
+            0xa7, 0xf1
+        ])
+
+        let registrationId = UUID()
+        let oneTimePreKeyId = UUID()
+        let signedPreKeyId = UUID()
+
+        let content = "DXProtocolSecureMessage"
+        let data = Data(content.utf8)
+
+        let secureMessage = try SecureMessage(
+                messageVersion: DXProtocolConstants.cipertextMessageCurrentVersion,
+                macKey: macKeyData,
+                senderRatchetKey: aliceRatchetKey,
+                counter: 3,
+                previousCounter: 2,
+                encrypted: data,
+                senderIdentityKey: aliceIdentityKey,
+                receiverIdentityKey: bobIdentityKey)
+        var container = MessageContainer.secureMessage(secureMessage)
+        XCTAssertEqual(container.senderRatchetKey, aliceRatchetKey)
+
+        let preKeyMessage = try PreKeySecureMessage(
+                messageVersion: DXProtocolConstants.cipertextMessageCurrentVersion,
+                registrationId: registrationId,
+                oneTimePreKeyId: oneTimePreKeyId,
+                signedPreKeyId: signedPreKeyId,
+                senderBaseKey: aliceBaseKey,
+                senderIdentityKey: aliceIdentityKey,
+                secureMessage: secureMessage)
+        container = MessageContainer.preKeySecureMessage(preKeyMessage)
+        XCTAssertEqual(container.senderRatchetKey, aliceRatchetKey)
+    }
+
     func testSerializeDeserializeSecureMessage() throws {
         let aliceRatchetKey = try KeyPair().publicKey
         let aliceIdentityKey = IdentityKeyPublic(publicKey: try KeyPair().publicKey)

diff --git a/Tests/DXProtocolTests/SessionTests/SessionTests.swift b/Tests/DXProtocolTests/SessionTests/SessionTests.swift
@@ -531,6 +531,17 @@ final class SessionTests: XCTestCase {
         XCTAssertFalse(session.hasCurrentState())
         XCTAssertEqual(session.previousSessionStates.count, 1)
     }
+
+    func testCurrentSenderRatchetKey() throws {
+        let senderClient = try TestClient(userId: UUID())  // Alice
+        let recipientClient = try TestClient(userId: UUID())  // Bob
+        try initializeSession(senderClient: senderClient, recipientClient: recipientClient)
+
+        let session = try recipientClient.sessionStore.loadSession(for: senderClient.protocolAddress)
+        let expectedRatchetKey = try XCTUnwrap(session?.state.senderChain.ratchetKeyPair.publicKey)
+        XCTAssertEqual(expectedRatchetKey, session?.currentSenderRatchetKey())
+    }
+
     // FIXME: - Need fix
     func testBasicSessionInteraction() throws {
         let senderClient = try TestClient(userId: UUID())  // Alice