1097 lines · 35.8 KB · Rust
RawDownload
use crate::{AudioStream, Participant, RemoteTrack, RoomEvent, TrackPublication}; use crate::mock_client::{participant::*, publication::*, track::*};use anyhow::{Context as _, Result};use async_trait::async_trait;use collections::{BTreeMap, HashMap, HashSet, btree_map::Entry as BTreeEntry, hash_map::Entry};use gpui::{App, AsyncApp, BackgroundExecutor};use livekit_api::{proto, token};use parking_lot::Mutex;use postage::{mpsc, sink::Sink};use std::sync::{    Arc, Weak,    atomic::{AtomicBool, AtomicU64, Ordering::SeqCst},}; #[derive(Clone, Debug, Eq, Hash, PartialEq, PartialOrd, Ord)]pub struct ParticipantIdentity(pub String); #[derive(Clone, Debug, Eq, Hash, PartialEq, PartialOrd, Ord)]pub struct TrackSid(pub(crate) String); impl std::fmt::Display for TrackSid {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        self.0.fmt(f)    }} impl TryFrom<String> for TrackSid {    type Error = anyhow::Error;     fn try_from(value: String) -> Result<Self, Self::Error> {        Ok(TrackSid(value))    }} #[derive(Copy, Clone, Debug, Eq, Hash, PartialEq, PartialOrd, Ord)]#[non_exhaustive]pub enum ConnectionState {    Connected,    Disconnected,} #[derive(Clone, Debug, Default)]pub struct SessionStats {    pub publisher_stats: Vec<RtcStats>,    pub subscriber_stats: Vec<RtcStats>,} #[derive(Clone, Debug)]pub enum RtcStats {} static SERVERS: Mutex<BTreeMap<String, Arc<TestServer>>> = Mutex::new(BTreeMap::new()); pub struct TestServer {    pub url: String,    pub api_key: String,    pub secret_key: String,    rooms: Mutex<HashMap<String, TestServerRoom>>,    executor: BackgroundExecutor,    timestamp_source: Arc<dyn token::UnixTimestampSource>,} pub struct ManualUnixTimestampSource(AtomicU64); impl ManualUnixTimestampSource {    pub fn new(timestamp: u64) -> Self {        Self(AtomicU64::new(timestamp))    }     pub fn advance(&self) {        self.0.fetch_add(1, SeqCst);    }} impl token::UnixTimestampSource for ManualUnixTimestampSource {    fn unix_timestamp(&self) -> Result<u64> {        Ok(self.0.load(SeqCst))    }} impl TestServer {    pub fn create(        url: String,        api_key: String,        secret_key: String,        executor: BackgroundExecutor,    ) -> Result<Arc<TestServer>> {        Self::create_with_timestamp_source(            url,            api_key,            secret_key,            executor,            Arc::new(token::SystemUnixTimestampSource),        )    }     pub fn create_with_timestamp_source(        url: String,        api_key: String,        secret_key: String,        executor: BackgroundExecutor,        timestamp_source: Arc<dyn token::UnixTimestampSource>,    ) -> Result<Arc<TestServer>> {        let mut servers = SERVERS.lock();        if let BTreeEntry::Vacant(e) = servers.entry(url.clone()) {            let server = Arc::new(TestServer {                url,                api_key,                secret_key,                rooms: Default::default(),                executor,                timestamp_source,            });            e.insert(server.clone());            Ok(server)        } else {            anyhow::bail!("a server with url {url:?} already exists");        }    }     fn get(url: &str) -> Result<Arc<TestServer>> {        Ok(SERVERS            .lock()            .get(url)            .context("no server found for url")?            .clone())    }     pub fn teardown(&self) -> Result<()> {        SERVERS            .lock()            .remove(&self.url)            .with_context(|| format!("server with url {:?} does not exist", self.url))?;        Ok(())    }     pub fn create_api_client(&self) -> TestApiClient {        TestApiClient {            url: self.url.clone(),        }    }     #[cfg(any(test, feature = "test-support"))]    fn validate_token<'a>(&self, token: &'a str) -> Result<token::ClaimGrants<'a>> {        token::validate_with_timestamp_source(            token,            &self.secret_key,            self.timestamp_source.as_ref(),        )    }     #[cfg(not(any(test, feature = "test-support")))]    fn validate_token<'a>(&self, token: &'a str) -> Result<token::ClaimGrants<'a>> {        token::validate(token, &self.secret_key)    }     pub async fn create_room(&self, room: String) -> Result<()> {        self.simulate_random_delay().await;         let mut server_rooms = self.rooms.lock();        if let Entry::Vacant(e) = server_rooms.entry(room.clone()) {            e.insert(Default::default());            Ok(())        } else {            anyhow::bail!("{room:?} already exists");        }    }     async fn delete_room(&self, room: String) -> Result<()> {        self.simulate_random_delay().await;         let mut server_rooms = self.rooms.lock();        server_rooms            .remove(&room)            .with_context(|| format!("room {room:?} does not exist"))?;        Ok(())    }     async fn join_room(&self, token: String, client_room: Room) -> Result<ParticipantIdentity> {        self.simulate_random_delay().await;         let claims = self.validate_token(&token)?;        let identity = ParticipantIdentity(            claims                .sub                .context("missing participant identity")?                .to_string(),        );        let room_name = claims.video.room.context("missing room name")?.to_string();        let mut server_rooms = self.rooms.lock();        let room = (*server_rooms).entry(room_name.clone()).or_default();        if let Some(revoked_before) = room.token_revocations.get(&identity) {            anyhow::ensure!(claims.nbf >= *revoked_before, "invalid token: revoked");        }         if let Entry::Vacant(e) = room.client_rooms.entry(identity.clone()) {            for server_track in &room.video_tracks {                let track = RemoteTrack::Video(RemoteVideoTrack {                    server_track: server_track.clone(),                    _room: client_room.downgrade(),                });                client_room                    .0                    .lock()                    .updates_tx                    .blocking_send(RoomEvent::TrackSubscribed {                        track: track.clone(),                        publication: RemoteTrackPublication {                            sid: server_track.sid.clone(),                            room: client_room.downgrade(),                            track,                        },                        participant: RemoteParticipant {                            room: client_room.downgrade(),                            identity: server_track.publisher_id.clone(),                        },                    })                    .unwrap();            }            for server_track in &room.audio_tracks {                let track = RemoteTrack::Audio(RemoteAudioTrack {                    server_track: server_track.clone(),                    room: client_room.downgrade(),                });                client_room                    .0                    .lock()                    .updates_tx                    .blocking_send(RoomEvent::TrackSubscribed {                        track: track.clone(),                        publication: RemoteTrackPublication {                            sid: server_track.sid.clone(),                            room: client_room.downgrade(),                            track,                        },                        participant: RemoteParticipant {                            room: client_room.downgrade(),                            identity: server_track.publisher_id.clone(),                        },                    })                    .unwrap();            }            e.insert(client_room);            Ok(identity)        } else {            anyhow::bail!("{identity:?} attempted to join room {room_name:?} twice");        }    }     async fn leave_room(&self, token: String) -> Result<()> {        self.simulate_random_delay().await;         let claims = self.validate_token(&token)?;        let identity = ParticipantIdentity(claims.sub.unwrap().to_string());        let room_name = claims.video.room.unwrap();        let mut server_rooms = self.rooms.lock();        let room = server_rooms            .get_mut(&*room_name)            .with_context(|| format!("room {room_name:?} does not exist"))?;        room.client_rooms.remove(&identity).with_context(|| {            format!("{identity:?} attempted to leave room {room_name:?} before joining it")        })?;        Ok(())    }     fn remote_participants(        &self,        token: String,    ) -> Result<HashMap<ParticipantIdentity, RemoteParticipant>> {        let claims = self.validate_token(&token)?;        let local_identity = ParticipantIdentity(claims.sub.unwrap().to_string());        let room_name = claims.video.room.unwrap().to_string();         if let Some(server_room) = self.rooms.lock().get(&room_name) {            let room = server_room                .client_rooms                .get(&local_identity)                .unwrap()                .downgrade();            Ok(server_room                .client_rooms                .iter()                .filter(|(identity, _)| *identity != &local_identity)                .map(|(identity, _)| {                    (                        identity.clone(),                        RemoteParticipant {                            room: room.clone(),                            identity: identity.clone(),                        },                    )                })                .collect())        } else {            Ok(Default::default())        }    }     async fn remove_participant(        &self,        room_name: String,        identity: ParticipantIdentity,    ) -> Result<()> {        self.simulate_random_delay().await;        let revoked_before = self.timestamp_source.unix_timestamp()?;         let mut server_rooms = self.rooms.lock();        let room = server_rooms            .get_mut(&room_name)            .with_context(|| format!("room {room_name} does not exist"))?;        let removed_room = room            .client_rooms            .remove(&identity)            .with_context(|| format!("participant {identity:?} did not join room {room_name:?}"))?;        room.token_revocations.insert(identity, revoked_before);        let mut removed_room = removed_room.0.lock();        removed_room.connection_state = ConnectionState::Disconnected;        removed_room            .updates_tx            .blocking_send(RoomEvent::Disconnected {                reason: "PARTICIPANT_REMOVED",            })            .ok();        Ok(())    }     async fn update_participant(        &self,        room_name: String,        identity: String,        permission: proto::ParticipantPermission,    ) -> Result<()> {        self.simulate_random_delay().await;        let revoked_before = self.timestamp_source.unix_timestamp()?;         let mut server_rooms = self.rooms.lock();        let room = server_rooms            .get_mut(&room_name)            .with_context(|| format!("room {room_name} does not exist"))?;        let identity = ParticipantIdentity(identity);        room.participant_permissions            .insert(identity.clone(), permission);        // Permission changes in LiveKit Cloud invalidate existing participant        // tokens, so the mock needs to reject tokens minted before the update.        room.token_revocations.insert(identity, revoked_before);        Ok(())    }     pub async fn disconnect_client(&self, client_identity: String) {        let client_identity = ParticipantIdentity(client_identity);         self.simulate_random_delay().await;         let mut server_rooms = self.rooms.lock();        for room in server_rooms.values_mut() {            if let Some(room) = room.client_rooms.remove(&client_identity) {                let mut room = room.0.lock();                room.connection_state = ConnectionState::Disconnected;                room.updates_tx                    .blocking_send(RoomEvent::Disconnected {                        reason: "SIGNAL_CLOSED",                    })                    .ok();            }        }    }     pub(crate) async fn publish_video_track(        &self,        token: String,        _local_track: LocalVideoTrack,    ) -> Result<TrackSid> {        self.simulate_random_delay().await;         let claims = self.validate_token(&token)?;        let identity = ParticipantIdentity(claims.sub.unwrap().to_string());        let room_name = claims.video.room.unwrap();         let mut server_rooms = self.rooms.lock();        let room = server_rooms            .get_mut(&*room_name)            .with_context(|| format!("room {room_name} does not exist"))?;         let can_publish = room            .participant_permissions            .get(&identity)            .map(|permission| permission.can_publish)            .or(claims.video.can_publish)            .unwrap_or(true);         anyhow::ensure!(can_publish, "user is not allowed to publish");         let sid: TrackSid = format!("TR_{}", nanoid::nanoid!(17)).try_into().unwrap();        let server_track = Arc::new(TestServerVideoTrack {            sid: sid.clone(),            publisher_id: identity.clone(),        });         room.video_tracks.push(server_track.clone());         for (room_identity, client_room) in &room.client_rooms {            if *room_identity != identity {                let track = RemoteTrack::Video(RemoteVideoTrack {                    server_track: server_track.clone(),                    _room: client_room.downgrade(),                });                let publication = RemoteTrackPublication {                    sid: sid.clone(),                    room: client_room.downgrade(),                    track: track.clone(),                };                let participant = RemoteParticipant {                    identity: identity.clone(),                    room: client_room.downgrade(),                };                client_room                    .0                    .lock()                    .updates_tx                    .blocking_send(RoomEvent::TrackSubscribed {                        track,                        publication,                        participant,                    })                    .unwrap();            }        }         Ok(sid)    }     pub(crate) async fn publish_audio_track(        &self,        token: String,        _local_track: &LocalAudioTrack,    ) -> Result<TrackSid> {        self.simulate_random_delay().await;         let claims = self.validate_token(&token)?;        let identity = ParticipantIdentity(claims.sub.unwrap().to_string());        let room_name = claims.video.room.unwrap();         let mut server_rooms = self.rooms.lock();        let room = server_rooms            .get_mut(&*room_name)            .with_context(|| format!("room {room_name} does not exist"))?;         let can_publish = room            .participant_permissions            .get(&identity)            .map(|permission| permission.can_publish)            .or(claims.video.can_publish)            .unwrap_or(true);         anyhow::ensure!(can_publish, "user is not allowed to publish");         let sid: TrackSid = format!("TR_{}", nanoid::nanoid!(17)).try_into().unwrap();        let server_track = Arc::new(TestServerAudioTrack {            sid: sid.clone(),            publisher_id: identity.clone(),            muted: AtomicBool::new(false),        });         room.audio_tracks.push(server_track.clone());         for (room_identity, client_room) in &room.client_rooms {            if *room_identity != identity {                let track = RemoteTrack::Audio(RemoteAudioTrack {                    server_track: server_track.clone(),                    room: client_room.downgrade(),                });                let publication = RemoteTrackPublication {                    sid: sid.clone(),                    room: client_room.downgrade(),                    track: track.clone(),                };                let participant = RemoteParticipant {                    identity: identity.clone(),                    room: client_room.downgrade(),                };                client_room                    .0                    .lock()                    .updates_tx                    .blocking_send(RoomEvent::TrackSubscribed {                        track,                        publication,                        participant,                    })                    .ok();            }        }         Ok(sid)    }     pub(crate) async fn unpublish_track(&self, token: String, track_sid: &TrackSid) -> Result<()> {        let claims = self.validate_token(&token)?;        let identity = ParticipantIdentity(claims.sub.unwrap().to_string());        let room_name = claims.video.room.unwrap();

Lines 1–500 of 1097.