livekit-client/src/room/data-track/outgoing/OutgoingDataTrackManager.test.ts

JavaScript/TypeScript client SDK for LiveKit

/* eslint-disable @typescript-eslint/no-unused-vars */
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import {
  type DecryptDataResponseMessage,
  type EncryptDataResponseMessage,
  LocalDataTrack,
} from '../../..';
import { type BaseE2EEManager } from '../../../e2ee/E2eeManager';
import { subscribeToEvents } from '../../../utils/subscribeToEvents';
import RTCEngine from '../../RTCEngine';
import Room from '../../Room';
import { DataTrackHandle } from '../handle';
import { DataTrackPacket, FrameMarker } from '../packet';
import OutgoingDataTrackManager, {
  type DataTrackOutgoingManagerCallbacks,
  Descriptor,
} from './OutgoingDataTrackManager';
import { DataTrackPublishError } from './errors';

/** Fake encryption provider for testing e2ee data track features. */
export class PrefixingEncryptionProvider implements BaseE2EEManager {
  isEnabled = true;

  isDataChannelEncryptionEnabled = true;

  setup(_room: Room) {}

  setupEngine(_engine: RTCEngine) {}

  setParticipantCryptorEnabled(_enabled: boolean, _participantIdentity: string) {}

  setSifTrailer(_trailer: Uint8Array) {}

  on(_event: any, _listener: any): this {
    return this;
  }

  /** A fake "encryption" provider used for test purposes. Adds a prefix to the payload. */
  async encryptData(data: Uint8Array): Promise<EncryptDataResponseMessage['data']> {
    const prefix = new Uint8Array([0xde, 0xad, 0xbe, 0xef]);

    const output = new Uint8Array(prefix.length + data.length);
    output.set(prefix, 0);
    output.set(data, prefix.length);

    return {
      uuid: crypto.randomUUID(),
      payload: output,
      iv: new Uint8Array(12), // Just leaving this empty, is this a bad idea?
      keyIndex: 0,
    };
  }

  /** A fake "decryption" provider used for test purposes. Assumes the payload is prefixed with
   * 0xdeafbeef, which is stripped off. */
  async handleEncryptedData(
    payload: Uint8Array,
    _iv: Uint8Array,
    _participantIdentity: string,
    _keyIndex: number,
  ): Promise<DecryptDataResponseMessage['data']> {
    if (payload[0] !== 0xde || payload[1] !== 0xad || payload[2] !== 0xbe || payload[3] !== 0xef) {
      throw new Error(
        `PrefixingEncryptionProvider: first four bytes of payload were not 0xdeadbeef, found ${payload.slice(0, 4)}`,
      );
    }

    return {
      uuid: crypto.randomUUID(),
      payload: payload.slice(4),
    };
  }
}

describe('DataTrackOutgoingManager', () => {
  it('should test track publishing (ok case)', async () => {
    const manager = new OutgoingDataTrackManager();
    const managerEvents = subscribeToEvents<DataTrackOutgoingManagerCallbacks>(manager, [
      'sfuPublishRequest',
    ]);

    const localDataTrack = new LocalDataTrack({ name: 'test' }, manager);
    expect(localDataTrack.isPublished()).toStrictEqual(false);

    // 1. Publish a data track
    const publishRequestPromise = localDataTrack.publish();

    // 2. This publish request should be sent along to the SFU
    const sfuPublishEvent = await managerEvents.waitFor('sfuPublishRequest');
    expect(sfuPublishEvent.name).toStrictEqual('test');
    expect(sfuPublishEvent.usesE2ee).toStrictEqual(false);
    const handle = sfuPublishEvent.handle;

    // 3. Respond to the SFU publish request with an OK response
    manager.receivedSfuPublishResponse(handle, {
      type: 'ok',
      data: {
        sid: 'bogus-sid',
        pubHandle: sfuPublishEvent.handle,
        name: 'test',
        usesE2ee: false,
      },
    });

    // Make sure that the original input event resolves.
    await publishRequestPromise;
    expect(localDataTrack.isPublished()).toStrictEqual(true);
  });

  it('should test track publishing (error case)', async () => {
    const manager = new OutgoingDataTrackManager();
    const managerEvents = subscribeToEvents<DataTrackOutgoingManagerCallbacks>(manager, [
      'sfuPublishRequest',
    ]);

    // 1. Publish a data track
    const localDataTrack = new LocalDataTrack({ name: 'test' }, manager);
    const publishRequestPromise = localDataTrack.publish();

    // 2. This publish request should be sent along to the SFU
    const sfuPublishEvent = await managerEvents.waitFor('sfuPublishRequest');

    // 3. Respond to the SFU publish request with an ERROR response
    manager.receivedSfuPublishResponse(sfuPublishEvent.handle, {
      type: 'error',
      error: DataTrackPublishError.limitReached(),
    });

    // Make sure that the rejection bubbles back to the caller
    await expect(publishRequestPromise).rejects.toThrowError(
      'Data track publication limit reached',
    );
  });

  it('should test track publishing (cancellation half way through)', async () => {
    const manager = new OutgoingDataTrackManager();
    const managerEvents = subscribeToEvents<DataTrackOutgoingManagerCallbacks>(manager, [
      'sfuPublishRequest',
      'sfuUnpublishRequest',
    ]);

    // 1. Publish a data track
    const controller = new AbortController();
    const localDataTrack = new LocalDataTrack({ name: 'test' }, manager);
    const publishRequestPromise = localDataTrack.publish(controller.signal);

    // 2. This publish request should be sent along to the SFU
    const sfuPublishEvent = await managerEvents.waitFor('sfuPublishRequest');
    expect(sfuPublishEvent.name).toStrictEqual('test');
    expect(sfuPublishEvent.usesE2ee).toStrictEqual(false);
    const handle = sfuPublishEvent.handle;

    // 3. Explictly cancel the publish
    controller.abort();

    // 4. Make sure an unpublish event is sent so that the SFU cleans up things properly
    // on its end as well
    const sfuUnpublishEvent = await managerEvents.waitFor('sfuUnpublishRequest');
    expect(sfuUnpublishEvent.handle).toStrictEqual(handle);

    // 5. Make sure cancellation is bubbled up as an error to stop further execution
    await expect(publishRequestPromise).rejects.toStrictEqual(DataTrackPublishError.cancelled());
  });

  it('should test track publishing (cancellation before it starts)', async () => {
    const manager = new OutgoingDataTrackManager();
    const managerEvents = subscribeToEvents<DataTrackOutgoingManagerCallbacks>(manager, [
      'sfuPublishRequest',
      'sfuUnpublishRequest',
    ]);

    // Publish a data track
    const localDataTrack = new LocalDataTrack({ name: 'test' }, manager);
    const publishRequestPromise = localDataTrack.publish(AbortSignal.abort(/* already aborted */));

    // Make sure cancellation is immediately bubbled up
    await expect(publishRequestPromise).rejects.toStrictEqual(DataTrackPublishError.cancelled());

    // And there were no pending sfu publish requests sent
    expect(managerEvents.areThereBufferedEvents('sfuPublishRequest')).toBe(false);
  });

  it('should test track publishing, unpublishing, and republishing again', async () => {
    const manager = new OutgoingDataTrackManager();
    const managerEvents = subscribeToEvents<DataTrackOutgoingManagerCallbacks>(manager, [
      'sfuPublishRequest',
      'sfuUnpublishRequest',
    ]);

    // 1. Create a local data track
    const localDataTrack = new LocalDataTrack({ name: 'test' }, manager);
    expect(localDataTrack.isPublished()).toStrictEqual(false);

    // 2. Publish it
    const publishRequestPromise = localDataTrack.publish();

    // 3. This publish request should be sent along to the SFU
    const sfuPublishEvent = await managerEvents.waitFor('sfuPublishRequest');
    expect(sfuPublishEvent.name).toStrictEqual('test');
    expect(sfuPublishEvent.usesE2ee).toStrictEqual(false);
    const handle = sfuPublishEvent.handle;

    // 4. Respond to the SFU publish request with an OK response
    manager.receivedSfuPublishResponse(handle, {
      type: 'ok',
      data: {
        sid: 'bogus-sid',
        pubHandle: sfuPublishEvent.handle,
        name: 'test',
        usesE2ee: false,
      },
    });

    // Make sure that the original input event resolves.
    await publishRequestPromise;

    // 5. Now the data track should be published
    expect(localDataTrack.isPublished()).toStrictEqual(true);

    // 6. Unpublish the data track
    const unpublishRequestPromise = localDataTrack.unpublish();
    const sfuUnpublishEvent = await managerEvents.waitFor('sfuUnpublishRequest');
    manager.receivedSfuUnpublishResponse(sfuUnpublishEvent.handle);
    await unpublishRequestPromise;

    // 7. Now the data track should be unpublished
    expect(localDataTrack.isPublished()).toStrictEqual(false);

    // 8. Now, republish the track and make sure that be done a second time
    const publishRequestPromise2 = localDataTrack.publish();
    const sfuPublishEvent2 = await managerEvents.waitFor('sfuPublishRequest');
    expect(sfuPublishEvent2.name).toStrictEqual('test');
    expect(sfuPublishEvent2.usesE2ee).toStrictEqual(false);
    const handle2 = sfuPublishEvent2.handle;
    manager.receivedSfuPublishResponse(handle2, {
      type: 'ok',
      data: {
        sid: 'bogus-sid',
        pubHandle: sfuPublishEvent2.handle,
        name: 'test',
        usesE2ee: false,
      },
    });
    await publishRequestPromise2;

    // 9. Ensure that the track is published again
    expect(localDataTrack.isPublished()).toStrictEqual(true);

    // 10. Also ensure that the handle used on the second publish attempt differs from the first
    // publish attempt.
    expect(handle).not.toStrictEqual(handle2);
  });

  it.each([
    // Single packet payload case
    [
      new Uint8Array([0x01, 0x02, 0x03, 0x04, 0x05]),
      [
        {
          header: {
            extensions: {
              e2ee: null,
              userTimestamp: null,
            },
            frameNumber: 0,
            marker: FrameMarker.Single,
            sequence: 0,
            timestamp: expect.anything(),
            trackHandle: 5,
          },
          payload: new Uint8Array([0x01, 0x02, 0x03, 0x04, 0x05]),
        },
      ],
    ],

    // Multi packet payload case
    [
      new Uint8Array(24_000).fill(0xbe),
      [
        {
          header: {
            extensions: {
              e2ee: null,
              userTimestamp: null,
            },
            frameNumber: 0,
            marker: FrameMarker.Start,
            sequence: 0,
            timestamp: expect.anything(),
            trackHandle: 5,
          },
          payload: new Uint8Array(15988 /* 16k mtu - 12 header bytes */).fill(0xbe),
        },
        {
          header: {
            extensions: {
              e2ee: null,
              userTimestamp: null,
            },
            frameNumber: 0,
            marker: FrameMarker.Final,
            sequence: 1,
            timestamp: expect.anything(),
            trackHandle: 5,
          },
          payload: new Uint8Array(8012 /* 24k payload - (16k mtu - 12 header bytes) */).fill(0xbe),
        },
      ],
    ],
  ])(
    'should test track payload sending',
    async (inputBytes: Uint8Array, outputPacketsJson: Array<unknown>) => {
      // Create a manager prefilled with a descriptor
      const manager = OutgoingDataTrackManager.withDescriptors(
        new Map([
          [
            DataTrackHandle.fromNumber(5),
            Descriptor.active(
              {
                sid: 'bogus-sid',
                pubHandle: 5,
                name: 'test',
                usesE2ee: false,
              },
              null,
            ),
          ],
        ]),
      );
      const managerEvents = subscribeToEvents<DataTrackOutgoingManagerCallbacks>(manager, [
        'packetAvailable',
      ]);

      const localDataTrack = LocalDataTrack.withExplicitHandle({ name: 'track name' }, manager, 5);

      // Kick off sending the bytes...
      localDataTrack.tryPush({ payload: inputBytes });

      // ... and make sure the corresponding events are emitted to tell the SFU to send the packets
      for (const outputPacketJson of outputPacketsJson) {
        const packetBytes = await managerEvents.waitFor('packetAvailable');
        const [packet] = DataTrackPacket.fromBinary(packetBytes.bytes);

        expect(packet.toJSON()).toStrictEqual(outputPacketJson);
      }
    },
  );

  it('should send e2ee encrypted datatrack payload', async () => {
    const manager = new OutgoingDataTrackManager({
      e2eeManager: new PrefixingEncryptionProvider(),
    });
    const managerEvents = subscribeToEvents<DataTrackOutgoingManagerCallbacks>(manager, [
      'sfuPublishRequest',
      'packetAvailable',
    ]);

    // 1. Publish a data track
    const localDataTrack = new LocalDataTrack({ name: 'test' }, manager);
    const publishRequestPromise = localDataTrack.publish();

    // 2. This publish request should be sent along to the SFU
    const sfuPublishEvent = await managerEvents.waitFor('sfuPublishRequest');
    expect(sfuPublishEvent.name).toStrictEqual('test');
    expect(sfuPublishEvent.usesE2ee).toStrictEqual(true); // NOTE: this is true, e2ee is enabled!
    const handle = sfuPublishEvent.handle;

    // 3. Respond to the SFU publish request with an OK response
    manager.receivedSfuPublishResponse(handle, {
      type: 'ok',
      data: {
        sid: 'bogus-sid',
        pubHandle: sfuPublishEvent.handle,
        name: 'test',
        usesE2ee: true, // NOTE: this is true, e2ee is enabled!
      },
    });

    // Get the connected local data track
    await publishRequestPromise;
    expect(localDataTrack.isPublished()).toStrictEqual(true);

    // Kick off sending the payload bytes
    localDataTrack.tryPush({ payload: new Uint8Array([0x01, 0x02, 0x03, 0x04, 0x05]) });

    // Make sure the packet that was sent was encrypted with the PrefixingEncryptionProvider
    const packetBytes = await managerEvents.waitFor('packetAvailable');
    const [packet] = DataTrackPacket.fromBinary(packetBytes.bytes);

    expect(packet.toJSON()).toStrictEqual({
      header: {
        extensions: {
          e2ee: {
            iv: new Uint8Array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]),
            keyIndex: 0,
            lengthBytes: 13,
            tag: 1,
          },
          userTimestamp: null,
        },
        frameNumber: 0,
        marker: 3,
        sequence: 0,
        timestamp: expect.anything(),
        trackHandle: 1,
      },
      payload: new Uint8Array([
        // Encryption added prefix
        0xde, 0xad, 0xbe, 0xef,
        // Actual payload
        0x01, 0x02, 0x03, 0x04, 0x05,
      ]),
    });
  });

  it('should test track unpublishing', async () => {
    // Create a manager prefilled with a descriptor
    const manager = OutgoingDataTrackManager.withDescriptors(
      new Map([
        [
          DataTrackHandle.fromNumber(5),
          Descriptor.active(
            {
              sid: 'bogus-sid',
              pubHandle: 5,
              name: 'test',
              usesE2ee: false,
            },
            null,
          ),
        ],
      ]),
    );
    const managerEvents = subscribeToEvents<DataTrackOutgoingManagerCallbacks>(manager, [
      'sfuUnpublishRequest',
    ]);

    // Make sure the descriptor is in there
    expect(manager.getDescriptor(5)?.type).toStrictEqual('active');

    // Unpublish data track
    const unpublishRequestPromise = manager.unpublishRequest(DataTrackHandle.fromNumber(5));

    const sfuUnpublishEvent = await managerEvents.waitFor('sfuUnpublishRequest');
    expect(sfuUnpublishEvent.handle).toStrictEqual(5);

    manager.receivedSfuUnpublishResponse(DataTrackHandle.fromNumber(5));

    await unpublishRequestPromise;

    // Make sure data track is no longer
    expect(manager.getDescriptor(5)).toStrictEqual(null);
  });

  it('should test a full reconnect', async () => {
    const pubHandle = 5;
    // Create a manager prefilled with a descriptor
    const manager = OutgoingDataTrackManager.withDescriptors(
      new Map([
        [
          DataTrackHandle.fromNumber(5),
          Descriptor.active(
            {
              sid: 'bogus-sid',
              pubHandle,
              name: 'test',
              usesE2ee: false,
            },
            null,
          ),
        ],
      ]),
    );
    const managerEvents = subscribeToEvents<DataTrackOutgoingManagerCallbacks>(manager, [
      'sfuPublishRequest',
      'packetAvailable',
      'sfuUnpublishRequest',
    ]);
    const localDataTrack = LocalDataTrack.withExplicitHandle({ name: 'track name' }, manager, 5);

    // Make sure the descriptor is in there
    expect(manager.getDescriptor(5)?.type).toStrictEqual('active');

    // Simulate a full reconnect, which means that any published tracks will need to be republished.
    manager.sfuWillRepublishTracks();

    // Even though behind the scenes the SFU publications are not active, the user should still see
    // it as "published", sfu reconnects are an implementation detail
    expect(localDataTrack.isPublished()).toStrictEqual(true);

    // But, even though `isPublished` is true, pushing data should drop (no sfu to send them to!)
    await expect(() =>
      localDataTrack.tryPush({ payload: new Uint8Array([0x01, 0x02, 0x03, 0x04, 0x05]) }),
    ).rejects.toThrowError('Frame was dropped');

    // 2. This publish request should be sent along to the SFU
    const sfuPublishEvent = await managerEvents.waitFor('sfuPublishRequest');
    expect(sfuPublishEvent.name).toStrictEqual('test');
    expect(sfuPublishEvent.usesE2ee).toStrictEqual(false);
    const handle = sfuPublishEvent.handle;
    expect(handle).toStrictEqual(pubHandle);

    // 3. Respond to the SFU publish request with an OK response
    manager.receivedSfuPublishResponse(handle, {
      type: 'ok',
      data: {
        sid: 'bogus-sid-REPUBLISHED',
        pubHandle: sfuPublishEvent.handle,
        name: 'test',
        usesE2ee: false,
      },
    });

    // After all this, the local data track should still be published
    expect(localDataTrack.isPublished()).toStrictEqual(true);

    // And the sid should be the new value
    expect(localDataTrack.info!.sid).toStrictEqual('bogus-sid-REPUBLISHED');

    // And now that the tracks are backed by the SFU again, pushes should function!
    await localDataTrack.tryPush({ payload: new Uint8Array([0x01, 0x02, 0x03, 0x04, 0x05]) });
    await managerEvents.waitFor('packetAvailable');
  });

  it('should query currently active descriptors', async () => {
    // Create a manager prefilled with a descriptor
    const manager = OutgoingDataTrackManager.withDescriptors(
      new Map([
        [
          DataTrackHandle.fromNumber(2),
          Descriptor.active(
            {
              sid: 'bogus-sid-2',
              pubHandle: 2,
              name: 'twotwotwo',
              usesE2ee: false,
            },
            null,
          ),
        ],
        [
          DataTrackHandle.fromNumber(6),
          Descriptor.active(
            {
              sid: 'bogus-sid-6',
              pubHandle: 6,
              name: 'sixsixsix',
              usesE2ee: false,
            },
            null,
          ),
        ],
      ]),
    );

    const result = await manager.queryPublished();

    expect(result).toStrictEqual([
      { sid: 'bogus-sid-2', pubHandle: 2, name: 'twotwotwo', usesE2ee: false },
      { sid: 'bogus-sid-6', pubHandle: 6, name: 'sixsixsix', usesE2ee: false },
    ]);
  });

  it('should shutdown cleanly', async () => {
    // Create a manager prefilled with a descriptor
    const pendingDescriptor = Descriptor.pending();
    const manager = OutgoingDataTrackManager.withDescriptors(
      new Map<DataTrackHandle, Descriptor>([
        [DataTrackHandle.fromNumber(2), pendingDescriptor],
        [
          DataTrackHandle.fromNumber(6),
          Descriptor.active(
            {
              sid: 'bogus-sid-6',
              pubHandle: 6,
              name: 'sixsixsix',
              usesE2ee: false,
            },
            null,
          ),
        ],
      ]),
    );
    const managerEvents = subscribeToEvents<DataTrackOutgoingManagerCallbacks>(manager, [
      'sfuUnpublishRequest',
    ]);

    // Shut down the manager
    const shutdownPromise = manager.shutdown();

    // The pending data track should be cancelled
    await expect(pendingDescriptor.completionFuture.promise).rejects.toThrowError(
      'Room disconnected',
    );

    // And the active data track should be requested to be unpublished
    const unpublishEvent = await managerEvents.waitFor('sfuUnpublishRequest');
    expect(unpublishEvent.handle).toStrictEqual(6);

    // Acknowledge that the unpublish has occurred
    manager.receivedSfuUnpublishResponse(DataTrackHandle.fromNumber(6));

    await shutdownPromise;
  });
});