use std::fmt::{Debug, Formatter}; use futures::future::poll_fn; use futures_util::{SinkExt, StreamExt}; use futures_util::stream::{SplitSink, SplitStream}; use prost::Message as ProstMessage; use thiserror::Error; use tokio::net::TcpStream; use tokio::sync::{mpsc, oneshot}; use tokio::task::JoinHandle; use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream}; use tokio_tungstenite::tungstenite::{Error as WsError, Message}; use tokio_tungstenite::tungstenite::protocol::CloseFrame; use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode; use tracing::{event, Level}; use crate::proto::{signal_request, signal_response, SignalRequest, SignalResponse}; use crate::signal_client::SendMessage::Pong; pub const PROTOCOL_VERSION: u32 = 8; #[derive(Error, Debug)] pub enum SignalError { #[error("websocket failure")] WsError(#[from] WsError), #[error("failed to parse the url")] UrlParse(#[from] url::ParseError), #[error("failed to decode messages from server")] ProtoParse(#[from] prost::DecodeError), } type SignalResult = Result; type WebSocket = WebSocketStream>; #[derive(Debug)] enum RecvMessage { Signal { response_chn: oneshot::Sender>, }, } #[derive(Debug)] enum SendMessage { Signal { signal: signal_request::Message, response_chn: oneshot::Sender>, }, Pong { ping_data: Vec, }, } pub struct SignalClient { read_tx: mpsc::Sender, write_tx: mpsc::Sender, read_handle: JoinHandle<()>, write_handle: JoinHandle<()>, } impl Debug for SignalClient { fn fmt(&self, f: &mut Formatter) -> std::fmt::Result { write!(f, "SignalClient") } } #[tracing::instrument(skip(url, token))] pub async fn connect(url: &str, token: &str) -> SignalResult { 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()); event!(Level::DEBUG, "connecting to websocket: {}", lk_url); let (ws_stream, _) = connect_async(lk_url.clone()).await?; event!(Level::DEBUG, "connected to SignalClient"); let (ws_writer, ws_reader) = ws_stream.split(); let (read_tx, read_rx) = mpsc::channel::(8); let (write_tx, write_rx) = mpsc::channel::(8); let read_handle = tokio::spawn(SignalClient::ws_read(read_rx, ws_reader, write_tx.clone())); let write_handle = tokio::spawn(SignalClient::ws_write(write_rx, ws_writer)); Ok(SignalClient { read_tx, write_tx, read_handle, write_handle, }) } impl SignalClient { pub async fn close(self) { drop(self.read_tx); drop(self.write_tx); let _ = self.read_handle.await; let _ = self.write_handle.await; } pub async fn recv(&self) -> Option { let (send, recv) = oneshot::channel(); let msg = RecvMessage::Signal { response_chn: send }; let _ = self.read_tx.send(msg).await; recv.await.expect("channel closed") } pub async fn send(&self, signal: signal_request::Message) -> SignalResult<()> { let (send, recv) = oneshot::channel(); let msg = SendMessage::Signal { signal, response_chn: send, }; let _ = self.write_tx.send(msg).await; recv.await.expect("channel closed") } async fn ws_write( mut write_receiver: mpsc::Receiver, mut ws_writer: SplitSink, ) { while let Some(msg) = write_receiver.recv().await { match msg { SendMessage::Signal { signal, response_chn, } => { event!(Level::TRACE, "sending: {:?}", signal); let req = SignalRequest { message: Some(signal), }; let write_res = ws_writer.send(Message::Binary(req.encode_to_vec())).await; if let Err(err) = write_res { event!(Level::ERROR, "failed to send signal: {:?}", err); let _ = response_chn.send(Err(err.into())); break; } let _ = response_chn.send(Ok(())); } Pong { ping_data } => { if let Err(err) = ws_writer.send(Message::Pong(ping_data)).await { event!(Level::ERROR, "failed to send pong message: {:?}", err); } } } } let _ = ws_writer .send(Message::Close(Some(CloseFrame { code: CloseCode::Normal, reason: "disconnected by client".into(), }))) .await; let _ = ws_writer.flush().await; } async fn ws_read( mut read_receiver: mpsc::Receiver, mut ws_reader: SplitStream, write_tx: mpsc::Sender, ) { while let Some(RecvMessage::Signal { mut response_chn }) = read_receiver.recv().await { tokio::select! { read = Self::handle_msg(&mut ws_reader, &write_tx) => { let _ = response_chn.send(read); } _ = poll_fn(|cx| response_chn.poll_closed(cx)) => { continue; // Cancelled }, } } } async fn handle_msg( ws_reader: &mut SplitStream, write_tx: &mpsc::Sender, ) -> Option { loop { let read = ws_reader.next().await?; match read { Ok(Message::Binary(data)) => { let res = SignalResponse::decode(data.as_slice()) .expect("failed to decode SignalResponse"); event!(Level::TRACE, "received: {:?}", res); return Some(res.message.unwrap()); } Ok(Message::Ping(data)) => { let _ = write_tx.send(Pong { ping_data: data }); continue; } Ok(Message::Close(close)) => { event!(Level::DEBUG, "server closed the connection: {:?}", close); return None; } _ => { event!(Level::ERROR, "unhandled websocket message {:?}", read); return None; } } } } }