Merge pull request #42 from Zomatree/msgpack

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",

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"

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 async_tungstenite::tungstenite::{Message, handshake::server};
-use futures::channel::mpsc::{unbounded, UnboundedSender};
+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<MSGFormat>
+}
+
+impl server::Callback for HeaderCallback {
+    fn on_request(self, request: &server::Request, response: server::Response) -> Result<server::Response, server::ErrorResponse> {
+        // 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::<Url>().unwrap().query_pairs().into_owned().collect();  // should be safe to use unwrap here as we just made the url ourself
+        let format_query: Option<String> = 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::<MSGFormat>();
+    
+    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(&notification) {
+        let res = match msg_format {
-            if let Err(_) = tx.unbounded_send(Message::Text(response)) {
+            MSGFormat::JSON => match serde_json::to_string(&notification) {
-                debug!("Failed unbounded_send to websocket stream.");
+                Ok(s) => Message::Text(s),
+                Err(_) => return
             }
+            MSGFormat::MSGPACK => match rmp_serde::to_vec(&notification) {
+                Ok(v) => Message::Binary(v),
+                Err(_) => return
+            }
+        };
+
+        if let Err(_) = tx.unbounded_send(res) {
+            debug!("Failed unbounded_send to websocket stream.");
         }
     };
 
@@ -73,21 +116,24 @@ 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::<ServerboundNotification>(&text);
 
-            // If serde fails to decode the data, return a `MalformedData` error
+        let maybe_decoded = match msg {
-            if let Err(why) = maybe_decoded {
+            Message::Text(text) => serde_json::from_str::<ServerboundNotification>(&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
+        };
 
-                return Ok(());
-            }
-
-            if let Ok(notification) = maybe_decoded {
         match notification {
             ServerboundNotification::Authenticate(auth) => {
                 {
@@ -251,9 +297,6 @@ async fn accept(stream: TcpStream) {
                 }
             }
         }
-            }
-        }
-
         Ok(())
     });