Improvements for E2EE. (#276)

* Improvements for E2EE.

* format.

* spotlessApply.

* Update livekit-android-sdk/src/main/java/io/livekit/android/e2ee/E2EEManager.kt

Co-authored-by: davidliu <davidliu@deviange.net>

* Update livekit-android-sdk/src/main/java/io/livekit/android/e2ee/E2EEManager.kt

Co-authored-by: davidliu <davidliu@deviange.net>

* Update livekit-android-sdk/src/main/java/io/livekit/android/e2ee/E2EEManager.kt

Co-authored-by: davidliu <davidliu@deviange.net>

* Update livekit-android-sdk/src/main/java/io/livekit/android/room/Room.kt

Co-authored-by: davidliu <davidliu@deviange.net>

* Update livekit-android-sdk/src/main/java/io/livekit/android/room/Room.kt

Co-authored-by: davidliu <davidliu@deviange.net>

* fix.

* fix.

* spotlessKotlinCheck.

* fix

---------

Co-authored-by: davidliu <davidliu@deviange.net>

@@ -37,7 +37,7 @@ class E2EEManager
 constructor(keyProvider: KeyProvider) {
     private var room: Room? = null
     private var keyProvider: KeyProvider
-    private var frameCryptors = mutableMapOf<String, FrameCryptor>()
+    private var frameCryptors = mutableMapOf<Pair<String, String>, FrameCryptor>()
     private var algorithm: FrameCryptorAlgorithm = FrameCryptorAlgorithm.AES_GCM
     private lateinit var emitEvent: (roomEvent: RoomEvent) -> Unit?
     var enabled: Boolean = false
@@ -45,6 +45,10 @@ constructor(keyProvider: KeyProvider) {
         this.keyProvider = keyProvider
     }
 
+    public fun keyProvider(): KeyProvider {
+        return this.keyProvider
+    }
+
     suspend fun setup(room: Room, emitEvent: (roomEvent: RoomEvent) -> Unit) {
         if (this.room != room) {
             // E2EEManager already setup, clean up first
@@ -72,8 +76,6 @@ constructor(keyProvider: KeyProvider) {
     }
 
     public fun addSubscribedTrack(track: Track, publication: TrackPublication, participant: RemoteParticipant, room: Room) {
-        var trackId = publication.sid
-        var participantId = participant.sid
         var rtpReceiver: RtpReceiver? = when (publication.track!!) {
             is RemoteAudioTrack -> (publication.track!! as RemoteAudioTrack).receiver
             is RemoteVideoTrack -> (publication.track!! as RemoteVideoTrack).receiver
@@ -81,7 +83,7 @@ constructor(keyProvider: KeyProvider) {
                 throw IllegalArgumentException("unsupported track type")
             }
         }
-        var frameCryptor = addRtpReceiver(rtpReceiver!!, participantId, trackId, publication.track!!.kind.name.lowercase())
+        var frameCryptor = addRtpReceiver(rtpReceiver!!, participant.identity!!, publication.sid, publication.track!!.kind.name.lowercase())
         frameCryptor.setObserver { trackId, state ->
             LKLog.i { "Receiver::onFrameCryptionStateChanged: $trackId, state:  $state" }
             emitEvent(
@@ -96,9 +98,18 @@ constructor(keyProvider: KeyProvider) {
         }
     }
 
-    public fun addPublishedTrack(track: Track, publication: TrackPublication, participant: LocalParticipant, room: Room) {
+    public fun removeSubscribedTrack(track: Track, publication: TrackPublication, participant: RemoteParticipant, room: Room) {
         var trackId = publication.sid
-        var participantId = participant.sid
+        var participantId = participant.identity
+        var frameCryptor = frameCryptors.get(trackId to participantId)
+        if (frameCryptor != null) {
+            frameCryptor.isEnabled = false
+            frameCryptor.dispose()
+            frameCryptors.remove(trackId to participantId)
+        }
+    }
+
+    public fun addPublishedTrack(track: Track, publication: TrackPublication, participant: LocalParticipant, room: Room) {
         var rtpSender: RtpSender? = when (publication.track!!) {
             is LocalAudioTrack -> (publication.track!! as LocalAudioTrack)?.sender
             is LocalVideoTrack -> (publication.track!! as LocalVideoTrack)?.sender
@@ -107,7 +118,7 @@ constructor(keyProvider: KeyProvider) {
             }
         } ?: throw IllegalArgumentException("rtpSender is null")
 
-        var frameCryptor = addRtpSender(rtpSender!!, participantId, trackId, publication.track!!.kind.name.lowercase())
+        var frameCryptor = addRtpSender(rtpSender!!, participant.identity!!, publication.sid, publication.track!!.kind.name.lowercase())
         frameCryptor.setObserver { trackId, state ->
             LKLog.i { "Sender::onFrameCryptionStateChanged: $trackId, state:  $state" }
             emitEvent(
@@ -122,6 +133,17 @@ constructor(keyProvider: KeyProvider) {
         }
     }
 
+    public fun removePublishedTrack(track: Track, publication: TrackPublication, participant: LocalParticipant, room: Room) {
+        var trackId = publication.sid
+        var participantId = participant.identity
+        var frameCryptor = frameCryptors.get(trackId to participantId)
+        if (frameCryptor != null) {
+            frameCryptor.isEnabled = false
+            frameCryptor.dispose()
+            frameCryptors.remove(trackId to participantId)
+        }
+    }
+
     private fun e2eeStateFromFrameCryptoState(state: FrameCryptionState?): E2EEState {
         return when (state) {
             FrameCryptionState.NEW -> E2EEState.NEW
@@ -136,41 +158,28 @@ constructor(keyProvider: KeyProvider) {
     }
 
     private fun addRtpSender(sender: RtpSender, participantId: String, trackId: String, kind: String): FrameCryptor {
-        var pid = "$kind-sender-$participantId-$trackId"
         var frameCryptor = FrameCryptorFactory.createFrameCryptorForRtpSender(
             sender,
-            pid,
+            participantId,
             algorithm,
             keyProvider.rtcKeyProvider,
         )
 
-        frameCryptors[trackId] = frameCryptor
+        frameCryptors[trackId to participantId] = frameCryptor
         frameCryptor.setEnabled(enabled)
-        if (keyProvider.enableSharedKey) {
-            keyProvider.rtcKeyProvider?.setKey(pid, 0, keyProvider?.sharedKey)
-            frameCryptor.setKeyIndex(0)
-        }
-
         return frameCryptor
     }
 
     private fun addRtpReceiver(receiver: RtpReceiver, participantId: String, trackId: String, kind: String): FrameCryptor {
-        var pid = "$kind-receiver-$participantId-$trackId"
         var frameCryptor = FrameCryptorFactory.createFrameCryptorForRtpReceiver(
             receiver,
-            pid,
+            participantId,
             algorithm,
             keyProvider.rtcKeyProvider,
         )
 
-        frameCryptors[trackId] = frameCryptor
+        frameCryptors[trackId to participantId] = frameCryptor
         frameCryptor.setEnabled(enabled)
-
-        if (keyProvider.enableSharedKey) {
-            keyProvider.rtcKeyProvider?.setKey(pid, 0, keyProvider?.sharedKey)
-            frameCryptor.setKeyIndex(0)
-        }
-
         return frameCryptor
     }
 
@@ -182,12 +191,7 @@ constructor(keyProvider: KeyProvider) {
         this.enabled = enabled
         for (item in frameCryptors.entries) {
             var frameCryptor = item.value
-            var participantId = item.key
             frameCryptor.setEnabled(enabled)
-            if (keyProvider.enableSharedKey) {
-                keyProvider.rtcKeyProvider?.setKey(participantId, 0, keyProvider?.sharedKey)
-                frameCryptor.setKeyIndex(0)
-            }
         }
     }
 
@@ -195,10 +199,8 @@ constructor(keyProvider: KeyProvider) {
      * Ratchet key for local participant
      */
     fun ratchetKey() {
-        for (participantId in frameCryptors.keys) {
-            var newKey = keyProvider.rtcKeyProvider?.ratchetKey(participantId, 0)
-            LKLog.d { "ratchetKey: newKey: $newKey" }
-        }
+        var newKey = keyProvider.ratchetSharedKey()
+        LKLog.d { "ratchetSharedKey: newKey: $newKey" }
     }
 
     fun cleanUp() {

@@ -28,13 +28,16 @@ constructor(var participantId: String, var keyIndex: Int, var key: String) {
 }
 
 public interface KeyProvider {
+    fun setSharedKey(key: String, keyIndex: Int? = 0): Boolean
+    fun ratchetSharedKey(keyIndex: Int? = 0): ByteArray
+    fun exportSharedKey(keyIndex: Int? = 0): ByteArray
     fun setKey(key: String, participantId: String?, keyIndex: Int? = 0)
-    fun ratchetKey(participantId: String, index: Int): ByteArray
+    fun ratchetKey(participantId: String, keyIndex: Int? = 0): ByteArray
+    fun exportKey(participantId: String, keyIndex: Int? = 0): ByteArray
+    fun setSifTrailer(trailer: ByteArray)
 
     val rtcKeyProvider: FrameCryptorKeyProvider
 
-    var sharedKey: ByteArray?
-
     var enableSharedKey: Boolean
 }
 
@@ -47,8 +50,18 @@ constructor(
     private var failureTolerance: Int,
 ) :
     KeyProvider {
-    override var sharedKey: ByteArray? = null
     private var keys: MutableMap<String, MutableMap<Int, String>> = mutableMapOf()
+    override fun setSharedKey(key: String, keyIndex: Int?): Boolean {
+        return rtcKeyProvider.setSharedKey(keyIndex ?: 0, key.toByteArray())
+    }
+
+    override fun ratchetSharedKey(keyIndex: Int?): ByteArray {
+        return rtcKeyProvider.ratchetSharedKey(keyIndex ?: 0)
+    }
+
+    override fun exportSharedKey(keyIndex: Int?): ByteArray {
+        return rtcKeyProvider.exportSharedKey(keyIndex ?: 0)
+    }
 
     /**
      * Set a key for a participant
@@ -58,7 +71,6 @@ constructor(
      */
     override fun setKey(key: String, participantId: String?, keyIndex: Int?) {
         if (enableSharedKey) {
-            sharedKey = key.toByteArray()
             return
         }
 
@@ -76,8 +88,16 @@ constructor(
         rtcKeyProvider.setKey(participantId, keyInfo.keyIndex, key.toByteArray())
     }
 
-    override fun ratchetKey(participantId: String, index: Int): ByteArray {
-        return rtcKeyProvider.ratchetKey(participantId, index)
+    override fun ratchetKey(participantId: String, keyIndex: Int?): ByteArray {
+        return rtcKeyProvider.ratchetKey(participantId, keyIndex ?: 0)
+    }
+
+    override fun exportKey(participantId: String, keyIndex: Int?): ByteArray {
+        return rtcKeyProvider.exportKey(participantId, keyIndex ?: 0)
+    }
+
+    override fun setSifTrailer(trailer: ByteArray) {
+        rtcKeyProvider.setSifTrailer(trailer)
     }
 
     override val rtcKeyProvider: FrameCryptorKeyProvider

@@ -322,6 +322,10 @@ constructor(
         name = response.room.name
         metadata = response.room.metadata
 
+        if (e2eeManager != null && !response.sifTrailer.isEmpty) {
+            e2eeManager!!.keyProvider().setSifTrailer(response.sifTrailer.toByteArray())
+        }
+
         if (response.room.activeRecording != isRecording) {
             isRecording = response.room.activeRecording
             eventBus.postEvent(RoomEvent.RecordingStatusChanged(this, isRecording), coroutineScope)
@@ -938,6 +942,9 @@ constructor(
      */
     override fun onTrackUnpublished(publication: LocalTrackPublication, participant: LocalParticipant) {
         listener?.onTrackUnpublished(publication, participant, this)
+        e2eeManager?.let { e2eeManager ->
+            e2eeManager!!.removePublishedTrack(publication.track!!, publication, participant, this)
+        }
         eventBus.postEvent(RoomEvent.TrackUnpublished(this, publication, participant), coroutineScope)
     }
 
@@ -973,6 +980,9 @@ constructor(
         participant: RemoteParticipant,
     ) {
         listener?.onTrackUnsubscribed(track, publication, participant, this)
+        e2eeManager?.let { e2eeManager ->
+            e2eeManager!!.removeSubscribedTrack(track, publication, participant, this)
+        }
         eventBus.postEvent(RoomEvent.TrackUnsubscribed(this, track, publication, participant), coroutineScope)
     }
 

+/*
+ * Copyright 2023 LiveKit, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package io.livekit.android.sample
 
 import android.app.Application
@@ -47,7 +63,7 @@ class CallViewModel(
         if (e2ee && e2eeKey != null) {
             e2eeOptions = E2EEOptions()
         }
-        e2eeOptions?.keyProvider?.setKey(e2eeKey!!, null, 0)
+        e2eeOptions?.keyProvider?.setSharedKey(e2eeKey!!)
         return e2eeOptions
     }