feat: move signal_client to livekit-api (#98)

2023-06-26 22:59:55 +02:00
parent 678e20b0e3
commit efc334af8e
11 changed files with 102 additions and 25 deletions
@@ -3,13 +3,38 @@ name = "livekit-api"
 version = "0.1.0"
 edition = "2021"

+[features]
+# By default ws TLS is not enabled
+default = ["services", "access-token"]
+
+signal-client = ["dep:tokio-tungstenite", "dep:tokio", "dep:futures-util" ]
+services = ["dep:reqwest"]
+access-token = ["dep:jsonwebtoken"]
+
+# Note that the following features only change the behavior of tokio-tungstenite.
+# It doesn't change the behavior of libwebrtc/webrtc-sys
+native-tls = ["tokio-tungstenite?/native-tls", "reqwest?/native-tls"]
+native-tls-vendored = ["tokio-tungstenite?/native-tls-vendored", "reqwest?/native-tls-vendored"]
+rustls-tls-native-roots = ["tokio-tungstenite?/rustls-tls-native-roots", "reqwest?/rustls-tls-native-roots"]
+rustls-tls-webpki-roots = ["tokio-tungstenite?/rustls-tls-webpki-roots", "reqwest?/rustls-tls-webpki-roots"]
+__rustls-tls = ["tokio-tungstenite?/__rustls-tls", "reqwest?/__rustls"]
+
 [dependencies]
-jsonwebtoken = {version = "8", default-features = false }
-thiserror = "1.0"
-serde = {version = "1.0", features = ["derive"] }
-serde_json = "1.0"
 livekit-protocol = { path = "../livekit-protocol", version = "0.1.0" }
+thiserror = "1.0"
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
 sha2 = "0.10.6"
-reqwest = { version = "0.11", features=["json"] }
 url = "2.3.1"
-prost = "0.11.9"
+log = "0.4"
+parking_lot = { version = "0.12.1" }
+prost = "0.11"
+
+# access_token & services
+jsonwebtoken = { version = "8", default-features = false, optional = true }
+reqwest = { version = "0.11", default-features = false, features = ["json"], optional = true }
+
+# signal_client
+tokio-tungstenite = { version = "0.19", optional = true }
+tokio = { version = "1", features = ["full"], optional = true }
+futures-util = { version = "0.3", default-features = false, features = ["sink"], optional = true }
@@ -1,7 +1,14 @@
+#[cfg(feature = "access-token")]
 pub mod access_token;
+
+#[cfg(feature = "services")]
 pub mod services;
+
+#[cfg(feature = "signal-client")]
+pub mod signal_client;
 pub mod webhook_receiver;

+#[allow(dead_code)]
 pub(crate) fn get_env_keys() -> Result<(String, String), std::env::VarError> {
    let api_key = std::env::var("LIVEKIT_API_KEY")?;
    let api_secret = std::env::var("LIVEKIT_API_SECRET")?;
@@ -0,0 +1,198 @@
+use crate::signal_client::signal_stream::SignalStream;
+use livekit_protocol as proto;
+use parking_lot::Mutex;
+use std::fmt::Debug;
+use std::time::Duration;
+use thiserror::Error;
+use tokio::sync::mpsc;
+use tokio::sync::RwLock as AsyncRwLock;
+use tokio_tungstenite::tungstenite::Error as WsError;
+
+mod signal_stream;
+
+pub type SignalEmitter = mpsc::Sender<SignalEvent>;
+pub type SignalEvents = mpsc::Receiver<SignalEvent>;
+pub type SignalResult<T> = Result<T, SignalError>;
+
+pub const JOIN_RESPONSE_TIMEOUT: Duration = Duration::from_secs(5);
+pub const PROTOCOL_VERSION: u32 = 8;
+
+#[derive(Error, Debug)]
+pub enum SignalError {
+    #[error("already connected")]
+    AlreadyConnected,
+    #[error("ws failure: {0}")]
+    WsError(#[from] WsError),
+    #[error("failed to parse the url")]
+    UrlParse(#[from] url::ParseError),
+    #[error("failed to decode messages from server")]
+    ProtoParse(#[from] prost::DecodeError),
+    #[error("{0}")]
+    Timeout(String),
+}
+
+/// Events used by the RTCEngine who will handle the reconnection logic
+#[derive(Debug)]
+pub enum SignalEvent {
+    Open,
+    Signal(proto::signal_response::Message),
+    Close,
+}
+
+#[derive(Debug, Clone)]
+pub struct SignalOptions {
+    pub auto_subscribe: bool,
+    pub adaptive_stream: bool,
+}
+
+impl Default for SignalOptions {
+    fn default() -> Self {
+        Self {
+            auto_subscribe: true,
+            adaptive_stream: false,
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct SignalClient {
+    stream: AsyncRwLock<Option<SignalStream>>,
+    url: String,
+    token: Mutex<String>, // TODO(theomonnom): Handle token refresh
+    join_response: proto::JoinResponse,
+    options: SignalOptions,
+    emitter: SignalEmitter,
+}
+
+impl SignalClient {
+    pub async fn connect(
+        url: &str,
+        token: &str,
+        options: SignalOptions,
+    ) -> SignalResult<(Self, proto::JoinResponse, SignalEvents)> {
+        let (emitter, mut events) = mpsc::channel(8);
+        let lk_url = get_livekit_url(url, token, &options)?;
+        let new_stream = SignalStream::connect(lk_url, emitter.clone()).await?;
+        let join_response = get_join_response(&mut events).await?;
+
+        Ok((
+            Self {
+                stream: AsyncRwLock::new(Some(new_stream)),
+                url: url.to_string(),
+                token: Mutex::new(token.to_string()),
+                join_response: join_response.clone(),
+                options,
+                emitter,
+            },
+            join_response,
+            events,
+        ))
+    }
+
+    // Restart is called when trying to resume the room (RtcSession resume)
+    // TODO(theomonom): Should this be renamed to resume?
+    pub async fn restart(&self) -> SignalResult<()> {
+        self.close().await;
+
+        let sid = &self.join_response.participant.as_ref().unwrap().sid;
+        let token = self.token.lock().clone();
+
+        let mut lk_url = get_livekit_url(&self.url, &token, &self.options)?;
+        lk_url
+            .query_pairs_mut()
+            .append_pair("reconnect", "1")
+            .append_pair("sid", sid);
+
+        let new_stream = SignalStream::connect(lk_url, self.emitter.clone()).await?;
+        *self.stream.write().await = Some(new_stream);
+        Ok(())
+    }
+
+    pub async fn close(&self) {
+        if let Some(stream) = self.stream.write().await.take() {
+            stream.close().await;
+        }
+    }
+
+    pub async fn send(&self, signal: proto::signal_request::Message) {
+        // TODO: Check if currently reconnecting and queue message
+
+        if let Some(stream) = self.stream.read().await.as_ref() {
+            if stream.send(signal).await.is_ok() {
+                return;
+            }
+        }
+        // TODO(theomonnom): return result?
+    }
+
+    #[allow(dead_code)]
+    pub async fn clear_queue(&self) {
+        // TODO(theomonnom): Clear the queue
+    }
+
+    pub async fn flush_queue(&self) {
+        // TODO(theomonnom): Send the queue
+    }
+
+    pub fn join_response(&self) -> proto::JoinResponse {
+        self.join_response.clone()
+    }
+
+    pub fn options(&self) -> SignalOptions {
+        self.options.clone()
+    }
+
+    pub fn url(&self) -> String {
+        self.url.clone()
+    }
+
+    pub fn token(&self) -> String {
+        self.token.lock().clone()
+    }
+}
+
+fn get_livekit_url(url: &str, token: &str, options: &SignalOptions) -> SignalResult<url::Url> {
+    let mut lk_url = url::Url::parse(url)?;
+    lk_url.set_path("/rtc");
+    lk_url
+        .query_pairs_mut()
+        .append_pair("access_token", token)
+        .append_pair("protocol", PROTOCOL_VERSION.to_string().as_str())
+        .append_pair(
+            "auto_subscribe",
+            if options.auto_subscribe { "1" } else { "0" },
+        )
+        .append_pair(
+            "adaptive_stream",
+            if options.adaptive_stream { "1" } else { "0" },
+        );
+
+    Ok(lk_url)
+}
+
+async fn get_join_response(receiver: &mut SignalEvents) -> SignalResult<proto::JoinResponse> {
+    let join = async {
+        while let Some(event) = receiver.recv().await {
+            match event {
+                SignalEvent::Signal(proto::signal_response::Message::Join(join)) => {
+                    return Ok(join)
+                }
+                SignalEvent::Close => break,
+                SignalEvent::Open => continue,
+                _ => {
+                    log::warn!(
+                        "received unexpected message while waiting for JoinResponse: {:?}",
+                        event
+                    );
+                    continue;
+                }
+            }
+        }
+
+        Err(WsError::ConnectionClosed)?
+    };
+
+    tokio::time::timeout(JOIN_RESPONSE_TIMEOUT, join)
+        .await
+        .map_err(|_| SignalError::Timeout("failed to receive JoinResponse".to_string()))?
+}
@@ -0,0 +1,183 @@
+use crate::signal_client::{SignalEmitter, SignalEvent, SignalResult};
+use futures_util::stream::{SplitSink, SplitStream};
+use futures_util::{SinkExt, StreamExt};
+use livekit_protocol as proto;
+use prost::Message as ProstMessage;
+use tokio::net::TcpStream;
+use tokio::sync::{mpsc, oneshot};
+use tokio::task::JoinHandle;
+use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
+use tokio_tungstenite::tungstenite::protocol::CloseFrame;
+use tokio_tungstenite::tungstenite::Message;
+use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream};
+
+type WebSocket = WebSocketStream<MaybeTlsStream<TcpStream>>;
+
+#[derive(Debug)]
+enum InternalMessage {
+    Signal {
+        signal: proto::signal_request::Message,
+        response_chn: oneshot::Sender<SignalResult<()>>,
+    },
+    Pong {
+        ping_data: Vec<u8>,
+    },
+    Close {
+        close_frame: Option<CloseFrame<'static>>,
+    },
+}
+
+/// SignalStream hold the WebSocket connection
+///
+/// It is replaced by [SignalClient] at each reconnection.
+#[derive(Debug)]
+pub(super) struct SignalStream {
+    internal_tx: mpsc::Sender<InternalMessage>,
+    read_handle: JoinHandle<()>,
+    write_handle: JoinHandle<()>,
+}
+
+impl SignalStream {
+    /// Connect to livekit websocket.
+    /// Return SignalError if the connections failed
+    ///
+    /// SignalStream will never try to reconnect if the connection has been
+    /// closed.
+    pub async fn connect(url: url::Url, emitter: SignalEmitter) -> SignalResult<Self> {
+        log::info!("connecting to SignalClient: {}", url);
+
+        let (ws_stream, _) = connect_async(url).await?;
+        let _ = emitter.send(SignalEvent::Open).await;
+
+        let (ws_writer, ws_reader) = ws_stream.split();
+        let (internal_tx, internal_rx) = mpsc::channel::<InternalMessage>(8);
+
+        let write_handle = tokio::spawn(Self::write_task(internal_rx, ws_writer, emitter.clone()));
+        let read_handle = tokio::spawn(Self::read_task(internal_tx.clone(), ws_reader, emitter));
+
+        Ok(Self {
+            internal_tx,
+            read_handle,
+            write_handle,
+        })
+    }
+
+    /// Close the websocket
+    /// It sends a CloseFrame to the server before closing
+    pub async fn close(self) {
+        let _ = self
+            .internal_tx
+            .send(InternalMessage::Close {
+                close_frame: Some(CloseFrame {
+                    code: CloseCode::Normal,
+                    reason: "disconnected by client".into(),
+                }),
+            })
+            .await;
+
+        let _ = self.write_handle.await;
+        let _ = self.read_handle.await;
+    }
+
+    /// Send a SignalRequest to the websocket
+    /// It also waits for the message to be sent
+    pub async fn send(&self, signal: proto::signal_request::Message) -> SignalResult<()> {
+        let (send, recv) = oneshot::channel();
+        let msg = InternalMessage::Signal {
+            signal,
+            response_chn: send,
+        };
+        let _ = self.internal_tx.send(msg).await;
+        recv.await.expect("channel closed")
+    }
+
+    /// This task is used to send messages to the websocket
+    /// It is also responsible for closing the connection
+    async fn write_task(
+        mut internal_rx: mpsc::Receiver<InternalMessage>,
+        mut ws_writer: SplitSink<WebSocket, Message>,
+        emitter: SignalEmitter,
+    ) {
+        while let Some(msg) = internal_rx.recv().await {
+            match msg {
+                InternalMessage::Signal {
+                    signal,
+                    response_chn,
+                } => {
+                    log::debug!("sending SignalRequest: {:?}", signal);
+
+                    let data = Message::Binary(
+                        proto::SignalRequest {
+                            message: Some(signal),
+                        }
+                        .encode_to_vec(),
+                    );
+
+                    if let Err(err) = ws_writer.send(data).await {
+                        log::error!("failed to send signal: {:?}", err);
+                        let _ = response_chn.send(Err(err.into()));
+                        break;
+                    }
+
+                    let _ = response_chn.send(Ok(()));
+                }
+                InternalMessage::Pong { ping_data } => {
+                    if let Err(err) = ws_writer.send(Message::Pong(ping_data)).await {
+                        log::error!("failed to send pong message: {:?}", err);
+                    }
+                }
+                InternalMessage::Close { close_frame } => {
+                    if let Some(close_frame) = close_frame {
+                        let _ = ws_writer.send(Message::Close(Some(close_frame))).await;
+                        let _ = ws_writer.flush().await;
+                    }
+                    break;
+                }
+            }
+        }
+
+        let _ = ws_writer.close().await;
+        let _ = emitter.send(SignalEvent::Close).await;
+    }
+
+    /// This task is used to read incoming messages from the websocket
+    /// and dispatch them through the EventEmitter.
+    ///
+    /// It can also send messages to [handle_write] task ( Used e.g. answer to pings )
+    async fn read_task(
+        internal_tx: mpsc::Sender<InternalMessage>,
+        mut ws_reader: SplitStream<WebSocket>,
+        emitter: SignalEmitter,
+    ) {
+        while let Some(msg) = ws_reader.next().await {
+            match msg {
+                Ok(Message::Binary(data)) => {
+                    let res = proto::SignalResponse::decode(data.as_slice())
+                        .expect("failed to decode SignalResponse");
+
+                    let msg = res.message.unwrap();
+                    log::debug!("received SignalResponse: {:?}", msg);
+                    let _ = emitter.send(SignalEvent::Signal(msg)).await;
+                }
+                Ok(Message::Ping(data)) => {
+                    let _ = internal_tx
+                        .send(InternalMessage::Pong { ping_data: data })
+                        .await;
+                    continue;
+                }
+                Ok(Message::Close(close)) => {
+                    log::debug!("server closed the connection: {:?}", close);
+                    break;
+                }
+                _ => {
+                    log::error!("unhandled websocket message {:?}", msg);
+                    break;
+                }
+            }
+        }
+
+        let _ = internal_tx
+            .send(InternalMessage::Close { close_frame: None })
+            .await;
+    }
+}