fix(voice): working decryption

kordlib · Jun 5, 2024 · 999685c · 999685c
1 parent 65de968
commit 999685c
Show file tree

Hide file tree

Showing 5 changed files with 42 additions and 24 deletions.
diff --git a/voice/src/main/kotlin/encryption/AeadAes256Gcm.kt b/voice/src/main/kotlin/encryption/AeadAes256Gcm.kt
@@ -32,17 +32,14 @@ public data object AeadAes256Gcm : VoiceEncryption {
         protected val iv = ByteArray(IV_LEN)
         protected val ivCursor = iv.mutableCursor()
 
-        protected open val nonceBuffer: ByteArray = ByteArray(NONCE_LEN)
-        protected val nonceCursor by lazy { nonceBuffer.mutableCursor() }
-        protected val nonceView by lazy { nonceBuffer.view() }
-
         protected val cipherKey = SecretKeySpec(key, "AES")
         protected val cipher: Cipher = Cipher.getInstance("AES/GCM/NoPadding")
 
         fun apply(
             mode: Int,
             src: ByteArrayView,
             dst: MutableByteArrayCursor,
+            aead: ByteArrayView,
             nonce: ByteArray,
             writeNonce: MutableByteArrayCursor.(nonce: ByteArray) -> Unit,
         ): Boolean {
@@ -51,7 +48,7 @@ public data object AeadAes256Gcm : VoiceEncryption {
             ivCursor.apply { writeNonce(nonce) }
 
             init(mode)
-            cipher.updateAAD(dst.data.copyOfRange(0, dst.cursor))
+            cipher.updateAAD(aead.data, aead.dataStart, aead.viewSize)
             dst.cursor += cipher.doFinal(src.data, src.dataStart, src.viewSize, dst.data, dst.cursor)
 
             return true
@@ -63,11 +60,19 @@ public data object AeadAes256Gcm : VoiceEncryption {
     }
 
     private class BoxImpl(key: ByteArray) : Box, Common(key) {
+        val nonceBuffer: ByteArray = ByteArray(NONCE_LEN)
+        val nonceCursor by lazy { nonceBuffer.mutableCursor() }
+        val nonceView by lazy { nonceBuffer.view() }
+
         override val overhead: Int
             get() = AUTH_TAG_LEN + NONCE_LEN
 
-        override fun apply(src: ByteArrayView, nonce: ByteArray, dst: MutableByteArrayCursor): Boolean =
-            apply(Cipher.ENCRYPT_MODE, src, dst, nonce, MutableByteArrayCursor::writeByteArray)
+        override fun apply(
+            src: ByteArrayView,
+            dst: MutableByteArrayCursor,
+            aead: ByteArrayView,
+            nonce: ByteArray,
+        ): Boolean = apply(Cipher.ENCRYPT_MODE, src, dst, aead, nonce, MutableByteArrayCursor::writeByteArray)
 
         override fun appendNonce(nonce: ByteArrayView, dst: MutableByteArrayCursor) {
             dst.writeByteView(nonce)
@@ -81,17 +86,22 @@ public data object AeadAes256Gcm : VoiceEncryption {
     }
 
     private class UnboxImpl(key: ByteArray) : Unbox, Common(key) {
-        // Since RTPPacket expects a container big enough to fit the entire header this will
-        // need to be 12 bytes (or the min length of an RTP header). We will only use the
-        // first NONCE_LEN bytes, though.
-        override val nonceBuffer: ByteArray = ByteArray(12)
+        override fun apply(
+            src: ByteArrayView,
+            dst: MutableByteArrayCursor,
+            aead: ByteArrayView,
+            nonce: ByteArray,
+        ): Boolean = apply(Cipher.DECRYPT_MODE, src, dst, aead, nonce) { writeByteView(it.view(0, NONCE_LEN)!!) }
 
-        override fun apply(src: ByteArrayView, nonce: ByteArray, dst: MutableByteArrayCursor): Boolean =
-            apply(Cipher.DECRYPT_MODE, src, dst, nonce) { writeByteView(it.view(0, NONCE_LEN)!!) }
+        override fun getNonce(packet: RTPPacket): ByteArrayView = with(packet.payload) {
+            // grab the last NONCE_LEN bytes of the packet payload.
+            val nonce = view(dataEnd - NONCE_LEN, dataEnd)
+                ?: error("Failed to strip nonce from RTP packet payload.")
 
-        override fun getNonce(packet: RTPPacket): ByteArrayView {
-            packet.writeHeader(nonceCursor)
-            return nonceView
+            // resize the payload view to exclude the nonce.
+            resize(0, dataEnd - NONCE_LEN)
+
+            return nonce
         }
     }
 }
diff --git a/voice/src/main/kotlin/encryption/VoiceEncryption.kt b/voice/src/main/kotlin/encryption/VoiceEncryption.kt
@@ -36,7 +36,7 @@ public interface VoiceEncryption {
      */
     public sealed interface Method {
 
-        public fun apply(src: ByteArrayView, nonce: ByteArray, dst: MutableByteArrayCursor): Boolean
+        public fun apply(src: ByteArrayView, dst: MutableByteArrayCursor, aead: ByteArrayView, nonce: ByteArray): Boolean
     }
 
     public interface Box : Method {

diff --git a/voice/src/main/kotlin/encryption/XSalsa20Poly1305.kt b/voice/src/main/kotlin/encryption/XSalsa20Poly1305.kt
@@ -37,7 +37,7 @@ public data class XSalsa20Poly1305(public val nonceStrategyFactory: NonceStrateg
             override val overhead: Int
                 get() = boxzerobytesLength + nonceStrategyFactory.length
 
-            override fun apply(src: ByteArrayView, nonce: ByteArray, dst: MutableByteArrayCursor): Boolean {
+            override fun apply(src: ByteArrayView, dst: MutableByteArrayCursor, aead: ByteArrayView, nonce: ByteArray): Boolean {
                 m.fill(0)
                 c.fill(0)
 
@@ -65,7 +65,7 @@ public data class XSalsa20Poly1305(public val nonceStrategyFactory: NonceStrateg
             private val key: ByteArray,
             nonceStrategyFactory: NonceStrategy.Factory,
         ) : VoiceEncryption.Unbox, Impl(nonceStrategyFactory) {
-            override fun apply(src: ByteArrayView, nonce: ByteArray, dst: MutableByteArrayCursor): Boolean {
+            override fun apply(src: ByteArrayView, dst: MutableByteArrayCursor, aead: ByteArrayView, nonce: ByteArray): Boolean {
                 c.fill(0)
                 m.fill(0)
 

diff --git a/voice/src/main/kotlin/streams/DefaultStreams.kt b/voice/src/main/kotlin/streams/DefaultStreams.kt
@@ -88,18 +88,26 @@ private fun Flow<RTPPacket>.decrypt(encryption: VoiceEncryption, key: ByteArray)
     val unbox = encryption.createUnbox(key)
     val nonceBuffer = ByteArray(encryption.nonceLength).mutableCursor()
 
+    val aeadBuffer = ByteArray(12)
+    val aeadCursor = aeadBuffer.mutableCursor()
+    val aeadView = aeadBuffer.view()
+
     val decryptedBuffer = ByteArray(512)
     val decryptedCursor = decryptedBuffer.mutableCursor()
     val decryptedView = decryptedBuffer.view()
 
     return mapNotNull { packet ->
+        // strip and write the nonce.
         nonceBuffer.reset()
-        decryptedCursor.reset()
-
         nonceBuffer.writeByteView(unbox.getNonce(packet))
 
-        val decrypted = unbox.apply(packet.payload, nonceBuffer.data, decryptedCursor)
-        if (!decrypted) {
+        // write the RTP packet header to the AEAD buffer.
+        aeadCursor.reset()
+        packet.writeHeader(aeadCursor)
+
+        // decrypt the RTP packet payload.
+        decryptedCursor.reset()
+        if (!unbox.apply(packet.payload, decryptedCursor, aeadView, nonceBuffer.data)) {
             defaultStreamsLogger.trace { "failed to decrypt the packet with data ${packet.payload.data.contentToString()} at offset ${packet.payload.dataStart} and length ${packet.payload.viewSize - 4}" }
             return@mapNotNull null
         }

diff --git a/voice/src/main/kotlin/udp/AudioPacketProvider.kt b/voice/src/main/kotlin/udp/AudioPacketProvider.kt
@@ -53,7 +53,7 @@ public class DefaultAudioPacketProvider(key: ByteArray, encryption: VoiceEncrypt
                 nonceBuffer.writeByteView(rawNonce)
 
                 // encrypt data and write into our buffer
-                val encrypted = box.apply(data.view(), nonceBuffer.data, this)
+                val encrypted = box.apply(data.view(), this, rtpHeaderView, nonceBuffer.data)
                 if (!encrypted) throw CouldNotEncryptDataException(data)
 
                 box.appendNonce(rawNonce, this)