working container with adding keys

src/commonMain/kotlin/net/sergeych/crypto2/Container.kt

@@ -5,57 +5,129 @@ import kotlinx.serialization.Serializable
 import kotlinx.serialization.Transient
 import net.sergeych.bipack.BipackDecoder
 import net.sergeych.bipack.BipackEncoder
+import net.sergeych.crypto2.Container.Companion.createWith
 
-/*
-The problem with container is the following. When we encrypt it with asymmetric key,
-we provide public key as the recipient. But public key alone is not effective
-as it is restricted to anonymous encryption which is very slow.
-
-We might need an alternative to specify the (sender key,recipient key) pair also. We need a good
-solution for it.
 
+/**
+ * Multi-key encrypted container with simple adding new keys function that does not need to
+ * know all existing keys, e.g., you can add recipients to the container data if you can
+ * decrypt it. This is sometimes very important to be able to add recipients to the message
+ * keeping existing recipients you know no keys of.
+ *
+ * See:
+ * - [createWith] for more on create a new container
+ * - [addRecipients] and various [plus] operators to add recipients
+ * - [decryptWith] to decrypt
  */
-
-
 @Serializable
 sealed class Container {
 
+    /**
+     * Exception thrown when container inner structure is bad. It could mean it was altered.
+     */
     class InvalidContainerException : Crypto2Exception("the container is invalid")
 
+    /**
+     * Attempt to decrypt container contents. This function caches decrypted data, so it is ok
+     * to call it more than once.
+     *
+     * @param keyRing a key ring with keys that caller wants to be used
+     * @return decrypted data or null if this ring contains no proper key for it
+     */
     abstract fun decryptWith(keyRing: UniversalRing): UByteArray?
 
-    fun decryptWith(vararg decryptingKeys: DecryptingKey): UByteArray? =
+    /**
-        decryptWith(UniversalRing(*decryptingKeys))
+     * Attempt to decrypt container contents. This function caches decrypted data, so it is ok
+     * to call it more than once.
+     *
+     * @param keys keys that caller wants to be used to decrypt with
+     * @return decrypted data or null if this ring contains no proper key for it
+     */
+    fun decryptWith(vararg keys: DecryptingKey): UByteArray? =
+        decryptWith(UniversalRing(*keys))
 
     @Transient
     var decryptedData: UByteArray? = null
 
+    /**
+     * Check the container is already decrypted. It is important, for example, to be sure it is
+     * before adding more keys.
+     */
     val isDecrypted: Boolean get() = decryptedData != null
 
+    /**
+     * @suppress
+     * Single-key variant, to conserve space it does not use the main key logic and just encrypts the data.
+     */
     @Serializable
     @SerialName("1")
-    class Single(val keyTag: KeyTag, val encryptedMessage: UByteArray) : Container() {
+    internal class Single(val keyTag: KeyTag, val encryptedMessage: UByteArray) : Container() {
+
+        @Transient
+        private var decryptedWithKey: DecryptingKey? = null
+
         override fun decryptWith(keyRing: UniversalRing): UByteArray? {
             decryptedData?.let { return it }
             for (k in keyRing) {
                 if (k.tag == keyTag) {
                     kotlin.runCatching { k.decrypt(encryptedMessage) }.getOrNull()?.let {
                         decryptedData = it
+                        decryptedWithKey = k
                         return it
                     }
                 }
             }
             return null
         }
+
+        internal val asOpenMulti: Container by lazy {
+            check(isDecrypted) { "container should be decrypted" }
+            create(decryptedData!!) {
+                alwaysMulti()
+                when (val k = decryptedWithKey!!) {
+                    is Asymmetric.SecretKey -> {
+                        key(k.publicKey)
+                    }
+
+                    is EncryptingKey -> {
+                        key(k)
+                    }
+
+                    is UniversalKey.Secret -> {
+                        key(k.key.publicKey)
+                    }
+
+                    else -> {
+                        throw IllegalStateException("unknown key type to convert container: ${k::class.simpleName}")
+                    }
+                }
+            }.apply { decryptWith(decryptedWithKey!!) ?: throw Crypto2Exception("internal error in container update (1)") }
+        }
     }
 
+    /**
+     * @suppress
+     * Implementation of 2+ keys container
+     */
     @Serializable
     @SerialName("*")
-    class Multi(val encryptedKeys: List<EncryptedKey>, val encryptedMessage: UByteArray) : Container() {
+    internal class Multi(val encryptedKeys: List<EncryptedKey>, val encryptedMessage: UByteArray) : Container() {
         @Serializable
-        class EncryptedKey(val tag: KeyTag, val cipherData: UByteArray)
+        class EncryptedKey(val tag: KeyTag, val cipherData: UByteArray) {
+            constructor(key: EncryptingKey, encodeMainKey: UByteArray) :
+                    this(key.tag, key.encrypt(encodeMainKey))
 
-        private var mainKey: SymmetricKey? = null
+            constructor(sender: Asymmetric.SecretKey?, recipient: Asymmetric.PublicKey, encodeMainKey: UByteArray) :
+                    this(
+                        recipient.tag,
+                        recipient.encryptMessage(
+                            encodeMainKey,
+                            senderKey = sender ?: Asymmetric.randomSecretKey(),
+                        ).encoded
+                    )
+        }
+
+        internal var mainKey: SymmetricKey? = null
 
         override fun decryptWith(keyRing: UniversalRing): UByteArray? {
             decryptedData?.let { return it }
@@ -68,7 +140,8 @@ sealed class Container {
                                 key.decrypt(encryptedKey.cipherData).toByteArray()
                             )
                         }.getOrNull()?.let { k ->
-                            if (kotlin.runCatching { decryptedData = k.decrypt(encryptedKey.cipherData) }.isFailure)
+                            println(k)
+                            if (kotlin.runCatching { decryptedData = k.decrypt(encryptedMessage) }.isFailure)
                                 throw InvalidContainerException()
                             mainKey = k
                         }
@@ -80,61 +153,204 @@ sealed class Container {
         }
     }
 
+    fun addRecipients(builder: Builder.() -> Unit): Container =
+        if (this is Single) asOpenMulti.addRecipients(builder)
+        else {
+            Builder(this).apply(builder).build()
+        }
 
+    operator fun plus(recipient: Asymmetric.PublicKey) = addRecipients { key(recipient) }
+    operator fun plus(recipient: EncryptingKey) = addRecipients { key(recipient) }
+    operator fun plus(pair: Pair<Asymmetric.SecretKey,Asymmetric.PublicKey>) = addRecipients { key(pair) }
+
+    /**
+     * Binary encoded version. It is desirable to include [Container] as an object, though,
+     * especially when using custom serialization (Json, Boss, etc), it is serializable.
+     * Still, if you need it in binary form, this is a shortcut. You can use [decode] or call
+     * [BipackDecoder.decode] to deserialize the binary form.
+     */
     val encoded: UByteArray by lazy {
         BipackEncoder.encode(this).toUByteArray()
     }
 
     companion object {
 
-        class Builder internal constructor(private val plainData: UByteArray) {
+        /**
+         * The builder to create container with various parameters.
+         * Use [create] to create container using a builder. Usage sample:
+         *
+         * ```kotlin
+         * Container.create(plainData) {
+         *     // optional: add a random filling from 10 to 20 bytes
+         *     fill( 10 .. 20 )
+         *
+         *     key(symmetricKey1, symmetricKey2) // add two SymmetricKey recipients
+         *
+         *     key(publicKey1) // add a Asymmetric.PublicKey recipient anonymously
+         *
+         *     // More interesting: add publicKey2 and publicKey3 recipients using my
+         *     // secret key as authority. IT is faster and allow to owner of the listed public keys
+         *     // to know it was me who added them to this container.
+         *     key(mySecretKey to publicKey2,mySecretKey to publicKey3)
+         * }
+         * ```
+
+         */
+        class Builder internal constructor(
+            private val plainData: UByteArray,
+            private var parent: Container? = null,
+        ) {
+
+            internal constructor(parent: Container) :
+                    this(
+                        parent.decryptedData ?: throw IllegalStateException("container is not decrypted"),
+                        parent
+                    )
+
+
             private val plainKeys = mutableListOf<EncryptingKey>()
-            private val keyPairs = mutableListOf<Pair<Asymmetric.SecretKey?,Asymmetric.PublicKey>>()
+            private val keyPairs = mutableListOf<Pair<Asymmetric.SecretKey?, Asymmetric.PublicKey>>()
             private var fillRange: IntRange? = null
 
-            fun key(vararg keys: EncryptingKey) { plainKeys.addAll(keys) }
+            /**
+             * Add one or more encrypting keys
+             */
+            fun key(vararg keys: EncryptingKey) {
+                plainKeys.addAll(keys)
+            }
 
-            fun key(vararg pairs: Pair<Asymmetric.SecretKey,Asymmetric.PublicKey>) {
+            /**
+             * Add one or more [Asymmetric.SecretKey] as sender authority coupled with [Asymmetric.PublicKey] as
+             * a recipient. This is faster than anonymous usage of [Asymmetric.PublicKey] only
+             */
+            fun key(vararg pairs: Pair<Asymmetric.SecretKey, Asymmetric.PublicKey>) {
                 keyPairs.addAll(pairs)
             }
 
+            /**
+             * Add one or more public keys as recipients. This is slower than using pairs of sender -> recipient.
+             */
             fun key(vararg publicKeys: Asymmetric.PublicKey) {
                 keyPairs.addAll(publicKeys.map { null to it })
             }
 
+            /**
+             * Causes random filling of the encrypted message in the specified interval
+             */
             @Suppress("unused")
             fun fill(range: IntRange) {
-                require(range.first >= 0 ) { "range must be positive"}
+                require(range.first >= 0) { "range must be positive" }
                 fillRange = range
             }
 
-            fun build(): Container {
+            private var makeMulti = false
-                return when( plainKeys.size + keyPairs.size ) {
+
-                    0 -> throw IllegalArgumentException("Container needs at least one key")
+            /**
-                    1 -> {
+             * @suppress
-                        plainKeys.firstOrNull()?.let {
+             * will produce multikey internal variant even with only one key. User internally
-                            Single(it.tag, it.encrypt(plainData, fillRange))
+             */
-                        } ?: run {
+            internal fun alwaysMulti() {
-                            val (sk, pk) = keyPairs.first()
+                makeMulti = true
-                            Single(pk.tag, pk.encryptMessage(plainData,
+            }
-                                senderKey = sk ?: Asymmetric.randomSecretKey(),
+
-                                randomFill = fillRange).encoded)
+            /**
-                        }
+             * Create a Container
+             */
+            internal fun build(): Container {
+                val countNewKeys = plainKeys.size + keyPairs.size
+                if (parent != null) require(parent is Multi) { "parent container mut be a multikey variant" }
+                return when {
+                    countNewKeys == 0 -> throw IllegalArgumentException("Container needs at least one key")
+                    countNewKeys == 1 && makeMulti == false && parent == null -> {
+                        createSingle()
                     }
+
                     else -> {
-                        TODO("multikey")
+                        val eks: MutableList<Multi.EncryptedKey>
+                        val mainKey: SymmetricKey
+                        val p = parent
+                        if (p != null) {
+                            p as Multi
+                            eks = p.encryptedKeys.toMutableList()
+                            mainKey = p.mainKey ?: throw IllegalStateException("parent container must be decrypted")
+                        } else {
+                            eks = mutableListOf<Multi.EncryptedKey>()
+                            mainKey = SymmetricKey.random()
+                        }
+                        val encodedMainKey = BipackEncoder.encode(mainKey).toUByteArray()
+                        createMulti(eks, encodedMainKey, mainKey)
                     }
                 }
             }
+
+            private fun createSingle() = plainKeys.firstOrNull()?.let {
+                Single(it.tag, it.encrypt(plainData, fillRange))
+            } ?: run {
+                val (sk, pk) = keyPairs.first()
+                Single(
+                    pk.tag, pk.encryptMessage(
+                        plainData,
+                        senderKey = sk ?: Asymmetric.randomSecretKey(),
+                        randomFill = fillRange
+                    ).encoded
+                )
+            }
+
+            private fun createMulti(
+                eks: MutableList<Multi.EncryptedKey>,
+                encodedMainKey: UByteArray,
+                mainKey: SymmetricKey,
+            ): Multi {
+                for (k in plainKeys)
+                    eks += Multi.EncryptedKey(k, encodedMainKey)
+                for (p in keyPairs) {
+                    val (sender, recipient) = p
+                    eks += Multi.EncryptedKey(sender, recipient, encodedMainKey)
+                }
+                return Multi(eks, mainKey.encrypt(plainData, fillRange))
+            }
         }
 
-        fun create(plainData: UByteArray,builder: Builder.()->Unit) =
+        /**
+         * Create a container using a [Builder] instance.
+         * Usage sample:
+         * ```kotlin
+         * Container.create(plainData) {
+         *     // optional: add a random filling from 10 to 20 bytes
+         *     fill( 10 .. 20 )
+         *
+         *     key(symmetricKey1, symmetricKey2) // add two SymmetricKey recipients
+         *
+         *     key(publicKey1) // add a Asymmetric.PublicKey recipient anonymously
+         *
+         *     // More interesting: add publicKey2 and publicKey3 recipients using my
+         *     // secret key as authority. IT is faster and allow to owner of the listed public keys
+         *     // to know it was me who added them to this container.
+         *     key(mySecretKey to publicKey2,mySecretKey to publicKey3)
+         * }
+         * ```
+         * At least one key should be provided.
+         *
+         * @param plainData data to encrypt
+         */
+        fun create(plainData: UByteArray, builder: Builder.() -> Unit) =
             Builder(plainData).also { it.builder() }.build()
 
-        fun create(plainData: UByteArray, vararg keys: EncryptingKey): Container =
+        /**
-            create(plainData) { key(*keys) }
+         * Create container using one or more [EncryptingKey] and a builder, see [create]
+         * for builder usage sample.
+         */
+        fun createWith(
+            plainData: UByteArray, vararg keys: EncryptingKey,
+            builder: (Builder.() -> Unit)? = null,
+        ): Container =
+            create(plainData) { key(*keys); builder?.invoke(this) }
 
-        fun create(plainData: UByteArray, vararg keys: Pair<Asymmetric.SecretKey,Asymmetric.PublicKey>) =
+        /**
+         * Create the container using one or more `sender to recipient` asymmetric keys and a builder. See [create]
+         * for builder usage sample.
+         */
+        fun createWith(plainData: UByteArray, vararg keys: Pair<Asymmetric.SecretKey, Asymmetric.PublicKey>) =
             create(plainData) { key(*keys) }
 
         fun decode(encoded: UByteArray): Container {

src/commonMain/kotlin/net/sergeych/crypto2/UniversalKey.kt

@@ -45,6 +45,7 @@ sealed class UniversalKey : DecryptingKey {
     companion object {
         fun from(key: DecryptingKey) =
             when (key) {
+                is UniversalKey -> key
                 is Asymmetric.SecretKey -> Secret(key)
                 is SymmetricKey -> Symmetric(key)
                 is SafeKeyExchange.SessionKey -> Session(key)

src/commonTest/kotlin/ContainerTest.kt

@@ -1,9 +1,6 @@
 import com.ionspin.kotlin.crypto.util.encodeToUByteArray
 import kotlinx.coroutines.test.runTest
-import net.sergeych.crypto2.Asymmetric
+import net.sergeych.crypto2.*
-import net.sergeych.crypto2.Container
-import net.sergeych.crypto2.SymmetricKey
-import net.sergeych.crypto2.initCrypto
 import kotlin.test.*
 
 class ContainerTest {
@@ -14,7 +11,7 @@ class ContainerTest {
         val syk2 = SymmetricKey.random()
         val data = "sergeych, ohm many.".encodeToUByteArray()
 
-        val c = Container.create(data, syk1)
+        val c = Container.createWith(data, syk1)
         assertFalse { c.isDecrypted }
         val c1 = Container.decode(c.encoded)
         assertFalse { c.isDecrypted }
@@ -35,7 +32,7 @@ class ContainerTest {
         val p3 = Asymmetric.generateKeys()
         val data = "sergeych, ohm many.".encodeToUByteArray()
 
-        val c = Container.create(data, p1.secretKey to p2.publicKey)
+        val c = Container.createWith(data, p1.secretKey to p2.publicKey)
         assertFalse { c.isDecrypted }
         val c1 = Container.decode(c.encoded)
         assertFalse { c.isDecrypted }
@@ -74,13 +71,125 @@ class ContainerTest {
         initCrypto()
         val syk1 = SymmetricKey.random()
         val syk2 = SymmetricKey.random()
-//        val syk3 = SymmetricKey.random()
+        val syk3 = SymmetricKey.random()
+        val p1 = Asymmetric.generateKeys()
+        val p2 = Asymmetric.generateKeys()
+        val p3 = Asymmetric.generateKeys()
+        val p4 = Asymmetric.generateKeys()
         val data = "Translating the name 'Sergey Chernov' from Russian to archaic Sanskrit would be 'Ramo Krishna'"
             .encodeToUByteArray()
 
-        val c = Container.create(data, syk1, syk2)
+        val c = Container.createWith(data, syk1, syk2) {
+            key(p1.secretKey to p3.publicKey)
+            key(p4.publicKey)
+        }
         assertFalse { c.isDecrypted }
-        val c1 = Container.decode(c.encoded)
+        var c1 = Container.decode(c.encoded)
         assertFalse { c1.isDecrypted }
+
+        assertNull(c1.decryptWith(syk3))
+        assertFalse { c1.isDecrypted }
+
+        assertContentEquals(data, c1.decryptWith(syk3, syk1))
+        assertTrue { c1.isDecrypted }
+
+        c1 = Container.decode(c.encoded)
+        assertFalse { c1.isDecrypted }
+        assertNull(c1.decryptWith(p2.secretKey, p1.secretKey))
+        assertContentEquals(data, c1.decryptWith(syk3, p3.secretKey))
+
+        c1 = Container.decode(c.encoded)
+        assertFalse { c1.isDecrypted }
+        assertContentEquals(data, c1.decryptWith(syk3, p4.secretKey))
+    }
+
+    @Test
+    fun testSingleGrowSymmetric() = runTest {
+        initCrypto()
+        val syk1 = SymmetricKey.random()
+        val syk2 = SymmetricKey.random()
+        val syk3 = SymmetricKey.random()
+        val p1 = Asymmetric.generateKeys()
+        val p3 = Asymmetric.generateKeys()
+        val p4 = Asymmetric.generateKeys()
+        val data = "Translating the name 'Sergey Chernov' from Russian to archaic Sanskrit would be 'Ramo Krishna'"
+            .encodeToUByteArray()
+
+        var c = Container.createWith(data, syk1)
+
+        fun expectOpen(k: DecryptingKey) {
+            val c1 = Container.decode(c.encoded)
+            assertContentEquals(data, c1.decryptWith(k))
+        }
+
+        fun expectNotOpen(k: DecryptingKey) {
+            val c1 = Container.decode(c.encoded)
+            assertNull(c1.decryptWith(k))
+        }
+
+        expectOpen(syk1)
+        expectNotOpen(syk2)
+        expectNotOpen(p3.secretKey)
+
+        c.decryptWith(syk1)
+        assertTrue { c.isDecrypted }
+        assertNotNull(c.decryptedData)
+        c += syk2
+        expectOpen(syk1)
+        expectOpen(syk2)
+        expectNotOpen(syk3)
+        expectNotOpen(p3.secretKey)
+
+        c.decryptWith(syk1)
+        c += p3.publicKey
+        expectOpen(syk1)
+        expectOpen(syk2)
+        expectOpen(p3.secretKey)
+        expectNotOpen(syk3)
+        expectNotOpen(p4.secretKey)
+
+        c.decryptWith(syk1)
+        c += p1.secretKey to p4.publicKey
+        expectOpen(syk1)
+        expectOpen(syk2)
+        expectOpen(p3.secretKey)
+        expectNotOpen(syk3)
+        expectOpen(p4.secretKey)
+    }
+
+    @Test
+    fun testSingleGrowAsymmetric() = runTest {
+        initCrypto()
+        val syk1 = SymmetricKey.random()
+        val syk2 = SymmetricKey.random()
+        val syk3 = SymmetricKey.random()
+        val p1 = Asymmetric.generateKeys()
+        val p2 = Asymmetric.generateKeys()
+        val p3 = Asymmetric.generateKeys()
+        val p4 = Asymmetric.generateKeys()
+        val data = "Translating the name 'Sergey Chernov' from Russian to archaic Sanskrit would be 'Ramo Krishna'"
+            .encodeToUByteArray()
+
+        var c = Container.createWith(data, p1.secretKey to p3.publicKey)
+
+        fun expectOpen(k: DecryptingKey) {
+            val c1 = Container.decode(c.encoded)
+            assertContentEquals(data, c1.decryptWith(k))
+        }
+
+        fun expectNotOpen(k: DecryptingKey) {
+            val c1 = Container.decode(c.encoded)
+            assertNull(c1.decryptWith(k))
+        }
+
+        expectNotOpen(syk1)
+        expectOpen(p3.secretKey)
+
+        c.decryptWith(p3.secretKey)
+
+        c += syk1
+        expectOpen(syk1)
+        expectNotOpen(syk2)
+        expectOpen(p3.secretKey)
     }
 }