Faster Hkdf functions with less array copying and allocations (which can be impactful if the ciphertext is large)

quartz/src/androidMain/kotlin/com/vitorpamplona/quartz/nip44Encryption/Nip44v2.kt

@@ -216,14 +216,7 @@ class Nip44v2 {
     fun getMessageKeys(
         conversationKey: ByteArray,
         nonce: ByteArray,
-    ): MessageKey {
-        val keys = hkdf.expand(conversationKey, nonce, 76)
-        return MessageKey(
-            chachaKey = keys.copyOfRange(0, 32),
-            chachaNonce = keys.copyOfRange(32, 44),
-            hmacKey = keys.copyOfRange(44, 76),
-        )
-    }
+    ): MessageKey = hkdf.fastExpand(conversationKey, nonce)
 
     class MessageKey(
         val chachaKey: ByteArray,

quartz/src/androidMain/kotlin/com/vitorpamplona/quartz/nip44Encryption/crypto/Hkdf.kt

@@ -20,8 +20,10 @@
  */
 package com.vitorpamplona.quartz.nip44Encryption.crypto
 
+import com.vitorpamplona.quartz.nip44Encryption.Nip44v2.MessageKey
 import java.nio.ByteBuffer
 import javax.crypto.Mac
+import javax.crypto.SecretKey
 import javax.crypto.spec.SecretKeySpec
 
 class Hkdf(
@@ -33,7 +35,7 @@ class Hkdf(
         salt: ByteArray,
     ): ByteArray {
         val mac = Mac.getInstance(algorithm)
-        mac.init(SecretKeySpec(salt, algorithm))
+        mac.init(HMacKey(salt))
         return mac.doFinal(key)
     }
 
@@ -43,7 +45,7 @@ class Hkdf(
         salt: ByteArray,
     ): ByteArray {
         val mac = Mac.getInstance(algorithm)
-        mac.init(SecretKeySpec(salt, algorithm))
+        mac.init(HMacKey(salt))
         mac.update(key1)
         mac.update(key2)
         return mac.doFinal()
@@ -53,7 +55,7 @@ class Hkdf(
         key: ByteArray,
         nonce: ByteArray,
         outputLength: Int,
-    ): ByteArray {
+    ): MessageKey {
         check(key.size == hashLen)
         check(nonce.size == hashLen)
 
@@ -74,6 +76,74 @@ class Hkdf(
         val result = ByteArray(outputLength)
         generatedBytes.rewind()
         generatedBytes[result, 0, outputLength]
-        return result
+
+        return MessageKey(
+            chachaKey = result.copyOfRange(0, 32),
+            chachaNonce = result.copyOfRange(32, 44),
+            hmacKey = result.copyOfRange(44, 76),
+        )
+    }
+
+    /**
+     * Expands with outputLength == 76 while using the least amount of memory allocation
+     */
+    fun fastExpand(
+        key: ByteArray,
+        nonce: ByteArray,
+    ): MessageKey {
+        check(key.size == hashLen)
+        check(nonce.size == hashLen)
+
+        val mac = Mac.getInstance(algorithm)
+        mac.init(HMacKey(key))
+
+        // First round: T(1) = HMAC-SHA256(key, nonce || 0x01)
+        mac.update(nonce)
+        mac.update(1)
+        val round1 = mac.doFinal()
+
+        // Second round: T(2) = HMAC-SHA256(key, T(1) || nonce || 0x02)
+        mac.update(round1)
+        mac.update(nonce)
+        mac.update(2)
+        val round2 = mac.doFinal()
+
+        // Third round: T(3) = HMAC-SHA256(key, T(2) || nonce || 0x03)
+        mac.update(round2)
+        mac.update(nonce)
+        mac.update(3)
+        val round3 = mac.doFinal()
+
+        val hmacKey = ByteArray(32)
+        System.arraycopy(round2, 12, hmacKey, 0, 20)
+        System.arraycopy(round3, 0, hmacKey, 20, 12)
+
+        return MessageKey(
+            chachaKey = round1,
+            chachaNonce = round2.copyOfRange(0, 12),
+            hmacKey = hmacKey,
+        )
     }
 }
+
+class HMacKey(
+    val key: ByteArray,
+) : SecretKey {
+    override fun getAlgorithm() = "HmacSHA256"
+
+    override fun getEncoded() = key
+
+    override fun getFormat() = "RAW"
+
+    override fun hashCode() = key.contentHashCode()
+
+    override fun equals(other: Any?): Boolean {
+        if (this === other) return true
+        if (other !is HMacKey) return false
+        return key.contentEquals(other.key)
+    }
+
+    override fun destroy() = key.fill(0)
+
+    override fun isDestroyed() = key.all { it.toInt() == 0 }
+}