Migrate old ws crate to tungstenite, which is async and also removes over 20 old dependencies

Cargo.toml

@@ -45,7 +45,7 @@ backtrace = "0.3.65" # Logging panics to logfile instead stderr only
 dotenvy = { version = "0.15.1", default-features = false }
 
 # Lazy initialization
-once_cell = "1.10.0"
+once_cell = "1.11.0"
 
 # Numerical libraries
 num-traits = "0.2.15"
@@ -55,9 +55,9 @@ num-derive = "0.3.3"
 rocket = { version = "0.5.0-rc.2", features = ["tls", "json"], default-features = false }
 
 # WebSockets libraries
-ws = { version = "0.11.1", package = "parity-ws" }
+tokio-tungstenite = "0.17.1"
 rmpv = "1.0.0" # MessagePack library
-chashmap = "2.2.2" # Concurrent hashmap implementation
+dashmap = "5.3.3" # Concurrent hashmap implementation
 
 # Async futures
 futures = "0.3.21"
@@ -96,7 +96,7 @@ data-encoding = "2.3.2"
 jsonwebtoken = "8.1.0"
 
 # TOTP library
-totp-lite = "1.0.3"
+totp-lite = "2.0.0"
 
 # Yubico Library
 yubico = { version = "0.11.0", features = ["online-tokio"], default-features = false }
@@ -112,14 +112,14 @@ lettre = { version = "0.10.0-rc.7", features = ["smtp-transport", "builder", "se
 percent-encoding = "2.1.0" # URL encoding library used for URL's in the emails
 
 # Template library
-handlebars = { version = "4.2.2", features = ["dir_source"] }
+handlebars = { version = "4.3.0", features = ["dir_source"] }
 
 # HTTP client
 reqwest = { version = "0.11.10", features = ["stream", "json", "gzip", "brotli", "socks", "cookies", "trust-dns"] }
 
 # For favicon extraction from main website
 html5gum = "0.4.0"
-regex = { version = "1.5.5", features = ["std", "perf", "unicode-perl"], default-features = false }
+regex = { version = "1.5.6", features = ["std", "perf", "unicode-perl"], default-features = false }
 data-url = "0.1.1"
 bytes = "1.1.0"
 cached = "0.34.0"

src/api/core/ciphers.rs

@@ -464,7 +464,7 @@ pub async fn update_cipher_from_data(
     cipher.set_favorite(data.Favorite, &headers.user.uuid, conn).await?;
 
     if ut != UpdateType::None {
-        nt.send_cipher_update(ut, cipher, &cipher.update_users_revision(conn).await);
+        nt.send_cipher_update(ut, cipher, &cipher.update_users_revision(conn).await).await;
     }
 
     Ok(())
@@ -527,7 +527,7 @@ async fn post_ciphers_import(
 
     let mut user = headers.user;
     user.update_revision(&conn).await?;
-    nt.send_user_update(UpdateType::Vault, &user);
+    nt.send_user_update(UpdateType::Vault, &user).await;
     Ok(())
 }
 
@@ -1000,7 +1000,7 @@ async fn save_attachment(
 
     data.data.persist_to(file_path).await?;
 
-    nt.send_cipher_update(UpdateType::CipherUpdate, &cipher, &cipher.update_users_revision(&conn).await);
+    nt.send_cipher_update(UpdateType::CipherUpdate, &cipher, &cipher.update_users_revision(&conn).await).await;
 
     Ok((cipher, conn))
 }
@@ -1266,7 +1266,7 @@ async fn move_cipher_selected(
         // Move cipher
         cipher.move_to_folder(data.FolderId.clone(), &user_uuid, &conn).await?;
 
-        nt.send_cipher_update(UpdateType::CipherUpdate, &cipher, &[user_uuid.clone()]);
+        nt.send_cipher_update(UpdateType::CipherUpdate, &cipher, &[user_uuid.clone()]).await;
     }
 
     Ok(())
@@ -1313,7 +1313,7 @@ async fn delete_all(
                 Some(user_org) => {
                     if user_org.atype == UserOrgType::Owner {
                         Cipher::delete_all_by_organization(&org_data.org_id, &conn).await?;
-                        nt.send_user_update(UpdateType::Vault, &user);
+                        nt.send_user_update(UpdateType::Vault, &user).await;
                         Ok(())
                     } else {
                         err!("You don't have permission to purge the organization vault");
@@ -1334,7 +1334,7 @@ async fn delete_all(
             }
 
             user.update_revision(&conn).await?;
-            nt.send_user_update(UpdateType::Vault, &user);
+            nt.send_user_update(UpdateType::Vault, &user).await;
             Ok(())
         }
     }
@@ -1359,10 +1359,10 @@ async fn _delete_cipher_by_uuid(
     if soft_delete {
         cipher.deleted_at = Some(Utc::now().naive_utc());
         cipher.save(conn).await?;
-        nt.send_cipher_update(UpdateType::CipherUpdate, &cipher, &cipher.update_users_revision(conn).await);
+        nt.send_cipher_update(UpdateType::CipherUpdate, &cipher, &cipher.update_users_revision(conn).await).await;
     } else {
         cipher.delete(conn).await?;
-        nt.send_cipher_update(UpdateType::CipherDelete, &cipher, &cipher.update_users_revision(conn).await);
+        nt.send_cipher_update(UpdateType::CipherDelete, &cipher, &cipher.update_users_revision(conn).await).await;
     }
 
     Ok(())
@@ -1407,7 +1407,7 @@ async fn _restore_cipher_by_uuid(uuid: &str, headers: &Headers, conn: &DbConn, n
     cipher.deleted_at = None;
     cipher.save(conn).await?;
 
-    nt.send_cipher_update(UpdateType::CipherUpdate, &cipher, &cipher.update_users_revision(conn).await);
+    nt.send_cipher_update(UpdateType::CipherUpdate, &cipher, &cipher.update_users_revision(conn).await).await;
     Ok(Json(cipher.to_json(&headers.host, &headers.user.uuid, None, conn).await))
 }
 
@@ -1469,7 +1469,7 @@ async fn _delete_cipher_attachment_by_id(
 
     // Delete attachment
     attachment.delete(conn).await?;
-    nt.send_cipher_update(UpdateType::CipherUpdate, &cipher, &cipher.update_users_revision(conn).await);
+    nt.send_cipher_update(UpdateType::CipherUpdate, &cipher, &cipher.update_users_revision(conn).await).await;
     Ok(())
 }
 

src/api/core/folders.rs

@@ -50,7 +50,7 @@ async fn post_folders(data: JsonUpcase<FolderData>, headers: Headers, conn: DbCo
     let mut folder = Folder::new(headers.user.uuid, data.Name);
 
     folder.save(&conn).await?;
-    nt.send_folder_update(UpdateType::FolderCreate, &folder);
+    nt.send_folder_update(UpdateType::FolderCreate, &folder).await;
 
     Ok(Json(folder.to_json()))
 }
@@ -88,7 +88,7 @@ async fn put_folder(
     folder.name = data.Name;
 
     folder.save(&conn).await?;
-    nt.send_folder_update(UpdateType::FolderUpdate, &folder);
+    nt.send_folder_update(UpdateType::FolderUpdate, &folder).await;
 
     Ok(Json(folder.to_json()))
 }
@@ -112,6 +112,6 @@ async fn delete_folder(uuid: String, headers: Headers, conn: DbConn, nt: Notify<
     // Delete the actual folder entry
     folder.delete(&conn).await?;
 
-    nt.send_folder_update(UpdateType::FolderDelete, &folder);
+    nt.send_folder_update(UpdateType::FolderDelete, &folder).await;
     Ok(())
 }

src/api/core/sends.rs

@@ -173,7 +173,7 @@ async fn post_send(data: JsonUpcase<SendData>, headers: Headers, conn: DbConn, n
 
     let mut send = create_send(data, headers.user.uuid).await?;
     send.save(&conn).await?;
-    nt.send_send_update(UpdateType::SyncSendCreate, &send, &send.update_users_revision(&conn).await);
+    nt.send_send_update(UpdateType::SyncSendCreate, &send, &send.update_users_revision(&conn).await).await;
 
     Ok(Json(send.to_json()))
 }
@@ -237,7 +237,7 @@ async fn post_send_file(data: Form<UploadData<'_>>, headers: Headers, conn: DbCo
 
     // Save the changes in the database
     send.save(&conn).await?;
-    nt.send_send_update(UpdateType::SyncSendUpdate, &send, &send.update_users_revision(&conn).await);
+    nt.send_send_update(UpdateType::SyncSendUpdate, &send, &send.update_users_revision(&conn).await).await;
 
     Ok(Json(send.to_json()))
 }
@@ -418,7 +418,7 @@ async fn put_send(
     }
 
     send.save(&conn).await?;
-    nt.send_send_update(UpdateType::SyncSendUpdate, &send, &send.update_users_revision(&conn).await);
+    nt.send_send_update(UpdateType::SyncSendUpdate, &send, &send.update_users_revision(&conn).await).await;
 
     Ok(Json(send.to_json()))
 }
@@ -435,7 +435,7 @@ async fn delete_send(id: String, headers: Headers, conn: DbConn, nt: Notify<'_>)
     }
 
     send.delete(&conn).await?;
-    nt.send_send_update(UpdateType::SyncSendDelete, &send, &send.update_users_revision(&conn).await);
+    nt.send_send_update(UpdateType::SyncSendDelete, &send, &send.update_users_revision(&conn).await).await;
 
     Ok(())
 }
@@ -455,7 +455,7 @@ async fn put_remove_password(id: String, headers: Headers, conn: DbConn, nt: Not
 
     send.set_password(None);
     send.save(&conn).await?;
-    nt.send_send_update(UpdateType::SyncSendUpdate, &send, &send.update_users_revision(&conn).await);
+    nt.send_send_update(UpdateType::SyncSendUpdate, &send, &send.update_users_revision(&conn).await).await;
 
     Ok(Json(send.to_json()))
 }

src/api/notifications.rs

@@ -1,19 +1,41 @@
-use std::sync::atomic::{AtomicBool, Ordering};
+use std::{
+    net::SocketAddr,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
+    time::Duration,
+};
 
-use rocket::serde::json::Json;
-use rocket::Route;
+use chrono::NaiveDateTime;
+use futures::{SinkExt, StreamExt};
+use rmpv::Value;
+use rocket::{serde::json::Json, Route};
 use serde_json::Value as JsonValue;
+use tokio::{
+    net::{TcpListener, TcpStream},
+    sync::mpsc::Sender,
+};
+use tokio_tungstenite::{
+    accept_hdr_async,
+    tungstenite::{handshake, Message},
+};
 
-use crate::{api::EmptyResult, auth::Headers, Error, CONFIG};
+use crate::{
+    api::EmptyResult,
+    auth::Headers,
+    db::models::{Cipher, Folder, Send, User},
+    Error, CONFIG,
+};
 
 pub fn routes() -> Vec<Route> {
     routes![negotiate, websockets_err]
 }
 
-static SHOW_WEBSOCKETS_MSG: AtomicBool = AtomicBool::new(true);
-
 #[get("/hub")]
 fn websockets_err() -> EmptyResult {
+    static SHOW_WEBSOCKETS_MSG: AtomicBool = AtomicBool::new(true);
+
     if CONFIG.websocket_enabled()
         && SHOW_WEBSOCKETS_MSG.compare_exchange(true, false, Ordering::Relaxed, Ordering::Relaxed).is_ok()
     {
@@ -55,19 +77,6 @@ fn negotiate(_headers: Headers) -> Json<JsonValue> {
 //
 // Websockets server
 //
-use std::io;
-use std::sync::Arc;
-use std::thread;
-
-use ws::{self, util::Token, Factory, Handler, Handshake, Message, Sender};
-
-use chashmap::CHashMap;
-use chrono::NaiveDateTime;
-use serde_json::from_str;
-
-use crate::db::models::{Cipher, Folder, Send, User};
-
-use rmpv::Value;
 
 fn serialize(val: Value) -> Vec<u8> {
     use rmpv::encode::write_value;
@@ -118,192 +127,49 @@ fn convert_option<T: Into<Value>>(option: Option<T>) -> Value {
     }
 }
 
-// Server WebSocket handler
-pub struct WsHandler {
-    out: Sender,
-    user_uuid: Option<String>,
-    users: WebSocketUsers,
-}
-
 const RECORD_SEPARATOR: u8 = 0x1e;
 const INITIAL_RESPONSE: [u8; 3] = [0x7b, 0x7d, RECORD_SEPARATOR]; // {, }, <RS>
 
-#[derive(Deserialize)]
-struct InitialMessage {
-    protocol: String,
+#[derive(Deserialize, Copy, Clone, Eq, PartialEq)]
+struct InitialMessage<'a> {
+    protocol: &'a str,
     version: i32,
 }
 
-const PING_MS: u64 = 15_000;
-const PING: Token = Token(1);
-
-const ACCESS_TOKEN_KEY: &str = "access_token=";
-
-impl WsHandler {
-    fn err(&self, msg: &'static str) -> ws::Result<()> {
-        self.out.close(ws::CloseCode::Invalid)?;
-
-        // We need to specifically return an IO error so ws closes the connection
-        let io_error = io::Error::from(io::ErrorKind::InvalidData);
-        Err(ws::Error::new(ws::ErrorKind::Io(io_error), msg))
-    }
-
-    fn get_request_token(&self, hs: Handshake) -> Option<String> {
-        use std::str::from_utf8;
-
-        // Verify we have a token header
-        if let Some(header_value) = hs.request.header("Authorization") {
-            if let Ok(converted) = from_utf8(header_value) {
-                if let Some(token_part) = converted.split("Bearer ").nth(1) {
-                    return Some(token_part.into());
-                }
-            }
-        };
-
-        // Otherwise verify the query parameter value
-        let path = hs.request.resource();
-        if let Some(params) = path.split('?').nth(1) {
-            let params_iter = params.split('&').take(1);
-            for val in params_iter {
-                if let Some(stripped) = val.strip_prefix(ACCESS_TOKEN_KEY) {
-                    return Some(stripped.into());
-                }
-            }
-        };
-
-        None
-    }
-}
-
-impl Handler for WsHandler {
-    fn on_open(&mut self, hs: Handshake) -> ws::Result<()> {
-        // Path == "/notifications/hub?id=<id>==&access_token=<access_token>"
-        //
-        // We don't use `id`, and as of around 2020-03-25, the official clients
-        // no longer seem to pass `id` (only `access_token`).
-
-        // Get user token from header or query parameter
-        let access_token = match self.get_request_token(hs) {
-            Some(token) => token,
-            _ => return self.err("Missing access token"),
-        };
-
-        // Validate the user
-        use crate::auth;
-        let claims = match auth::decode_login(access_token.as_str()) {
-            Ok(claims) => claims,
-            Err(_) => return self.err("Invalid access token provided"),
-        };
-
-        // Assign the user to the handler
-        let user_uuid = claims.sub;
-        self.user_uuid = Some(user_uuid.clone());
-
-        // Add the current Sender to the user list
-        let handler_insert = self.out.clone();
-        let handler_update = self.out.clone();
-
-        self.users.map.upsert(user_uuid, || vec![handler_insert], |ref mut v| v.push(handler_update));
-
-        // Schedule a ping to keep the connection alive
-        self.out.timeout(PING_MS, PING)
-    }
-
-    fn on_message(&mut self, msg: Message) -> ws::Result<()> {
-        if let Message::Text(text) = msg.clone() {
-            let json = &text[..text.len() - 1]; // Remove last char
-
-            if let Ok(InitialMessage {
-                protocol,
-                version,
-            }) = from_str::<InitialMessage>(json)
-            {
-                if &protocol == "messagepack" && version == 1 {
-                    return self.out.send(&INITIAL_RESPONSE[..]); // Respond to initial message
-                }
-            }
-        }
-
-        // If it's not the initial message, just echo the message
-        self.out.send(msg)
-    }
-
-    fn on_timeout(&mut self, event: Token) -> ws::Result<()> {
-        if event == PING {
-            // send ping
-            self.out.send(create_ping())?;
-
-            // reschedule the timeout
-            self.out.timeout(PING_MS, PING)
-        } else {
-            Ok(())
-        }
-    }
-}
-
-struct WsFactory {
-    pub users: WebSocketUsers,
-}
-
-impl WsFactory {
-    pub fn init() -> Self {
-        WsFactory {
-            users: WebSocketUsers {
-                map: Arc::new(CHashMap::new()),
-            },
-        }
-    }
-}
-
-impl Factory for WsFactory {
-    type Handler = WsHandler;
-
-    fn connection_made(&mut self, out: Sender) -> Self::Handler {
-        WsHandler {
-            out,
-            user_uuid: None,
-            users: self.users.clone(),
-        }
-    }
-
-    fn connection_lost(&mut self, handler: Self::Handler) {
-        // Remove handler
-        if let Some(user_uuid) = &handler.user_uuid {
-            if let Some(mut user_conn) = self.users.map.get_mut(user_uuid) {
-                if let Some(pos) = user_conn.iter().position(|x| x == &handler.out) {
-                    user_conn.remove(pos);
-                }
-            }
-        }
-    }
-}
+static INITIAL_MESSAGE: InitialMessage<'static> = InitialMessage {
+    protocol: "messagepack",
+    version: 1,
+};
 
+// We attach the UUID to the sender so we can differentiate them when we need to remove them from the Vec
+type UserSenders = (uuid::Uuid, Sender<Message>);
 #[derive(Clone)]
 pub struct WebSocketUsers {
-    map: Arc<CHashMap<String, Vec<Sender>>>,
+    map: Arc<dashmap::DashMap<String, Vec<UserSenders>>>,
 }
 
 impl WebSocketUsers {
-    fn send_update(&self, user_uuid: &str, data: &[u8]) -> ws::Result<()> {
-        if let Some(user) = self.map.get(user_uuid) {
-            for sender in user.iter() {
-                sender.send(data)?;
+    async fn send_update(&self, user_uuid: &str, data: &[u8]) {
+        if let Some(user) = self.map.get(user_uuid).map(|v| v.clone()) {
+            for (_, sender) in user.iter() {
+                if sender.send(Message::binary(data)).await.is_err() {
+                    // TODO: Delete from map here too?
+                }
             }
         }
-        Ok(())
     }
 
     // NOTE: The last modified date needs to be updated before calling these methods
-    pub fn send_user_update(&self, ut: UpdateType, user: &User) {
+    pub async fn send_user_update(&self, ut: UpdateType, user: &User) {
         let data = create_update(
             vec![("UserId".into(), user.uuid.clone().into()), ("Date".into(), serialize_date(user.updated_at))],
             ut,
         );
 
-        self.send_update(&user.uuid, &data).ok();
+        self.send_update(&user.uuid, &data).await;
     }
 
-    pub fn send_folder_update(&self, ut: UpdateType, folder: &Folder) {
+    pub async fn send_folder_update(&self, ut: UpdateType, folder: &Folder) {
         let data = create_update(
             vec![
                 ("Id".into(), folder.uuid.clone().into()),
@@ -313,10 +179,10 @@ impl WebSocketUsers {
             ut,
         );
 
-        self.send_update(&folder.user_uuid, &data).ok();
+        self.send_update(&folder.user_uuid, &data).await;
     }
 
-    pub fn send_cipher_update(&self, ut: UpdateType, cipher: &Cipher, user_uuids: &[String]) {
+    pub async fn send_cipher_update(&self, ut: UpdateType, cipher: &Cipher, user_uuids: &[String]) {
         let user_uuid = convert_option(cipher.user_uuid.clone());
         let org_uuid = convert_option(cipher.organization_uuid.clone());
 
@@ -332,11 +198,11 @@ impl WebSocketUsers {
         );
 
         for uuid in user_uuids {
-            self.send_update(uuid, &data).ok();
+            self.send_update(uuid, &data).await;
         }
     }
 
-    pub fn send_send_update(&self, ut: UpdateType, send: &Send, user_uuids: &[String]) {
+    pub async fn send_send_update(&self, ut: UpdateType, send: &Send, user_uuids: &[String]) {
         let user_uuid = convert_option(send.user_uuid.clone());
 
         let data = create_update(
@@ -349,7 +215,7 @@ impl WebSocketUsers {
         );
 
         for uuid in user_uuids {
-            self.send_update(uuid, &data).ok();
+            self.send_update(uuid, &data).await;
         }
     }
 }
@@ -416,27 +282,129 @@ pub enum UpdateType {
     None = 100,
 }
 
-use rocket::State;
-pub type Notify<'a> = &'a State<WebSocketUsers>;
+pub type Notify<'a> = &'a rocket::State<WebSocketUsers>;
 
 pub fn start_notification_server() -> WebSocketUsers {
-    let factory = WsFactory::init();
-    let users = factory.users.clone();
+    let users = WebSocketUsers {
+        map: Arc::new(dashmap::DashMap::new()),
+    };
 
     if CONFIG.websocket_enabled() {
-        thread::spawn(move || {
-            let mut settings = ws::Settings::default();
-            settings.max_connections = 500;
-            settings.queue_size = 2;
-            settings.panic_on_internal = false;
+        let users2 = users.clone();
+        tokio::spawn(async move {
+            let addr = (CONFIG.websocket_address(), CONFIG.websocket_port());
+            let listener = TcpListener::bind(addr).await.expect("Can't listen on websocket port");
 
-            let ws = ws::Builder::new().with_settings(settings).build(factory).unwrap();
-            CONFIG.set_ws_shutdown_handle(ws.broadcaster());
-            ws.listen((CONFIG.websocket_address().as_str(), CONFIG.websocket_port())).unwrap();
+            let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel::<()>();
+            CONFIG.set_ws_shutdown_handle(shutdown_tx);
 
-            warn!("WS Server stopped!");
+            loop {
+                tokio::select! {
+                    Ok((stream, addr)) = listener.accept() => {
+                        tokio::spawn(handle_connection(stream, users2.clone(), addr));
+                    }
+
+                    _ = &mut shutdown_rx => {
+                        break;
+                    }
+                }
+            }
+
+            info!("Shutting down WebSockets server!")
         });
     }
 
     users
 }
+
+async fn handle_connection(stream: TcpStream, users: WebSocketUsers, _addr: SocketAddr) -> Result<(), Error> {
+    let mut user_uuid: Option<String> = None;
+
+    // Accept connection, do initial handshake, validate auth token and get the user ID
+    use handshake::server::{Request, Response};
+    let mut stream = accept_hdr_async(stream, |req: &Request, res: Response| {
+        if let Some(token) = get_request_token(req) {
+            if let Ok(claims) = crate::auth::decode_login(&token) {
+                user_uuid = Some(claims.sub);
+                return Ok(res);
+            }
+        }
+        Err(Response::builder().status(401).body(None).unwrap())
+    })
+    .await?;
+
+    let user_uuid = user_uuid.expect("User UUID should be set after the handshake");
+
+    // Add a channel to send messages to this client to the map
+    let entry_uuid = uuid::Uuid::new_v4();
+    let (tx, mut rx) = tokio::sync::mpsc::channel(100);
+    users.map.entry(user_uuid.clone()).or_default().push((entry_uuid, tx));
+
+    let mut interval = tokio::time::interval(Duration::from_secs(15));
+    loop {
+        tokio::select! {
+            res = stream.next() =>  {
+                match res {
+                    Some(Ok(message)) => {
+                        // We should receive an initial message with the protocol and version, and we will reply to it
+                        if let Message::Text(ref message) = message {
+                            let msg = message.strip_suffix('\u{1e}').unwrap_or(message);
+
+                            if serde_json::from_str(msg).ok() == Some(INITIAL_MESSAGE) {
+                                stream.send(Message::binary(INITIAL_RESPONSE)).await?;
+                                continue;
+                            }
+                        }
+
+                        // Just echo anything else the client sends
+                        if stream.send(message).await.is_err() {
+                            break;
+                        }
+                    }
+                    _ => break,
+                }
+            }
+
+            res = rx.recv() => {
+                match res {
+                    Some(res) => {
+                        if stream.send(res).await.is_err() {
+                            break;
+                        }
+                    },
+                    None => break,
+                }
+            }
+
+            _= interval.tick() => {
+                if stream.send(Message::Binary(create_ping())).await.is_err() {
+                    break;
+                }
+            }
+        }
+    }
+
+    //  Delete from map
+    users.map.entry(user_uuid).or_default().retain(|(uuid, _)| uuid != &entry_uuid);
+    Ok(())
+}
+
+fn get_request_token(req: &handshake::server::Request) -> Option<String> {
+    const ACCESS_TOKEN_KEY: &str = "access_token=";
+
+    if let Some(Ok(auth)) = req.headers().get("Authorization").map(|a| a.to_str()) {
+        if let Some(token_part) = auth.strip_prefix("Bearer ") {
+            return Some(token_part.to_owned());
+        }
+    }
+
+    if let Some(params) = req.uri().query() {
+        let params_iter = params.split('&').take(1);
+        for val in params_iter {
+            if let Some(stripped) = val.strip_prefix(ACCESS_TOKEN_KEY) {
+                return Some(stripped.to_owned());
+            }
+        }
+    }
+    None
+}

src/config.rs

@@ -37,7 +37,7 @@ macro_rules! make_config {
 
         struct Inner {
             rocket_shutdown_handle: Option<rocket::Shutdown>,
-            ws_shutdown_handle: Option<ws::Sender>,
+            ws_shutdown_handle: Option<tokio::sync::oneshot::Sender<()>>,
 
             templates: Handlebars<'static>,
             config: ConfigItems,
@@ -948,19 +948,17 @@ impl Config {
         self.inner.write().unwrap().rocket_shutdown_handle = Some(handle);
     }
 
-    pub fn set_ws_shutdown_handle(&self, handle: ws::Sender) {
+    pub fn set_ws_shutdown_handle(&self, handle: tokio::sync::oneshot::Sender<()>) {
         self.inner.write().unwrap().ws_shutdown_handle = Some(handle);
     }
 
     pub fn shutdown(&self) {
-        if let Ok(c) = self.inner.read() {
-            if let Some(handle) = c.ws_shutdown_handle.clone() {
-                handle.shutdown().ok();
+        if let Ok(mut c) = self.inner.write() {
+            if let Some(handle) = c.ws_shutdown_handle.take() {
+                handle.send(()).ok();
             }
-            // Wait a bit before stopping the web server
-            tokio::runtime::Handle::current()
-                .block_on(async move { tokio::time::sleep(tokio::time::Duration::from_secs(1)).await });
-            if let Some(handle) = c.rocket_shutdown_handle.clone() {
+
+            if let Some(handle) = c.rocket_shutdown_handle.take() {
                 handle.notify();
             }
         }

src/error.rs

@@ -49,6 +49,7 @@ use rocket::error::Error as RocketErr;
 use serde_json::{Error as SerdeErr, Value};
 use std::io::Error as IoErr;
 use std::time::SystemTimeError as TimeErr;
+use tokio_tungstenite::tungstenite::Error as TungstError;
 use webauthn_rs::error::WebauthnError as WebauthnErr;
 use yubico::yubicoerror::YubicoError as YubiErr;
 
@@ -88,6 +89,7 @@ make_error! {
     DieselCon(DieselConErr): _has_source, _api_error,
     DieselMig(DieselMigErr): _has_source, _api_error,
     Webauthn(WebauthnErr):   _has_source, _api_error,
+    WebSocket(TungstError):  _has_source, _api_error,
 }
 
 impl std::fmt::Debug for Error {