From d6b171229f5b986bcb15cb379deb284cf0efe519 Mon Sep 17 00:00:00 2001
From: sergeych <real.sergeych@gmail.com>
Date: Fri, 29 May 2026 15:43:44 +0300
Subject: [PATCH] Added support for SafeKeyExchange
 serialization/deserialization and associated tests.

---
 .../net/sergeych/crypto2/SafeKeyExchange.kt   | 65 +++++++++++++++++--
 src/commonTest/kotlin/KeysTest.kt             | 20 ++++++
 2 files changed, 81 insertions(+), 4 deletions(-)

diff --git a/src/commonMain/kotlin/net/sergeych/crypto2/SafeKeyExchange.kt b/src/commonMain/kotlin/net/sergeych/crypto2/SafeKeyExchange.kt
index 5d18f88..4461aba 100644
--- a/src/commonMain/kotlin/net/sergeych/crypto2/SafeKeyExchange.kt
+++ b/src/commonMain/kotlin/net/sergeych/crypto2/SafeKeyExchange.kt
@@ -11,7 +11,14 @@
 package net.sergeych.crypto2
 
 import com.ionspin.kotlin.crypto.keyexchange.KeyExchange
+import com.ionspin.kotlin.crypto.keyexchange.KeyExchangeKeyPair
+import com.ionspin.kotlin.crypto.keyexchange.crypto_kx_PUBLICKEYBYTES
+import com.ionspin.kotlin.crypto.keyexchange.crypto_kx_SECRETKEYBYTES
+import kotlinx.serialization.KSerializer
 import kotlinx.serialization.Serializable
+import kotlinx.serialization.descriptors.SerialDescriptor
+import kotlinx.serialization.encoding.Decoder
+import kotlinx.serialization.encoding.Encoder
 import net.sergeych.crypto2.SafeKeyExchange.SessionKey
 
 /**
@@ -31,9 +38,44 @@ import net.sergeych.crypto2.SafeKeyExchange.SessionKey
  * - while it is possible to generate several keys "ahead", the care should be taken when storing them,
  *   encrypt it with some other key to maintain safety.
  * - do not use [EncryptingPublicKey] for anything but creating session keys.
+ * - the serialized form includes the secret exchange key, so it should be encrypted when stored.
  */
-class SafeKeyExchange {
-    private val pair = KeyExchange.keypair()
+@Serializable(with = SafeKeyExchange.SafeKeyExchangeSerializer::class)
+class SafeKeyExchange private constructor(
+    private val pair: KeyExchangeKeyPair,
+) {
+
+    constructor() : this(KeyExchange.keypair())
+
+    @Serializable
+    private class Packed(
+        val publicKey: UByteArray,
+        val secretKey: UByteArray,
+    )
+
+    object SafeKeyExchangeSerializer : KSerializer<SafeKeyExchange> {
+        private val packedSerializer = Packed.serializer()
+
+        override val descriptor: SerialDescriptor = packedSerializer.descriptor
+
+        override fun serialize(encoder: Encoder, value: SafeKeyExchange) {
+            encoder.encodeSerializableValue(
+                packedSerializer,
+                Packed(value.pair.publicKey, value.pair.secretKey)
+            )
+        }
+
+        override fun deserialize(decoder: Decoder): SafeKeyExchange {
+            val packed = decoder.decodeSerializableValue(packedSerializer)
+            require(packed.publicKey.size == crypto_kx_PUBLICKEYBYTES) {
+                "SafeKeyExchange public key must be $crypto_kx_PUBLICKEYBYTES bytes"
+            }
+            require(packed.secretKey.size == crypto_kx_SECRETKEYBYTES) {
+                "SafeKeyExchange secret key must be $crypto_kx_SECRETKEYBYTES bytes"
+            }
+            return SafeKeyExchange(KeyExchangeKeyPair(packed.publicKey, packed.secretKey))
+        }
+    }
 
     /**
      * The session key. It uses a pair of keys to encrypt and decrypt messages to maintain high
@@ -90,9 +132,24 @@ class SafeKeyExchange {
      * The public key; it should be transmitted to the other party, this is serializable.
      * Do not use it except to get [SessionKey] with [clientSessionKey] or [serverSessionKey]. Storing and reusing
      * it is a great danger.
+     *
+     * Instances can be compared and used as hashtable keys.
      */
     @Serializable
-    class PublicKey(val keyBytes: UByteArray)
+    class PublicKey(val keyBytes: UByteArray) {
+        override fun equals(other: Any?): Boolean {
+            if (this === other) return true
+            if (other == null || this::class != other::class) return false
+
+            other as PublicKey
+
+            return keyBytes.contentEquals(other.keyBytes)
+        }
+
+        override fun hashCode(): Int {
+            return keyBytes.contentHashCode()
+        }
+    }
 
     /**
      * The public key to be sent to the other party. When received, get the session keys with [clientSessionKey]
@@ -120,4 +177,4 @@ class SafeKeyExchange {
             .let { SessionKey(SymmetricKey(it.sendKey), SymmetricKey(it.receiveKey), isClient = false) }
 
 
-}
\ No newline at end of file
+}
diff --git a/src/commonTest/kotlin/KeysTest.kt b/src/commonTest/kotlin/KeysTest.kt
index 8c9853b..20b274e 100644
--- a/src/commonTest/kotlin/KeysTest.kt
+++ b/src/commonTest/kotlin/KeysTest.kt
@@ -221,6 +221,26 @@ class KeysTest {
 
     }
 
+    @Test
+    fun keyExchangeSerializationTest() = runTest {
+        initCrypto()
+        val ske = SafeKeyExchange()
+        val packedExchange = pack(ske)
+
+        val cke = SafeKeyExchange()
+        val clientSessionKey = cke.clientSessionKey(ske.publicKey)
+
+        val restoredExchange = unpack<SafeKeyExchange>(packedExchange)
+        assertEquals(ske.publicKey, restoredExchange.publicKey)
+
+        val serverSessionKey = restoredExchange.serverSessionKey(cke.publicKey)
+
+        val src = "Hello after restore!"
+        assertEquals(src, serverSessionKey.decryptString(clientSessionKey.encrypt(src)))
+        assertEquals(src, clientSessionKey.decryptString(serverSessionKey.encrypt(src)))
+        assertContentEquals(clientSessionKey.sessionTag, serverSessionKey.sessionTag)
+    }
+
     @Test
     fun asymmetricKeyTest() = runTest {
         initCrypto()