Allow to swap KDF implementation using KdfProvider interface

kotpass/src/main/kotlin/app/keemobile/kotpass/cryptography/KeyTransform.kt

@@ -1,12 +1,9 @@
 package app.keemobile.kotpass.cryptography
 
-import app.keemobile.kotpass.cryptography.engines.Argon2Engine
-import app.keemobile.kotpass.cryptography.format.AesKdf
-import app.keemobile.kotpass.cryptography.format.Argon2Kdf
+import app.keemobile.kotpass.cryptography.format.KdfProvider
 import app.keemobile.kotpass.database.Credentials
 import app.keemobile.kotpass.database.header.DatabaseHeader
 import app.keemobile.kotpass.database.header.KdfParameters.Aes
-import app.keemobile.kotpass.database.header.KdfParameters.Argon2
 import app.keemobile.kotpass.extensions.clear
 import app.keemobile.kotpass.extensions.sha256
 import app.keemobile.kotpass.extensions.sha512
@@ -27,42 +24,30 @@ internal object KeyTransform {
             .also { composite.clear() }
     }
 
-    fun transformedKey(header: DatabaseHeader, credentials: Credentials): ByteArray {
-        return when (header) {
-            is DatabaseHeader.Ver3x -> {
-                AesKdf.transformKey(
-                    key = compositeKey(credentials),
-                    seed = header.transformSeed.toByteArray(),
-                    rounds = header.transformRounds
-                )
-            }
-            is DatabaseHeader.Ver4x -> {
-                when (header.kdfParameters) {
-                    is Aes -> {
-                        AesKdf.transformKey(
-                            key = compositeKey(credentials),
-                            seed = header.kdfParameters.seed.toByteArray(),
-                            rounds = header.kdfParameters.rounds
-                        )
-                    }
-                    is Argon2 -> {
-                        Argon2Kdf.transformKey(
-                            variant = when (header.kdfParameters.variant) {
-                                Argon2.Variant.Argon2d -> Argon2Engine.Variant.Argon2d
-                                Argon2.Variant.Argon2id -> Argon2Engine.Variant.Argon2id
-                            },
-                            version = Argon2Engine.Version.from(header.kdfParameters.version),
-                            password = compositeKey(credentials),
-                            salt = header.kdfParameters.salt.toByteArray(),
-                            secretKey = header.kdfParameters.secretKey?.toByteArray(),
-                            additional = header.kdfParameters.associatedData?.toByteArray(),
-                            iterations = header.kdfParameters.iterations,
-                            parallelism = header.kdfParameters.parallelism,
-                            memory = header.kdfParameters.memory
-                        )
-                    }
+    fun transformedKey(
+        kdfProvider: KdfProvider,
+        header: DatabaseHeader,
+        credentials: Credentials
+    ): ByteArray {
+        val compositeKey = compositeKey(credentials)
+
+        return try {
+            when (header) {
+                is DatabaseHeader.Ver3x -> {
+                    // KeePass 3.x supports only AES as key-derivation function
+                    kdfProvider.transformKey(
+                        kdfParameters = Aes(header.transformRounds, header.transformSeed),
+                        compositeKey = compositeKey
+                    )
+                }
+                is DatabaseHeader.Ver4x -> {
+                    kdfProvider.transformKey(header.kdfParameters, compositeKey)
                 }
             }
+        } catch (error: Exception) {
+            throw error
+        } finally {
+            compositeKey.clear()
         }
     }
 

kotpass/src/main/kotlin/app/keemobile/kotpass/cryptography/format/AesKdf.kt

@@ -1,7 +1,6 @@
 package app.keemobile.kotpass.cryptography.format
 
 import app.keemobile.kotpass.errors.CryptoError
-import app.keemobile.kotpass.extensions.clear
 import app.keemobile.kotpass.extensions.sha256
 import java.security.GeneralSecurityException
 import java.security.NoSuchAlgorithmException
@@ -18,15 +17,13 @@ internal object AesKdf {
         val keySpec = SecretKeySpec(seed, "AES")
         cipher.init(Cipher.ENCRYPT_MODE, keySpec)
 
-        for (r in 0 until rounds.toLong()) {
+        repeat(rounds.toInt()) {
             cipher.update(key, 0, 16, key, 0)
             cipher.update(key, 16, 16, key, 16)
         }
-        key.sha256().also {
-            key.clear()
-        }
-    } catch (e: GeneralSecurityException) {
-        if (e is NoSuchAlgorithmException) {
+        key.sha256()
+    } catch (error: GeneralSecurityException) {
+        if (error is NoSuchAlgorithmException) {
             throw CryptoError.AlgorithmUnavailable("AES/ECB encryption is not supported in current environment.")
         } else {
             throw CryptoError.InvalidKey("Wrong KDF seed used for decryption.")

kotpass/src/main/kotlin/app/keemobile/kotpass/cryptography/format/BaseKdfProvider.kt

@@ -0,0 +1,37 @@
+package app.keemobile.kotpass.cryptography.format
+
+import app.keemobile.kotpass.cryptography.engines.Argon2Engine
+import app.keemobile.kotpass.database.header.KdfParameters
+import app.keemobile.kotpass.database.header.KdfParameters.Aes
+import app.keemobile.kotpass.database.header.KdfParameters.Argon2
+
+internal object BaseKdfProvider : KdfProvider {
+    override fun transformKey(
+        kdfParameters: KdfParameters,
+        compositeKey: ByteArray
+    ): ByteArray = when (kdfParameters) {
+        is Aes -> {
+            AesKdf.transformKey(
+                key = compositeKey,
+                seed = kdfParameters.seed.toByteArray(),
+                rounds = kdfParameters.rounds
+            )
+        }
+        is Argon2 -> {
+            Argon2Kdf.transformKey(
+                variant = when (kdfParameters.variant) {
+                    Argon2.Variant.Argon2d -> Argon2Engine.Variant.Argon2d
+                    Argon2.Variant.Argon2id -> Argon2Engine.Variant.Argon2id
+                },
+                version = Argon2Engine.Version.from(kdfParameters.version),
+                password = compositeKey,
+                salt = kdfParameters.salt.toByteArray(),
+                secretKey = kdfParameters.secretKey?.toByteArray(),
+                additional = kdfParameters.associatedData?.toByteArray(),
+                iterations = kdfParameters.iterations,
+                parallelism = kdfParameters.parallelism,
+                memory = kdfParameters.memory
+            )
+        }
+    }
+}

kotpass/src/main/kotlin/app/keemobile/kotpass/cryptography/format/KdfProvider.kt

@@ -0,0 +1,13 @@
+package app.keemobile.kotpass.cryptography.format
+
+import app.keemobile.kotpass.database.header.KdfParameters
+
+interface KdfProvider {
+    /**
+     * Transforms `compositeKey` using key-derivation function.
+     *
+     * @param kdfParameters Describes key-derivation function type and parameters.
+     * @param compositeKey Key bytes.
+     */
+    fun transformKey(kdfParameters: KdfParameters, compositeKey: ByteArray): ByteArray
+}

kotpass/src/main/kotlin/app/keemobile/kotpass/database/Decoder.kt

@@ -4,7 +4,9 @@ import app.keemobile.kotpass.constants.Defaults
 import app.keemobile.kotpass.cryptography.EncryptionSaltGenerator
 import app.keemobile.kotpass.cryptography.KeyTransform
 import app.keemobile.kotpass.cryptography.format.BaseCiphers
+import app.keemobile.kotpass.cryptography.format.BaseKdfProvider
 import app.keemobile.kotpass.cryptography.format.CipherProvider
+import app.keemobile.kotpass.cryptography.format.KdfProvider
 import app.keemobile.kotpass.database.header.DatabaseHeader
 import app.keemobile.kotpass.database.header.DatabaseHeader.Compression
 import app.keemobile.kotpass.database.header.DatabaseInnerHeader
@@ -30,6 +32,7 @@ fun KeePassDatabase.Companion.decode(
     validateHashes: Boolean = true,
     contentParser: XmlContentParser = DefaultXmlContentParser,
     cipherProviders: List<CipherProvider> = BaseCiphers.entries,
+    kdfProvider: KdfProvider = BaseKdfProvider,
     untitledLabel: String = Defaults.UntitledLabel
 ): KeePassDatabase {
     val headerBuffer = Buffer()
@@ -47,7 +50,7 @@ fun KeePassDatabase.Companion.decode(
             throw FormatError.UnsupportedVersion("File version is not supported.")
         }
         val rawHeaderData = headerBuffer.snapshot()
-        val transformedKey = KeyTransform.transformedKey(header, credentials)
+        val transformedKey = KeyTransform.transformedKey(kdfProvider, header, credentials)
 
         return when (header) {
             is DatabaseHeader.Ver3x -> {

kotpass/src/main/kotlin/app/keemobile/kotpass/database/Encoder.kt

@@ -3,7 +3,9 @@ package app.keemobile.kotpass.database
 import app.keemobile.kotpass.cryptography.EncryptionSaltGenerator
 import app.keemobile.kotpass.cryptography.KeyTransform
 import app.keemobile.kotpass.cryptography.format.BaseCiphers
+import app.keemobile.kotpass.cryptography.format.BaseKdfProvider
 import app.keemobile.kotpass.cryptography.format.CipherProvider
+import app.keemobile.kotpass.cryptography.format.KdfProvider
 import app.keemobile.kotpass.database.header.DatabaseHeader
 import app.keemobile.kotpass.database.modifiers.binaries
 import app.keemobile.kotpass.database.modifiers.regenerateVectors
@@ -25,19 +27,18 @@ fun KeePassDatabase.encode(
     outputStream: OutputStream,
     contentParser: XmlContentParser = DefaultXmlContentParser,
     cipherProviders: List<CipherProvider> = BaseCiphers.entries,
+    kdfProvider: KdfProvider = BaseKdfProvider,
     random: SecureRandom = SecureRandom()
 ) = regenerateVectors(random, cipherProviders)
-    .encodeAsBinary(outputStream, contentParser, cipherProviders)
+    .encodeAsBinary(outputStream, contentParser, cipherProviders, kdfProvider)
 
 private fun KeePassDatabase.encodeAsBinary(
     outputStream: OutputStream,
     contentParser: XmlContentParser = DefaultXmlContentParser,
-    cipherProviders: List<CipherProvider>
+    cipherProviders: List<CipherProvider>,
+    kdfProvider: KdfProvider
 ) = apply {
-    val transformedKey = KeyTransform.transformedKey(
-        header = header,
-        credentials = credentials
-    )
+    val transformedKey = KeyTransform.transformedKey(kdfProvider, header, credentials)
     val headerBuffer = Buffer().apply {
         header.writeTo(this)
     }

kotpass/src/main/kotlin/app/keemobile/kotpass/database/KeePassDatabase.kt

@@ -4,6 +4,8 @@ package app.keemobile.kotpass.database
 
 import app.keemobile.kotpass.constants.GroupOverride
 import app.keemobile.kotpass.cryptography.KeyTransform
+import app.keemobile.kotpass.cryptography.format.BaseKdfProvider
+import app.keemobile.kotpass.cryptography.format.KdfProvider
 import app.keemobile.kotpass.database.header.DatabaseHeader
 import app.keemobile.kotpass.database.header.DatabaseInnerHeader
 import app.keemobile.kotpass.extensions.clear
@@ -224,8 +226,10 @@ fun KeePassDatabase.findEntries(
  * Measures KDF transform rounds performance based on
  * [header][KeePassDatabase.header] parameters.
  */
-fun KeePassDatabase.measureKeyTransformMillis() = measureTimeMillis {
+fun KeePassDatabase.measureKeyTransformMillis(
+    kdfProvider: KdfProvider = BaseKdfProvider
+): Long = measureTimeMillis {
     KeyTransform
-        .transformedKey(header, credentials)
+        .transformedKey(kdfProvider, header, credentials)
         .clear()
 }