From bcfde797451763ae99cabe750f55cad8bad62c6d Mon Sep 17 00:00:00 2001 From: Zomatree Date: Tue, 17 Aug 2021 17:23:36 +0100 Subject: [PATCH 1/2] add msgpack support --- Cargo.lock | 31 ++- Cargo.toml | 2 + src/notifications/websocket.rs | 335 +++++++++++++++++++-------------- 3 files changed, 218 insertions(+), 150 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e5a1118a..85243f11 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1193,9 +1193,9 @@ dependencies = [ [[package]] name = "http" -version = "0.2.3" +version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7245cd7449cc792608c3c8a9eaf69bd4eabbabf802713748fd739c98b82f0747" +checksum = "527e8c9ac747e28542699a951517aa9a6945af506cd1f2e1b53a576c17b6cc11" dependencies = [ "bytes 1.0.1", "fnv", @@ -2728,12 +2728,14 @@ dependencies = [ "rauth", "regex", "reqwest 0.11.4", + "rmp-serde", "rocket", "rocket-governor", "rocket_cors", "serde", "serde_json", "ulid", + "url", "validator", "web-push", ] @@ -2753,6 +2755,27 @@ dependencies = [ "winapi 0.3.9", ] +[[package]] +name = "rmp" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4f55e5fa1446c4d5dd1f5daeed2a4fe193071771a2636274d0d7a3b082aa7ad6" +dependencies = [ + "byteorder", + "num-traits", +] + +[[package]] +name = "rmp-serde" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "723ecff9ad04f4ad92fe1c8ca6c20d2196d9286e9c60727c4cb5511629260e9d" +dependencies = [ + "byteorder", + "rmp", + "serde", +] + [[package]] name = "rocket" version = "0.5.0-rc.1" @@ -3831,9 +3854,9 @@ checksum = "a156c684c91ea7d62626509bce3cb4e1d9ed5c4d978f7b4352658f96a4c26b4a" [[package]] name = "url" -version = "2.2.0" +version = "2.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5909f2b0817350449ed73e8bcd81c8c3c8d9a7a5d8acba4b27db277f1868976e" +checksum = "a507c383b2d33b5fc35d1861e77e6b383d158b2da5e14fe51b83dfedf6fd578c" dependencies = [ "form_urlencoded", "idna", diff --git a/Cargo.toml b/Cargo.toml index 8564a6a9..c9be9731 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,6 +18,7 @@ once_cell = "1.4.1" env_logger = "0.7.1" lazy_static = "1.4.0" ctrlc = { version = "3.0", features = ["termination"] } +url = "2.2.2" # Lang. Utilities regex = "1" @@ -35,6 +36,7 @@ base64 = "0.13.0" serde_json = "1.0.57" serde = { version = "1.0.115", features = ["derive"] } validator = { version = "0.11", features = ["derive"] } +rmp-serde = "0.15.5" # async futures = "0.3.8" diff --git a/src/notifications/websocket.rs b/src/notifications/websocket.rs index 6e3f6dce..a63ea6d0 100644 --- a/src/notifications/websocket.rs +++ b/src/notifications/websocket.rs @@ -6,8 +6,8 @@ use super::subscriptions; use async_std::net::{TcpListener, TcpStream}; use async_std::task; -use async_tungstenite::tungstenite::Message; -use futures::channel::mpsc::{unbounded, UnboundedSender}; +use async_tungstenite::tungstenite::{Message, handshake::server}; +use futures::channel::{oneshot, mpsc::{unbounded, UnboundedSender}}; use futures::stream::TryStreamExt; use futures::{pin_mut, prelude::*}; use hive_pubsub::PubSub; @@ -18,6 +18,8 @@ use rauth::{ auth::{Auth}, options::Options, }; +use rmp_serde; +use url::Url; use std::collections::HashMap; use std::net::SocketAddr; use std::sync::{Arc, Mutex, RwLock}; @@ -46,14 +48,46 @@ pub async fn launch_server() { } } +#[derive(Debug)] +enum MSGFormat { + JSON, + MSGPACK +} + +struct HeaderCallback { + sender: oneshot::Sender +} + +impl server::Callback for HeaderCallback { + fn on_request(self, request: &server::Request, response: server::Response) -> Result { + // we dont get some of the data sometimes so im generating a fake url with the only data we actually need + let url = format!("ws://example.com?{}", request.uri().query().unwrap_or("?format=json")); + let mut query: HashMap<_, _> = url.parse::().unwrap().query_pairs().into_owned().collect(); // should be safe to use unwrap here as it was just a Uri + let format_query: Option = query.remove("format"); + + let format = match format_query.as_deref().unwrap_or("json") { + "msgpack" => MSGFormat::MSGPACK, + "json" => MSGFormat::JSON, + _ => panic!("unknown format") // TODO: not use panic + }; + + self.sender.send(format).unwrap(); // TODO: not use unwrap + Ok(response) + } +} + async fn accept(stream: TcpStream) { let addr = stream .peer_addr() .expect("Connected streams should have a peer address."); - let ws_stream = async_tungstenite::accept_async(stream) + let (sender, receiver) = oneshot::channel::(); + + let ws_stream = async_tungstenite::accept_hdr_async_with_config(stream, HeaderCallback { sender }, None) .await .expect("Error during websocket handshake."); + let msg_format = receiver.await.unwrap(); // TODO: not use unwrap + info!("User established WebSocket connection from {}.", &addr); let (write, read) = ws_stream.split(); @@ -61,10 +95,19 @@ async fn accept(stream: TcpStream) { CONNECTIONS.lock().unwrap().insert(addr, tx.clone()); let send = |notification: ClientboundNotification| { - if let Ok(response) = serde_json::to_string(¬ification) { - if let Err(_) = tx.unbounded_send(Message::Text(response)) { - debug!("Failed unbounded_send to websocket stream."); + let res = match msg_format { + MSGFormat::JSON => match serde_json::to_string(¬ification) { + Ok(s) => Message::Text(s), + Err(_) => return } + MSGFormat::MSGPACK => match rmp_serde::to_vec(¬ification) { + Ok(v) => Message::Binary(v), + Err(_) => return + } + }; + + if let Err(_) = tx.unbounded_send(res) { + debug!("Failed unbounded_send to websocket stream."); } }; @@ -73,187 +116,187 @@ async fn accept(stream: TcpStream) { let fwd = rx.map(Ok).forward(write); let incoming = read.try_for_each(async move |msg| { let mutex = mutex_generator(); - if let Message::Text(text) = msg { - let maybe_decoded = serde_json::from_str::(&text); - // If serde fails to decode the data, return a `MalformedData` error - if let Err(why) = maybe_decoded { + let maybe_decoded = match msg { + Message::Text(text) => serde_json::from_str::(&text).map_err(|e| e.to_string()), + Message::Binary(vec) => rmp_serde::decode::from_read::<&[u8], ServerboundNotification>(vec.as_slice()).map_err(|e| e.to_string()), + _ => return Ok(()) + }; + + let notification = match maybe_decoded { + Err(why) => { send(ClientboundNotification::Error( WebSocketError::MalformedData { msg: why.to_string() + })); + return Ok(()) + }, + Ok(n) => n + }; + + match notification { + ServerboundNotification::Authenticate(auth) => { + { + if mutex.lock().unwrap().is_some() { + send(ClientboundNotification::Error( + WebSocketError::AlreadyAuthenticated, + )); + + return Ok(()); } - )); + } - return Ok(()); - } - - if let Ok(notification) = maybe_decoded { - match notification { - ServerboundNotification::Authenticate(auth) => { + if let Some(id) = match auth { + AuthType::User(new_session) => { + if let Ok(validated_session) = + Auth::new(get_collection("accounts"), Options::new()) + .verify_session(new_session) + .await { - if mutex.lock().unwrap().is_some() { - send(ClientboundNotification::Error( - WebSocketError::AlreadyAuthenticated, - )); - - return Ok(()); - } + Some(validated_session.user_id.clone()) + } else { + None } - - if let Some(id) = match auth { - AuthType::User(new_session) => { - if let Ok(validated_session) = - Auth::new(get_collection("accounts"), Options::new()) - .verify_session(new_session) - .await - { - Some(validated_session.user_id.clone()) + } + AuthType::Bot(BotAuth { token }) => { + if let Ok(doc) = get_collection("bots") + .find_one( + doc! { "token": token }, + None + ).await { + if let Some(doc) = doc { + Some(doc.get_str("_id").unwrap().to_string()) } else { None } + } else { + None } - AuthType::Bot(BotAuth { token }) => { - if let Ok(doc) = get_collection("bots") - .find_one( - doc! { "token": token }, - None - ).await { - if let Some(doc) = doc { - Some(doc.get_str("_id").unwrap().to_string()) - } else { - None - } - } else { - None - } + } + } { + if let Ok(user) = (Ref { id: id.clone() }).fetch_user().await { + let is_invisible = if let Some(status) = &user.status { + if let Some(presence) = &status.presence { + presence == &Presence::Invisible + } else { + false } - } { - if let Ok(user) = (Ref { id: id.clone() }).fetch_user().await { - let is_invisible = if let Some(status) = &user.status { - if let Some(presence) = &status.presence { - presence == &Presence::Invisible - } else { - false - } - } else { - false - }; + } else { + false + }; - let was_online = is_online(&id); + let was_online = is_online(&id); - { - match USERS.write() { - Ok(mut map) => { - map.insert(id.clone(), addr); - } - Err(_) => { - send(ClientboundNotification::Error( - WebSocketError::InternalError { - at: "Writing users map.".to_string(), - }, - )); - - return Ok(()); - } - } + { + match USERS.write() { + Ok(mut map) => { + map.insert(id.clone(), addr); } - - *mutex.lock().unwrap() = Some(id.clone()); - - if let Err(_) = subscriptions::generate_subscriptions(&user).await { + Err(_) => { send(ClientboundNotification::Error( WebSocketError::InternalError { - at: "Generating subscriptions.".to_string(), + at: "Writing users map.".to_string(), }, )); return Ok(()); } + } + } - send(ClientboundNotification::Authenticated); + *mutex.lock().unwrap() = Some(id.clone()); - match super::payload::generate_ready(user).await { - Ok(payload) => { - send(payload); + if let Err(_) = subscriptions::generate_subscriptions(&user).await { + send(ClientboundNotification::Error( + WebSocketError::InternalError { + at: "Generating subscriptions.".to_string(), + }, + )); - if !was_online && !is_invisible { - ClientboundNotification::UserUpdate { - id: id.clone(), - data: json!({ - "online": true - }), - clear: None - } - .publish_as_user(id); - } - } - Err(_) => { - send(ClientboundNotification::Error( - WebSocketError::InternalError { - at: "Generating payload.".to_string(), - }, - )); - - return Ok(()); + return Ok(()); + } + + send(ClientboundNotification::Authenticated); + + match super::payload::generate_ready(user).await { + Ok(payload) => { + send(payload); + + if !was_online && !is_invisible { + ClientboundNotification::UserUpdate { + id: id.clone(), + data: json!({ + "online": true + }), + clear: None } + .publish_as_user(id); } - } else { + } + Err(_) => { send(ClientboundNotification::Error( - WebSocketError::OnboardingNotFinished, + WebSocketError::InternalError { + at: "Generating payload.".to_string(), + }, )); + + return Ok(()); } - } else { - send(ClientboundNotification::Error( - WebSocketError::InvalidSession, - )); } + } else { + send(ClientboundNotification::Error( + WebSocketError::OnboardingNotFinished, + )); } - // ! TEMP: verify user part of channel - // ! Could just run permission check here. - ServerboundNotification::BeginTyping { channel } => { - if mutex.lock().unwrap().is_some() { - let user = { - let mutex = mutex.lock().unwrap(); - mutex.as_ref().unwrap().clone() - }; + } else { + send(ClientboundNotification::Error( + WebSocketError::InvalidSession, + )); + } + } + // ! TEMP: verify user part of channel + // ! Could just run permission check here. + ServerboundNotification::BeginTyping { channel } => { + if mutex.lock().unwrap().is_some() { + let user = { + let mutex = mutex.lock().unwrap(); + mutex.as_ref().unwrap().clone() + }; - ClientboundNotification::ChannelStartTyping { - id: channel.clone(), - user, - } - .publish(channel); - } else { - send(ClientboundNotification::Error( - WebSocketError::AlreadyAuthenticated, - )); - - return Ok(()); - } + ClientboundNotification::ChannelStartTyping { + id: channel.clone(), + user, } - ServerboundNotification::EndTyping { channel } => { - if mutex.lock().unwrap().is_some() { - let user = { - let mutex = mutex.lock().unwrap(); - mutex.as_ref().unwrap().clone() - }; + .publish(channel); + } else { + send(ClientboundNotification::Error( + WebSocketError::AlreadyAuthenticated, + )); - ClientboundNotification::ChannelStopTyping { - id: channel.clone(), - user, - } - .publish(channel); - } else { - send(ClientboundNotification::Error( - WebSocketError::AlreadyAuthenticated, - )); + return Ok(()); + } + } + ServerboundNotification::EndTyping { channel } => { + if mutex.lock().unwrap().is_some() { + let user = { + let mutex = mutex.lock().unwrap(); + mutex.as_ref().unwrap().clone() + }; - return Ok(()); - } + ClientboundNotification::ChannelStopTyping { + id: channel.clone(), + user, } + .publish(channel); + } else { + send(ClientboundNotification::Error( + WebSocketError::AlreadyAuthenticated, + )); + + return Ok(()); } } } - Ok(()) }); From 867962cc3cd52b30ad66104214439277121d3765 Mon Sep 17 00:00:00 2001 From: Zomatree Date: Tue, 17 Aug 2021 17:35:05 +0100 Subject: [PATCH 2/2] reword comment --- src/notifications/websocket.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/notifications/websocket.rs b/src/notifications/websocket.rs index a63ea6d0..66322407 100644 --- a/src/notifications/websocket.rs +++ b/src/notifications/websocket.rs @@ -62,7 +62,7 @@ impl server::Callback for HeaderCallback { fn on_request(self, request: &server::Request, response: server::Response) -> Result { // we dont get some of the data sometimes so im generating a fake url with the only data we actually need let url = format!("ws://example.com?{}", request.uri().query().unwrap_or("?format=json")); - let mut query: HashMap<_, _> = url.parse::().unwrap().query_pairs().into_owned().collect(); // should be safe to use unwrap here as it was just a Uri + let mut query: HashMap<_, _> = url.parse::().unwrap().query_pairs().into_owned().collect(); // should be safe to use unwrap here as we just made the url ourself let format_query: Option = query.remove("format"); let format = match format_query.as_deref().unwrap_or("json") {