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

2025-09-27 09:17:46 +02:00 · 2025-09-15 15:15:44 -04:00
parent bd70058f78
commit 98aaef61c7
5 changed files with 252 additions and 40 deletions
--- a/benchmark/src/androidTest/java/com/vitorpamplona/quartz/benchmark/EncryptDecrypt.kt
+++ b/benchmark/src/androidTest/java/com/vitorpamplona/quartz/benchmark/EncryptDecrypt.kt
@@ -25,7 +25,6 @@ import androidx.benchmark.junit4.measureRepeated
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.vitorpamplona.quartz.nip01Core.crypto.Nip01
 import com.vitorpamplona.quartz.nip44Encryption.Nip44v2
-import com.vitorpamplona.quartz.nip44Encryption.crypto.Hkdf
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -48,39 +47,14 @@ class EncryptDecrypt {
    @Test
    fun encrypt() {
        benchmarkRule.measureRepeated {
-            nip44v2.encrypt(msg, sharedKey)
+            nip44v2.encrypt(msg, privateKey, publicKey)
        }
    }

    @Test
    fun decrypt() {
        benchmarkRule.measureRepeated {
-            nip44v2.decrypt(encrypted, sharedKey)
-        }
-    }
-
-    @Test
-    fun messageKeys() {
-        benchmarkRule.measureRepeated {
-            nip44v2.getMessageKeys(sharedKey, encrypted.nonce)
-        }
-    }
-
-    @Test
-    fun checkMac() {
-        val messageKey = nip44v2.getMessageKeys(sharedKey, encrypted.nonce)
-
-        benchmarkRule.measureRepeated {
-            nip44v2.checkHMacAad(messageKey, encrypted)
-        }
-    }
-
-    @Test
-    fun hkdfExtract() {
-        val messageKey = nip44v2.getMessageKeys(sharedKey, encrypted.nonce)
-        val hkdf = Hkdf()
-        benchmarkRule.measureRepeated {
-            hkdf.extract(encrypted.nonce, encrypted.ciphertext, messageKey.hmacKey)
+            nip44v2.decrypt(encrypted, privateKey, publicKey)
        }
    }
 }
--- a/benchmark/src/androidTest/java/com/vitorpamplona/quartz/benchmark/Hkdf.kt
+++ b/benchmark/src/androidTest/java/com/vitorpamplona/quartz/benchmark/Hkdf.kt
@@ -0,0 +1,72 @@
+/**
+ * Copyright (c) 2025 Vitor Pamplona
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy of
+ * this software and associated documentation files (the "Software"), to deal in
+ * the Software without restriction, including without limitation the rights to use,
+ * copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the
+ * Software, and to permit persons to whom the Software is furnished to do so,
+ * subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+ * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+ * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN
+ * AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
+ * WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ */
+package com.vitorpamplona.quartz.benchmark
+
+import androidx.benchmark.junit4.BenchmarkRule
+import androidx.benchmark.junit4.measureRepeated
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.vitorpamplona.quartz.nip01Core.crypto.Nip01
+import com.vitorpamplona.quartz.nip44Encryption.Nip44v2
+import com.vitorpamplona.quartz.nip44Encryption.crypto.Hkdf
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(AndroidJUnit4::class)
+class Hkdf {
+    @get:Rule val benchmarkRule = BenchmarkRule()
+
+    companion object {
+        val nip44v2 = Nip44v2()
+        val msg = "Hi, how are you? this is supposed to be representative of an average message on Nostr"
+
+        val privateKey = Nip01.privKeyCreate()
+        val publicKey = Nip01.pubKeyCreate(privateKey)
+
+        val sharedKey = nip44v2.getConversationKey(privateKey, publicKey)
+        val encrypted = nip44v2.encrypt(msg, sharedKey)
+    }
+
+    @Test
+    fun hkdfExpand() {
+        val hkdf = Hkdf()
+        benchmarkRule.measureRepeated {
+            hkdf.fastExpand(sharedKey, encrypted.nonce)
+        }
+    }
+
+    @Test
+    fun hkdfExpandOld() {
+        val hkdf = Hkdf()
+        benchmarkRule.measureRepeated {
+            hkdf.expand(sharedKey, encrypted.nonce, 76)
+        }
+    }
+
+    @Test
+    fun hkdfExtract() {
+        val messageKey = nip44v2.getMessageKeys(sharedKey, encrypted.nonce)
+        val hkdf = Hkdf()
+        benchmarkRule.measureRepeated {
+            hkdf.extract(encrypted.nonce, encrypted.ciphertext, messageKey.hmacKey)
+        }
+    }
+}
--- a/quartz/src/androidMain/kotlin/com/vitorpamplona/quartz/nip44Encryption/Nip44v2.kt
+++ b/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,
--- a/quartz/src/androidMain/kotlin/com/vitorpamplona/quartz/nip44Encryption/crypto/Hkdf.kt
+++ b/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 }
+}
--- a/quartz/src/androidUnitTest/kotlin/com/vitorpamplona/quartz/nip44Encryption/crypto/HkdfText.kt
+++ b/quartz/src/androidUnitTest/kotlin/com/vitorpamplona/quartz/nip44Encryption/crypto/HkdfText.kt
@@ -0,0 +1,103 @@
+/**
+ * Copyright (c) 2025 Vitor Pamplona
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy of
+ * this software and associated documentation files (the "Software"), to deal in
+ * the Software without restriction, including without limitation the rights to use,
+ * copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the
+ * Software, and to permit persons to whom the Software is furnished to do so,
+ * subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+ * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+ * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN
+ * AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
+ * WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ */
+package com.vitorpamplona.quartz.nip44Encryption.crypto
+
+import com.vitorpamplona.quartz.nip01Core.core.toHexKey
+import org.junit.Test
+import kotlin.test.assertEquals
+
+class HkdfText {
+    companion object {
+        val hkdf = Hkdf()
+    }
+
+    @Test
+    fun testExpand1() {
+        val result =
+            hkdf.expand(
+                "b0a6232d3e5444dbf2291f3290504989716f1dc78ac2da812230720fd76c8f06".hexToByteArray(),
+                "c9300041754c97f030ec9c35db607cb642db0f4305f178692954f3a883ea3209".hexToByteArray(),
+                76,
+            )
+
+        assertEquals("9077071b92dfbd725f5a1dd2dbbf349615f1d0ce22b436e2f505cc05984ef588", result.hmacKey.toHexKey())
+        assertEquals("3256a46b762158a6e08338daf2c5b3eefd2bed343cb22df8b6315567b9207cef", result.chachaKey.toHexKey())
+        assertEquals("670bb239ac07dbb801653aa5", result.chachaNonce.toHexKey())
+
+        val result2 =
+            hkdf.fastExpand(
+                "b0a6232d3e5444dbf2291f3290504989716f1dc78ac2da812230720fd76c8f06".hexToByteArray(),
+                "c9300041754c97f030ec9c35db607cb642db0f4305f178692954f3a883ea3209".hexToByteArray(),
+            )
+
+        assertEquals("9077071b92dfbd725f5a1dd2dbbf349615f1d0ce22b436e2f505cc05984ef588", result2.hmacKey.toHexKey())
+        assertEquals("3256a46b762158a6e08338daf2c5b3eefd2bed343cb22df8b6315567b9207cef", result2.chachaKey.toHexKey())
+        assertEquals("670bb239ac07dbb801653aa5", result2.chachaNonce.toHexKey())
+    }
+
+    @Test
+    fun testExpand2() {
+        val result =
+            hkdf.expand(
+                "bc5d4e032696ef107ef1c7b6fc5f00c6e7b31ae4f86ee486ce24aa0d84847d83".hexToByteArray(),
+                "830426c471c0870afb0208cabcaa0e4d66fe51af7163336b7b9ec1846c31e900".hexToByteArray(),
+                76,
+            )
+
+        assertEquals("26ccd9471681fc42459dbf7b14fc54d3f5276ddad482f7c9f81dc021ccbe5592", result.hmacKey.toHexKey())
+        assertEquals("a13a519db107cf9ecbca6ba79ab2428fcd624286025b7ee452f153ae770f07a8", result.chachaKey.toHexKey())
+        assertEquals("4c8218640c43de928f4c52e1", result.chachaNonce.toHexKey())
+
+        val result2 =
+            hkdf.fastExpand(
+                "bc5d4e032696ef107ef1c7b6fc5f00c6e7b31ae4f86ee486ce24aa0d84847d83".hexToByteArray(),
+                "830426c471c0870afb0208cabcaa0e4d66fe51af7163336b7b9ec1846c31e900".hexToByteArray(),
+            )
+
+        assertEquals("26ccd9471681fc42459dbf7b14fc54d3f5276ddad482f7c9f81dc021ccbe5592", result2.hmacKey.toHexKey())
+        assertEquals("a13a519db107cf9ecbca6ba79ab2428fcd624286025b7ee452f153ae770f07a8", result2.chachaKey.toHexKey())
+        assertEquals("4c8218640c43de928f4c52e1", result2.chachaNonce.toHexKey())
+    }
+
+    @Test
+    fun testExpand3() {
+        val result =
+            hkdf.expand(
+                "0000000000000000000000000000000000000000000000000000000000000000".hexToByteArray(),
+                "0000000000000000000000000000000000000000000000000000000000000000".hexToByteArray(),
+                76,
+            )
+
+        assertEquals("a30b80c908d01aa28eae9dcdec55ac104f9c889b989ee563919985a84f66360d", result.hmacKey.toHexKey())
+        assertEquals("3769af12ff4dbf44e516a22d1d0512e8bc42516d59e8bf401ea346a4d60dccf7", result.chachaKey.toHexKey())
+        assertEquals("77938d29bb13ea73f677ac27", result.chachaNonce.toHexKey())
+
+        val result2 =
+            hkdf.fastExpand(
+                "0000000000000000000000000000000000000000000000000000000000000000".hexToByteArray(),
+                "0000000000000000000000000000000000000000000000000000000000000000".hexToByteArray(),
+            )
+
+        assertEquals("a30b80c908d01aa28eae9dcdec55ac104f9c889b989ee563919985a84f66360d", result2.hmacKey.toHexKey())
+        assertEquals("3769af12ff4dbf44e516a22d1d0512e8bc42516d59e8bf401ea346a4d60dccf7", result2.chachaKey.toHexKey())
+        assertEquals("77938d29bb13ea73f677ac27", result2.chachaNonce.toHexKey())
+    }
+}