Track E2EE key validity on a per index basis (#1284)

* Additional tests for E2EE * Additional tests for E2EE * Track decryption failures per key index * Changeset * Revert "Additional tests for E2EE" This reverts commit 73b7bc2. * Format * Revert from merge * Apply suggestions from code review Co-authored-by: lukasIO <[email protected]> --------- Co-authored-by: lukasIO <[email protected]>
livekit · Oct 14, 2024 · cbcb593 · cbcb593
1 parent 9b13728
commit cbcb593
Show file tree

Hide file tree

Showing 5 changed files with 205 additions and 24 deletions.
diff --git a/.changeset/bright-pots-lick.md b/.changeset/bright-pots-lick.md
@@ -0,0 +1,5 @@
+---
+'livekit-client': patch
+---
+
+Track E2EE key validity on a per index basis
diff --git a/src/e2ee/worker/FrameCryptor.test.ts b/src/e2ee/worker/FrameCryptor.test.ts
@@ -297,6 +297,7 @@ describe('FrameCryptor', () => {
         await vitest.advanceTimersToNextTimerAsync();
 
         expect(keys.decryptionFailure).toHaveBeenCalledTimes(1);
+        expect(keys.decryptionFailure).toHaveBeenCalledWith(1);
         expect(errorListener).toHaveBeenCalled();
       } finally {
         vitest.useRealTimers();
@@ -350,6 +351,7 @@ describe('FrameCryptor', () => {
         expect(output.chunks).toEqual([]);
         expect(errorListener).toHaveBeenCalled();
         expect(keys.decryptionFailure).toHaveBeenCalledTimes(1);
+        expect(keys.decryptionFailure).toHaveBeenCalledWith(1);
       } finally {
         vitest.useRealTimers();
       }
@@ -377,6 +379,7 @@ describe('FrameCryptor', () => {
         expect(output.chunks).toEqual([frame]);
 
         expect(keys.decryptionSuccess).toHaveBeenCalledTimes(1);
+        expect(keys.decryptionSuccess).toHaveBeenCalledWith(1);
 
         expect(frame.data.byteLength).toBe(16);
 
@@ -387,5 +390,75 @@ describe('FrameCryptor', () => {
         vitest.useRealTimers();
       }
     });
+
+    it('recovers from delayed use of rotated key', async () => {
+      vitest.useFakeTimers();
+      try {
+        // 1. we (the local participant) have just joined a room and do not have the existing key (index 0) for the existing/remote participant
+        const { keys, input, output } = prepareParticipantTestDecoder(participantIdentity, {
+          failureTolerance: 1,
+          ratchetWindowSize: 0,
+        });
+        vitest.spyOn(keys, 'decryptionFailure');
+
+        // 2. we receive some frames from the existing participant encrypted with the existing key 0 that we don't have
+        input.write(mockEncryptedRTCEncodedVideoFrame(0));
+        input.write(mockEncryptedRTCEncodedVideoFrame(0));
+        input.write(mockEncryptedRTCEncodedVideoFrame(0));
+        input.write(mockEncryptedRTCEncodedVideoFrame(0));
+
+        // 3. we should have marked key at index 0 as invalid by now and dropped all the frames
+        await vitest.waitFor(() => expect(keys.decryptionFailure).toHaveBeenCalledTimes(2));
+        expect(keys.hasInvalidKeyAtIndex(0)).toBe(true);
+        expect(output.chunks).toEqual([]);
+
+        // 4. the existing participant then notices that we have joined the room and generates a new key (with a new key index 1)
+        // and distributes it out of band to us
+        await keys.setKey(await createKeyMaterialFromString('key1'), 1);
+
+        // 5. the existing participant waits a period of time before using the new key and continues sending media using the previous key 0.
+        // we receive these frames and should drop them as we still don't have the key.
+        input.write(mockEncryptedRTCEncodedVideoFrame(0));
+        input.write(mockEncryptedRTCEncodedVideoFrame(0));
+        input.write(mockEncryptedRTCEncodedVideoFrame(0));
+        input.write(mockEncryptedRTCEncodedVideoFrame(0));
+
+        await vitest.advanceTimersToNextTimerAsync();
+        expect(output.chunks).toEqual([]);
+
+        // 6. the existing participant moves over to the new key index 1 and we start to receive frames for index 1 that we
+        // should be able to decrypt even though we had the previous failures.
+        input.write(
+          mockRTCEncodedVideoFrame(
+            new Uint8Array([
+              1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 254, 96, 91, 111, 187, 132, 31, 12, 207, 136, 17, 221,
+              233, 116, 174, 6, 50, 37, 214, 71, 119, 196, 255, 255, 255, 255, 0, 0, 0, 0, 255, 255,
+              199, 51, 12, 1,
+            ]),
+          ),
+        );
+
+        input.write(
+          mockRTCEncodedVideoFrame(
+            new Uint8Array([
+              99, 2, 3, 4, 5, 6, 7, 8, 9, 10, 154, 108, 209, 239, 253, 33, 72, 111, 13, 125, 10,
+              101, 28, 209, 141, 162, 0, 238, 189, 254, 66, 156, 255, 255, 255, 255, 0, 0, 0, 0,
+              255, 255, 96, 247, 12, 1,
+            ]),
+          ),
+        );
+
+        await vitest.waitFor(() => expect(output.chunks.length).toEqual(2));
+
+        expect(new Uint8Array(output.chunks[0].data)).toEqual(
+          new Uint8Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]),
+        );
+        expect(new Uint8Array(output.chunks[1].data)).toEqual(
+          new Uint8Array([99, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]),
+        );
+      } finally {
+        vitest.useRealTimers();
+      }
+    });
   });
 });
diff --git a/src/e2ee/worker/FrameCryptor.ts b/src/e2ee/worker/FrameCryptor.ts
@@ -357,10 +357,15 @@ export class FrameCryptor extends BaseFrameCryptor {
     const data = new Uint8Array(encodedFrame.data);
     const keyIndex = data[encodedFrame.data.byteLength - 1];
 
-    if (this.keys.getKeySet(keyIndex) && this.keys.hasValidKey) {
+    if (this.keys.hasInvalidKeyAtIndex(keyIndex)) {
+      // drop frame
+      return;
+    }
+
+    if (this.keys.getKeySet(keyIndex)) {
       try {
         const decodedFrame = await this.decryptFrame(encodedFrame, keyIndex);
-        this.keys.decryptionSuccess();
+        this.keys.decryptionSuccess(keyIndex);
         if (decodedFrame) {
           return controller.enqueue(decodedFrame);
         }
@@ -369,13 +374,13 @@ export class FrameCryptor extends BaseFrameCryptor {
           // emit an error if the key handler thinks we have a valid key
           if (this.keys.hasValidKey) {
             this.emit(CryptorEvent.Error, error);
-            this.keys.decryptionFailure();
+            this.keys.decryptionFailure(keyIndex);
           }
         } else {
           workerLogger.warn('decoding frame failed', { error });
         }
       }
-    } else if (!this.keys.getKeySet(keyIndex) && this.keys.hasValidKey) {
+    } else {
       // emit an error if the key index is out of bounds but the key handler thinks we still have a valid key
       workerLogger.warn(`skipping decryption due to missing key at index ${keyIndex}`);
       this.emit(
@@ -386,7 +391,7 @@ export class FrameCryptor extends BaseFrameCryptor {
           this.participantIdentity,
         ),
       );
-      this.keys.decryptionFailure();
+      this.keys.decryptionFailure(keyIndex);
     }
   }
 

diff --git a/src/e2ee/worker/ParticipantKeyHandler.test.ts b/src/e2ee/worker/ParticipantKeyHandler.test.ts
@@ -92,6 +92,7 @@ describe('ParticipantKeyHandler', () => {
     keyHandler.setCurrentKeyIndex(10);
 
     expect(keyHandler.hasValidKey).toBe(true);
+    expect(keyHandler.hasInvalidKeyAtIndex(0)).toBe(false);
 
     keyHandler.decryptionFailure();
 
@@ -102,6 +103,52 @@ describe('ParticipantKeyHandler', () => {
     expect(keyHandler.hasValidKey).toBe(true);
   });
 
+  it('marks specific key invalid if more than failureTolerance failures', async () => {
+    const keyHandler = new ParticipantKeyHandler(participantIdentity, {
+      ...KEY_PROVIDER_DEFAULTS,
+      failureTolerance: 2,
+    });
+
+    // set the current key to something different from what we are testing
+    keyHandler.setCurrentKeyIndex(10);
+
+    expect(keyHandler.hasInvalidKeyAtIndex(5)).toBe(false);
+
+    // 1
+    keyHandler.decryptionFailure(5);
+    expect(keyHandler.hasInvalidKeyAtIndex(5)).toBe(false);
+
+    // 2
+    keyHandler.decryptionFailure(5);
+    expect(keyHandler.hasInvalidKeyAtIndex(5)).toBe(false);
+
+    // 3
+    keyHandler.decryptionFailure(5);
+    expect(keyHandler.hasInvalidKeyAtIndex(5)).toBe(true);
+
+    expect(keyHandler.hasInvalidKeyAtIndex(10)).toBe(false);
+  });
+
+  it('marks specific key valid on encryption success', async () => {
+    const keyHandler = new ParticipantKeyHandler(participantIdentity, {
+      ...KEY_PROVIDER_DEFAULTS,
+      failureTolerance: 0,
+    });
+
+    // set the current key to something different from what we are testing
+    keyHandler.setCurrentKeyIndex(10);
+
+    expect(keyHandler.hasInvalidKeyAtIndex(5)).toBe(false);
+
+    keyHandler.decryptionFailure(5);
+
+    expect(keyHandler.hasInvalidKeyAtIndex(5)).toBe(true);
+
+    keyHandler.decryptionSuccess(5);
+
+    expect(keyHandler.hasInvalidKeyAtIndex(5)).toBe(false);
+  });
+
   it('marks valid on new key', async () => {
     const keyHandler = new ParticipantKeyHandler(participantIdentity, {
       ...KEY_PROVIDER_DEFAULTS,
@@ -111,6 +158,7 @@ describe('ParticipantKeyHandler', () => {
     keyHandler.setCurrentKeyIndex(10);
 
     expect(keyHandler.hasValidKey).toBe(true);
+    expect(keyHandler.hasInvalidKeyAtIndex(0)).toBe(false);
 
     keyHandler.decryptionFailure();
 
@@ -159,6 +207,26 @@ describe('ParticipantKeyHandler', () => {
     }
   });
 
+  describe('resetKeyStatus', () => {
+    it('marks all keys as valid if no index is provided', () => {
+      const keyHandler = new ParticipantKeyHandler(participantIdentity, {
+        ...KEY_PROVIDER_DEFAULTS,
+        failureTolerance: 0,
+      });
+
+      for (let i = 0; i < KEY_PROVIDER_DEFAULTS.keyringSize; i++) {
+        keyHandler.decryptionFailure(i);
+        expect(keyHandler.hasInvalidKeyAtIndex(i)).toBe(true);
+      }
+
+      keyHandler.resetKeyStatus();
+
+      for (let i = 0; i < KEY_PROVIDER_DEFAULTS.keyringSize; i++) {
+        expect(keyHandler.hasInvalidKeyAtIndex(i)).toBe(false);
+      }
+    });
+  });
+
   describe('ratchetKey', () => {
     it('emits event', async () => {
       const keyHandler = new ParticipantKeyHandler(participantIdentity, KEY_PROVIDER_DEFAULTS);

diff --git a/src/e2ee/worker/ParticipantKeyHandler.ts b/src/e2ee/worker/ParticipantKeyHandler.ts
@@ -21,18 +21,19 @@ export class ParticipantKeyHandler extends (EventEmitter as new () => TypedEvent
 
   private cryptoKeyRing: Array<KeySet | undefined>;
 
+  private decryptionFailureCounts: Array<number>;
+
   private keyProviderOptions: KeyProviderOptions;
 
   private ratchetPromiseMap: Map<number, Promise<CryptoKey>>;
 
   private participantIdentity: string;
 
-  private decryptionFailureCount = 0;
-
-  private _hasValidKey: boolean = true;
-
-  get hasValidKey() {
-    return this._hasValidKey;
+  /**
+   * true if the current key has not been marked as invalid
+   */
+  get hasValidKey(): boolean {
+    return !this.hasInvalidKeyAtIndex(this.currentKeyIndex);
   }
 
   constructor(participantIdentity: string, keyProviderOptions: KeyProviderOptions) {
@@ -42,35 +43,64 @@ export class ParticipantKeyHandler extends (EventEmitter as new () => TypedEvent
       throw new TypeError('Keyring size needs to be between 1 and 256');
     }
     this.cryptoKeyRing = new Array(keyProviderOptions.keyringSize).fill(undefined);
+    this.decryptionFailureCounts = new Array(keyProviderOptions.keyringSize).fill(0);
     this.keyProviderOptions = keyProviderOptions;
     this.ratchetPromiseMap = new Map();
     this.participantIdentity = participantIdentity;
-    this.resetKeyStatus();
   }
 
-  decryptionFailure() {
+  /**
+   * Returns true if the key at the given index is marked as invalid.
+   *
+   * @param keyIndex the index of the key
+   */
+  hasInvalidKeyAtIndex(keyIndex: number): boolean {
+    return (
+      this.keyProviderOptions.failureTolerance >= 0 &&
+      this.decryptionFailureCounts[keyIndex] > this.keyProviderOptions.failureTolerance
+    );
+  }
+
+  /**
+   * Informs the key handler that a decryption failure occurred for an encryption key.
+   * @internal
+   * @param keyIndex the key index for which the failure occurred. Defaults to the current key index.
+   */
+  decryptionFailure(keyIndex: number = this.currentKeyIndex): void {
     if (this.keyProviderOptions.failureTolerance < 0) {
       return;
     }
-    this.decryptionFailureCount += 1;
 
-    if (this.decryptionFailureCount > this.keyProviderOptions.failureTolerance) {
-      workerLogger.warn(`key for ${this.participantIdentity} is being marked as invalid`);
-      this._hasValidKey = false;
+    this.decryptionFailureCounts[keyIndex] += 1;
+
+    if (this.decryptionFailureCounts[keyIndex] > this.keyProviderOptions.failureTolerance) {
+      workerLogger.warn(
+        `key for ${this.participantIdentity} at index ${keyIndex} is being marked as invalid`,
+      );
     }
   }
 
-  decryptionSuccess() {
-    this.resetKeyStatus();
+  /**
+   * Informs the key handler that a frame was successfully decrypted using an encryption key.
+   * @internal
+   * @param keyIndex the key index for which the success occurred. Defaults to the current key index.
+   */
+  decryptionSuccess(keyIndex: number = this.currentKeyIndex): void {
+    this.resetKeyStatus(keyIndex);
   }
 
   /**
    * Call this after user initiated ratchet or a new key has been set in order to make sure to mark potentially
    * invalid keys as valid again
+   *
+   * @param keyIndex the index of the key. Defaults to the current key index.
    */
-  resetKeyStatus() {
-    this.decryptionFailureCount = 0;
-    this._hasValidKey = true;
+  resetKeyStatus(keyIndex?: number): void {
+    if (keyIndex === undefined) {
+      this.decryptionFailureCounts.fill(0);
+    } else {
+      this.decryptionFailureCounts[keyIndex] = 0;
+    }
   }
 
   /**
@@ -130,7 +160,7 @@ export class ParticipantKeyHandler extends (EventEmitter as new () => TypedEvent
    */
   async setKey(material: CryptoKey, keyIndex = 0) {
     await this.setKeyFromMaterial(material, keyIndex);
-    this.resetKeyStatus();
+    this.resetKeyStatus(keyIndex);
   }
 
   /**
@@ -161,7 +191,7 @@ export class ParticipantKeyHandler extends (EventEmitter as new () => TypedEvent
 
   async setCurrentKeyIndex(index: number) {
     this.currentKeyIndex = index % this.cryptoKeyRing.length;
-    this.resetKeyStatus();
+    this.resetKeyStatus(index);
   }
 
   getCurrentKeyIndex() {