NoiseKK.kt

1package com.peachpasswords.sync23import org.bouncycastle.crypto.digests.Blake2sDigest4import org.bouncycastle.crypto.generators.HKDFBytesGenerator5import org.bouncycastle.crypto.modes.ChaCha20Poly13056import org.bouncycastle.crypto.params.AEADParameters7import org.bouncycastle.crypto.params.HKDFParameters8import org.bouncycastle.crypto.params.KeyParameter9import org.bouncycastle.crypto.params.X25519PrivateKeyParameters10import java.nio.ByteBuffer11import java.nio.ByteOrder12import java.security.SecureRandom1314private const val NOISE_KK_PROTOCOL_NAME = "Noise_KK_25519_ChaChaPoly_BLAKE2s"1516internal data class NoiseKKKeyPair(17    val publicKey: ByteArray,18    val secretKey: ByteArray,19)2021private fun initKKHash(): ByteArray {22    val name = NOISE_KK_PROTOCOL_NAME.encodeToByteArray()23    if (name.size <= 32) return ByteArray(32).also { name.copyInto(it) }24    return Blake2sDigest(256).run {25        update(name, 0, name.size)26        ByteArray(32).also { doFinal(it, 0) }27    }28}2930private fun mixKKHash(hash: ByteArray, data: ByteArray): ByteArray =31    Blake2sDigest(256).run {32        update(hash, 0, hash.size)33        update(data, 0, data.size)34        ByteArray(32).also { doFinal(it, 0) }35    }3637private fun kkHkdf2(chainingKey: ByteArray, input: ByteArray): Pair<ByteArray, ByteArray> {38    val generator = HKDFBytesGenerator(Blake2sDigest(256))39    generator.init(HKDFParameters(input, chainingKey, ByteArray(0)))40    val output = ByteArray(64)41    generator.generateBytes(output, 0, output.size)42    return output.copyOfRange(0, 32) to output.copyOfRange(32, 64)43}4445private fun kkNonce(value: Long): ByteArray =46    ByteArray(12).also {47        ByteBuffer.wrap(it).order(ByteOrder.LITTLE_ENDIAN).putLong(4, value)48    }4950private fun kkSeal(key: ByteArray, plaintext: ByteArray, aad: ByteArray): ByteArray {51    val cipher = ChaCha20Poly1305()52    cipher.init(true, AEADParameters(KeyParameter(key), 128, ByteArray(12), aad))53    val output = ByteArray(cipher.getOutputSize(plaintext.size))54    var size = cipher.processBytes(plaintext, 0, plaintext.size, output, 0)55    size += cipher.doFinal(output, size)56    return output.copyOf(size)57}5859private fun kkOpen(key: ByteArray, ciphertext: ByteArray, aad: ByteArray): ByteArray? =60    runCatching {61        val cipher = ChaCha20Poly1305()62        cipher.init(false, AEADParameters(KeyParameter(key), 128, ByteArray(12), aad))63        val output = ByteArray(cipher.getOutputSize(ciphertext.size))64        var size = cipher.processBytes(ciphertext, 0, ciphertext.size, output, 0)65        size += cipher.doFinal(output, size)66        output.copyOf(size)67    }.getOrNull()6869private fun freshKKKeyPair(): NoiseKKKeyPair {70    val privateKey = X25519PrivateKeyParameters(SecureRandom())71    return NoiseKKKeyPair(privateKey.generatePublicKey().encoded, privateKey.encoded)72}7374/**75 * Standard Noise KK reconnect initiator.76 * Pre-messages: -> s, <- s; messages: -> e,es,ss then <- e,ee,se.77 */78internal suspend fun noiseKKInitiator(79    ourStaticPrivateKey: ByteArray,80    ourStaticPublicKey: ByteArray,81    responderStaticPublicKey: ByteArray,82    send: suspend (ByteArray) -> Unit,83    receive: suspend () -> ByteArray,84    ephemeral: NoiseKKKeyPair = freshKKKeyPair(),85): NoiseNKSession {86    requireKKKeys(87        ourStaticPrivateKey,88        ourStaticPublicKey,89        responderStaticPublicKey,90        ephemeral.secretKey,91        ephemeral.publicKey,92    )93    var hash = initKKHash()94    var chainingKey = hash.copyOf()95    hash = mixKKHash(hash, ourStaticPublicKey)96    hash = mixKKHash(hash, responderStaticPublicKey)9798    // -> e, es, ss99    hash = mixKKHash(hash, ephemeral.publicKey)100    var keys = kkHkdf2(101        chainingKey,102        x25519SharedKey(ephemeral.secretKey, responderStaticPublicKey),103    )104    chainingKey = keys.first105    keys = kkHkdf2(106        chainingKey,107        x25519SharedKey(ourStaticPrivateKey, responderStaticPublicKey),108    )109    chainingKey = keys.first110    val encrypted1 = kkSeal(keys.second, ByteArray(0), hash)111    hash = mixKKHash(hash, encrypted1)112    send(ephemeral.publicKey + encrypted1)113114    // <- e, ee, se115    val message2 = receive()116    require(message2.size >= 48) { "Noise KK: msg2 too short" }117    val responderEphemeral = message2.copyOfRange(0, 32)118    val encrypted2 = message2.copyOfRange(32, message2.size)119    hash = mixKKHash(hash, responderEphemeral)120    keys = kkHkdf2(121        chainingKey,122        x25519SharedKey(ephemeral.secretKey, responderEphemeral),123    )124    chainingKey = keys.first125    keys = kkHkdf2(126        chainingKey,127        x25519SharedKey(ourStaticPrivateKey, responderEphemeral),128    )129    chainingKey = keys.first130    requireNotNull(kkOpen(keys.second, encrypted2, hash)) {131        "Noise KK: msg2 auth failed — wrong peer key or MITM"132    }133    hash = mixKKHash(hash, encrypted2)134135    val transportKeys = kkHkdf2(chainingKey, ByteArray(0))136    return kkTransportSession(transportKeys.first, transportKeys.second, hash)137}138139/**140 * Standard Noise KK responder. [expectedInitiatorStaticPublicKey] must be141 * selected from local paired state before the handshake; it is never learned142 * from a peer capability or accepted via fallback.143 */144internal suspend fun noiseKKResponder(145    ourStaticPrivateKey: ByteArray,146    ourStaticPublicKey: ByteArray,147    expectedInitiatorStaticPublicKey: ByteArray,148    send: suspend (ByteArray) -> Unit,149    receive: suspend () -> ByteArray,150    ephemeral: NoiseKKKeyPair = freshKKKeyPair(),151): NoiseNKSession {152    requireKKKeys(153        ourStaticPrivateKey,154        ourStaticPublicKey,155        expectedInitiatorStaticPublicKey,156        ephemeral.secretKey,157        ephemeral.publicKey,158    )159    var hash = initKKHash()160    var chainingKey = hash.copyOf()161    hash = mixKKHash(hash, expectedInitiatorStaticPublicKey)162    hash = mixKKHash(hash, ourStaticPublicKey)163164    // -> e, es, ss165    val message1 = receive()166    require(message1.size >= 48) { "Noise KK: msg1 too short" }167    val initiatorEphemeral = message1.copyOfRange(0, 32)168    val encrypted1 = message1.copyOfRange(32, message1.size)169    hash = mixKKHash(hash, initiatorEphemeral)170    var keys = kkHkdf2(171        chainingKey,172        x25519SharedKey(ourStaticPrivateKey, initiatorEphemeral),173    )174    chainingKey = keys.first175    keys = kkHkdf2(176        chainingKey,177        x25519SharedKey(ourStaticPrivateKey, expectedInitiatorStaticPublicKey),178    )179    chainingKey = keys.first180    requireNotNull(kkOpen(keys.second, encrypted1, hash)) {181        "Noise KK: msg1 auth failed — wrong peer key or MITM"182    }183    hash = mixKKHash(hash, encrypted1)184185    // <- e, ee, se186    hash = mixKKHash(hash, ephemeral.publicKey)187    keys = kkHkdf2(188        chainingKey,189        x25519SharedKey(ephemeral.secretKey, initiatorEphemeral),190    )191    chainingKey = keys.first192    keys = kkHkdf2(193        chainingKey,194        x25519SharedKey(ephemeral.secretKey, expectedInitiatorStaticPublicKey),195    )196    chainingKey = keys.first197    val encrypted2 = kkSeal(keys.second, ByteArray(0), hash)198    hash = mixKKHash(hash, encrypted2)199    send(ephemeral.publicKey + encrypted2)200201    val transportKeys = kkHkdf2(chainingKey, ByteArray(0))202    return kkTransportSession(transportKeys.second, transportKeys.first, hash)203}204205private fun requireKKKeys(vararg keys: ByteArray) {206    require(keys.all { it.size == 32 }) { "Noise KK keys must be 32 bytes" }207}208209private fun kkTransportSession(210    sendKey: ByteArray,211    receiveKey: ByteArray,212    hash: ByteArray,213): NoiseNKSession {214    var sendNonce = 0L215    var receiveNonce = 0L216    return object : NoiseNKSession {217        override fun encrypt(plaintext: ByteArray): ByteArray {218            val cipher = ChaCha20Poly1305()219            cipher.init(true, AEADParameters(KeyParameter(sendKey), 128, kkNonce(sendNonce++), ByteArray(0)))220            val output = ByteArray(cipher.getOutputSize(plaintext.size))221            var size = cipher.processBytes(plaintext, 0, plaintext.size, output, 0)222            size += cipher.doFinal(output, size)223            return output.copyOf(size)224        }225226        override fun decrypt(ciphertext: ByteArray): ByteArray {227            val cipher = ChaCha20Poly1305()228            cipher.init(false, AEADParameters(KeyParameter(receiveKey), 128, kkNonce(receiveNonce++), ByteArray(0)))229            val output = ByteArray(cipher.getOutputSize(ciphertext.size))230            var size = cipher.processBytes(ciphertext, 0, ciphertext.size, output, 0)231            size += cipher.doFinal(output, size)232            return output.copyOf(size)233        }234235        override val handshakeHash: ByteArray = hash.copyOf()236    }237}