Skip to content

Commit

Permalink
implement all encryption modes documented by discord
Browse files Browse the repository at this point in the history
  • Loading branch information
lost-illusi0n committed Oct 18, 2021
1 parent 85943df commit a9fc26a
Show file tree
Hide file tree
Showing 12 changed files with 190 additions and 40 deletions.
4 changes: 4 additions & 0 deletions voice/src/main/kotlin/VoiceConnection.kt
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import dev.kord.common.annotation.KordVoice
import dev.kord.common.entity.Snowflake
import dev.kord.gateway.Gateway
import dev.kord.gateway.UpdateVoiceStatus
import dev.kord.voice.encryption.strategies.NonceStrategy
import dev.kord.voice.gateway.VoiceGateway
import dev.kord.voice.gateway.VoiceGatewayConfiguration
import dev.kord.voice.handlers.StreamsHandler
Expand Down Expand Up @@ -40,6 +41,8 @@ data class VoiceConnectionData(
* @param data the data representing this [VoiceConnection].
* @param voiceGatewayConfiguration the configuration used on each new [connect] for the [voiceGateway].
* @param audioProvider a [AudioProvider] that will provide [AudioFrame] when required.
* @param frameSender the [AudioFrameSender] that will handle the sending of audio packets.
* @param nonceStrategy the [NonceStrategy] that is used during encryption of audio.
* @param frameInterceptorFactory a factory for [FrameInterceptor]s that is used whenever audio is ready to be sent. See [FrameInterceptor] and [DefaultFrameInterceptor].
*/
@KordVoice
Expand All @@ -52,6 +55,7 @@ class VoiceConnection(
val streams: Streams,
val audioProvider: AudioProvider,
val frameSender: AudioFrameSender,
val nonceStrategy: NonceStrategy,
val frameInterceptorFactory: (FrameInterceptorContext) -> FrameInterceptor,
) {
val scope: CoroutineScope =
Expand Down
12 changes: 11 additions & 1 deletion voice/src/main/kotlin/VoiceConnectionBuilder.kt
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ import dev.kord.gateway.Gateway
import dev.kord.gateway.UpdateVoiceStatus
import dev.kord.gateway.VoiceServerUpdate
import dev.kord.gateway.VoiceStateUpdate
import dev.kord.voice.encryption.strategies.LiteNonceStrategy
import dev.kord.voice.encryption.strategies.NonceStrategy
import dev.kord.voice.exception.VoiceConnectionInitializationException
import dev.kord.voice.gateway.DefaultVoiceGatewayBuilder
import dev.kord.voice.gateway.VoiceGateway
Expand Down Expand Up @@ -46,6 +48,12 @@ class VoiceConnectionBuilder(
*/
var audioSender: AudioFrameSender? = null

/**
* The nonce strategy to be used for the encryption of audio packets.
* If `null`, [dev.kord.voice.encryption.strategies.LiteNonceStrategy] will be used.
*/
var nonceStrategy: NonceStrategy? = null

fun audioProvider(provider: AudioProvider) {
this.audioProvider = provider
}
Expand Down Expand Up @@ -154,9 +162,10 @@ class VoiceConnectionBuilder(
val audioProvider = audioProvider ?: EmptyAudioPlayerProvider
val audioSender =
audioSender ?: DefaultAudioFrameSender(DefaultAudioFrameSenderData(udpSocket))
val nonceStrategy = nonceStrategy ?: LiteNonceStrategy()
val frameInterceptorFactory = frameInterceptorFactory ?: { DefaultFrameInterceptor(it) }
val streams =
streams ?: if (receiveVoice) DefaultStreams(voiceGateway, udpSocket) else NOPStreams
streams ?: if (receiveVoice) DefaultStreams(voiceGateway, udpSocket, nonceStrategy) else NOPStreams

return VoiceConnection(
voiceConnectionData,
Expand All @@ -167,6 +176,7 @@ class VoiceConnectionBuilder(
streams,
audioProvider,
audioSender,
nonceStrategy,
frameInterceptorFactory,
)
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,9 +23,6 @@ internal class XSalsa20Poly1305Encryption(private val key: ByteArray) {

output.writeByteArray(c, boxzerobytesLength, messageBufferLength - boxzerobytesLength)

// TODO: implement a nonce strategy
output.writeByteArray(nonce, length = 4)

return true
}

Expand Down
36 changes: 36 additions & 0 deletions voice/src/main/kotlin/encryption/strategies/LiteNonceStrategy.kt
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
package dev.kord.voice.encryption.strategies

import dev.kord.voice.io.ByteArrayView
import dev.kord.voice.io.MutableByteArrayCursor
import dev.kord.voice.io.mutableCursor
import dev.kord.voice.io.view
import dev.kord.voice.udp.RTPPacket
import kotlinx.atomicfu.atomic

class LiteNonceStrategy : NonceStrategy {
override val nonceLength: Int = 4

private var count: Int by atomic(0)
private val nonceBuffer: ByteArray = ByteArray(4)
private val nonceView = nonceBuffer.view()
private val nonceCursor = nonceBuffer.mutableCursor()

override fun strip(packet: RTPPacket): ByteArrayView {
return with(packet.payload) {
val nonce = view(dataEnd - 4, dataEnd)!!
resize(dataStart, dataEnd - 4)
nonce
}
}

override fun generate(header: () -> ByteArrayView): ByteArrayView {
count++
nonceCursor.reset()
nonceCursor.writeInt(count)
return nonceView
}

override fun append(nonce: ByteArrayView, cursor: MutableByteArrayCursor) {
cursor.writeByteView(nonce)
}
}
30 changes: 30 additions & 0 deletions voice/src/main/kotlin/encryption/strategies/NonceStrategy.kt
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
package dev.kord.voice.encryption.strategies

import dev.kord.voice.io.ByteArrayView
import dev.kord.voice.io.MutableByteArrayCursor
import dev.kord.voice.udp.RTPPacket

/**
* An [encryption mode, regarding the nonce](https://discord.com/developers/docs/topics/voice-connections#establishing-a-voice-udp-connection-encryption-modes), supported by Discord.
*/
sealed interface NonceStrategy {
/**
* The amount of bytes this nonce will take up.
*/
val nonceLength: Int

/**
* Reads the nonce from this [packet] (also removes it if it resides in the payload), and returns a [ByteArrayView] of it.
*/
fun strip(packet: RTPPacket): ByteArrayView

/**
* Generates a nonce, may use the provided information.
*/
fun generate(header: () -> ByteArrayView): ByteArrayView

/**
* Writes the [nonce] to [cursor] in the correct relative position.
*/
fun append(nonce: ByteArrayView, cursor: MutableByteArrayCursor)
}
31 changes: 31 additions & 0 deletions voice/src/main/kotlin/encryption/strategies/NormalNonceStrategy.kt
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
package dev.kord.voice.encryption.strategies

import dev.kord.voice.io.ByteArrayView
import dev.kord.voice.io.MutableByteArrayCursor
import dev.kord.voice.io.mutableCursor
import dev.kord.voice.io.view
import dev.kord.voice.udp.RTPPacket
import dev.kord.voice.udp.RTP_HEADER_LENGTH

class NormalNonceStrategy : NonceStrategy {
// the nonce is already a part of the rtp header, which means this will take up no extra space.
override val nonceLength: Int = 0

private val rtpHeaderBuffer = ByteArray(RTP_HEADER_LENGTH)
private val rtpHeaderCursor = rtpHeaderBuffer.mutableCursor()
private val rtpHeaderView = rtpHeaderBuffer.view()

override fun strip(packet: RTPPacket): ByteArrayView {
rtpHeaderCursor.reset()
packet.writeHeader(rtpHeaderCursor)
return rtpHeaderView
}

override fun generate(header: () -> ByteArrayView): ByteArrayView {
return header()
}

override fun append(nonce: ByteArrayView, cursor: MutableByteArrayCursor) {
/* the nonce is the rtp header which means this should do nothing */
}
}
33 changes: 33 additions & 0 deletions voice/src/main/kotlin/encryption/strategies/SuffixNonceStrategy.kt
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
package dev.kord.voice.encryption.strategies

import dev.kord.voice.io.ByteArrayView
import dev.kord.voice.io.MutableByteArrayCursor
import dev.kord.voice.io.view
import dev.kord.voice.udp.RTPPacket
import kotlin.random.Random

private const val SUFFIX_NONCE_LENGTH = 24

class SuffixNonceStrategy : NonceStrategy {
override val nonceLength: Int = SUFFIX_NONCE_LENGTH

private val nonceBuffer: ByteArray = ByteArray(SUFFIX_NONCE_LENGTH)
private val nonceView = nonceBuffer.view()

override fun strip(packet: RTPPacket): ByteArrayView {
return with(packet.payload) {
val nonce = view(dataEnd - SUFFIX_NONCE_LENGTH, dataEnd)!!
resize(dataStart, dataEnd - SUFFIX_NONCE_LENGTH)
nonce
}
}

override fun generate(header: () -> ByteArrayView): ByteArrayView {
Random.Default.nextBytes(nonceBuffer)
return nonceView
}

override fun append(nonce: ByteArrayView, cursor: MutableByteArrayCursor) {
cursor.writeByteView(nonce)
}
}
12 changes: 11 additions & 1 deletion voice/src/main/kotlin/handlers/UdpLifeCycleHandler.kt
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,9 @@ import dev.kord.common.annotation.KordVoice
import dev.kord.voice.EncryptionMode
import dev.kord.voice.FrameInterceptorContextBuilder
import dev.kord.voice.VoiceConnection
import dev.kord.voice.encryption.strategies.LiteNonceStrategy
import dev.kord.voice.encryption.strategies.NormalNonceStrategy
import dev.kord.voice.encryption.strategies.SuffixNonceStrategy
import dev.kord.voice.gateway.*
import dev.kord.voice.udp.AudioFrameSenderConfiguration
import io.ktor.util.network.*
Expand Down Expand Up @@ -38,12 +41,18 @@ internal class UdpLifeCycleHandler(

udpLifeCycleLogger.trace { "ip discovered for voice successfully" }

val encryptionMode = when (connection.nonceStrategy) {
is LiteNonceStrategy -> EncryptionMode.XSalsa20Poly1305Lite
is NormalNonceStrategy -> EncryptionMode.XSalsa20Poly1305
is SuffixNonceStrategy -> EncryptionMode.XSalsa20Poly1305Suffix
}

val selectProtocol = SelectProtocol(
protocol = "udp",
data = SelectProtocol.Data(
address = ip.hostname,
port = ip.port,
mode = EncryptionMode.XSalsa20Poly1305Lite
mode = encryptionMode
)
)

Expand All @@ -55,6 +64,7 @@ internal class UdpLifeCycleHandler(
val config = AudioFrameSenderConfiguration(
ssrc = ssrc!!,
key = it.secretKey.toUByteArray().toByteArray(),
nonceStrategy = nonceStrategy,
provider = audioProvider,
baseFrameInterceptorContext = FrameInterceptorContextBuilder(gateway, voiceGateway),
interceptorFactory = frameInterceptorFactory,
Expand Down
15 changes: 10 additions & 5 deletions voice/src/main/kotlin/streams/DefaultStreams.kt
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import dev.kord.common.annotation.KordVoice
import dev.kord.common.entity.Snowflake
import dev.kord.voice.AudioFrame
import dev.kord.voice.encryption.XSalsa20Poly1305Codec
import dev.kord.voice.encryption.strategies.NonceStrategy
import dev.kord.voice.gateway.Speaking
import dev.kord.voice.gateway.VoiceGateway
import dev.kord.voice.io.*
Expand All @@ -27,13 +28,14 @@ private val defaultStreamsLogger = KotlinLogging.logger { }
class DefaultStreams(
private val voiceGateway: VoiceGateway,
private val udp: VoiceUdpSocket,
private val nonceStrategy: NonceStrategy
) : Streams {
private fun CoroutineScope.listenForIncoming(key: ByteArray, server: NetworkAddress) {
udp.incoming
.filter { it.address == server }
.mapNotNull { RTPPacket.fromPacket(it.packet) }
.filter { it.payloadType == PayloadType.Audio.raw }
.decrypt(key)
.decrypt(nonceStrategy, key)
.clean()
.onEach { _incomingAudioPackets.emit(it) }
.launchIn(this)
Expand Down Expand Up @@ -84,7 +86,7 @@ class DefaultStreams(
override val ssrcToUser: Map<UInt, Snowflake> by _ssrcToUser
}

private fun Flow<RTPPacket>.decrypt(key: ByteArray): Flow<RTPPacket> {
private fun Flow<RTPPacket>.decrypt(nonceStrategy: NonceStrategy, key: ByteArray): Flow<RTPPacket> {
val codec = XSalsa20Poly1305Codec(key)
val nonceBuffer = ByteArray(TweetNaclFast.SecretBox.nonceLength).mutableCursor()

Expand All @@ -94,10 +96,13 @@ private fun Flow<RTPPacket>.decrypt(key: ByteArray): Flow<RTPPacket> {

return mapNotNull {
nonceBuffer.reset()
nonceBuffer.writeByteArray(it.payload.data, it.payload.dataEnd - 4, 4)

decryptedCursor.reset()
val decrypted = with(it.payload) { codec.decrypt(data, dataStart, viewSize - 4, nonceBuffer.data, decryptedCursor) }

nonceBuffer.writeByteView(nonceStrategy.strip(it))

val decrypted = with(it.payload) {
codec.decrypt(data, dataStart, viewSize, nonceBuffer.data, decryptedCursor)
}

if (!decrypted) {
defaultStreamsLogger.trace { "failed to decrypt the packet with data ${it.payload.data.contentToString()} at offset ${it.payload.dataStart} and length ${it.payload.viewSize - 4}" }
Expand Down
2 changes: 2 additions & 0 deletions voice/src/main/kotlin/udp/AudioFrameSender.kt
Original file line number Diff line number Diff line change
Expand Up @@ -7,13 +7,15 @@ import dev.kord.voice.AudioProvider
import dev.kord.voice.FrameInterceptor
import dev.kord.voice.FrameInterceptorContext
import dev.kord.voice.FrameInterceptorContextBuilder
import dev.kord.voice.encryption.strategies.NonceStrategy
import io.ktor.util.network.*

@KordVoice
data class AudioFrameSenderConfiguration(
val server: NetworkAddress,
val ssrc: UInt,
val key: ByteArray,
val nonceStrategy: NonceStrategy,
val provider: AudioProvider,
val baseFrameInterceptorContext: FrameInterceptorContextBuilder,
val interceptorFactory: (FrameInterceptorContext) -> FrameInterceptor
Expand Down
Loading

0 comments on commit a9fc26a

Please sign in to comment.