Merge pull request #82 from revoltchat/rauth-v1

Cargo.toml

@@ -11,6 +11,7 @@ edition = "2018"
 
 [dependencies]
 # Utility
+url = "2.2.2"
 log = "0.4.11"
 dotenv = "0.15.0"
 linkify = "0.6.0"
@@ -18,7 +19,6 @@ 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"
@@ -48,7 +48,7 @@ async-std = { version = "1.8.0", features = ["tokio1", "tokio02", "attributes"]
 web-push = "0.7.2"
 many-to-many = "0.1.2"
 lettre = "0.10.0-alpha.4"
-rauth = { git = "https://github.com/insertish/rauth", rev = "19bef300ecc269c028e3d746f0a2f09842d3b37b" }
+rauth = { git = "https://github.com/insertish/rauth", rev = "df88fa8ec723663908a210f9da2d192826cbe736" }
 hive_pubsub = { git = "https://gitlab.insrt.uk/insert/hive", rev = "a89826df2b30166220e68a6ed01a58b751456604", features = ["mongo"] }
 
 # web

assets/templates/reset.txt

@@ -0,0 +1,5 @@
+You requested a password reset, if you did not perform this action you can safely ignore this email.
+
+Reset your password here: {{url}}
+
+Sent by Revolt.

assets/templates/verify.txt

@@ -0,0 +1,6 @@
+You're almost there!
+If you did not perform this action you can safely ignore this email.
+
+Please verify your account here: {{url}}
+
+Sent by Revolt.

set_version.sh

@@ -1,3 +1,3 @@
 #!/bin/bash
-export version=0.5.2-alpha.2
+export version=0.5.3-alpha.0
 echo "pub const VERSION: &str = \"${version}\";" > src/version.rs

src/database/entities/message.rs

@@ -9,14 +9,12 @@ use futures::StreamExt;
 use mongodb::options::UpdateOptions;
 use mongodb::{
     bson::{doc, to_bson, DateTime},
-    options::FindOptions,
 };
+use rauth::entities::{Model, Session};
 use rocket::serde::json::Value;
 use serde::{Deserialize, Serialize};
 use ulid::Ulid;
-use web_push::{
+use web_push::{ContentEncoding, SubscriptionInfo, SubscriptionKeys, VapidSignatureBuilder, WebPushClient, WebPushMessageBuilder};
-    ContentEncoding, SubscriptionInfo, VapidSignatureBuilder, WebPushClient, WebPushMessageBuilder,
-};
 use std::time::SystemTime;
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -335,60 +333,47 @@ impl Message {
             }
 
             // Fetch their corresponding sessions.
-            if let Ok(mut cursor) = get_collection("accounts")
+            if target_ids.len() > 0 {
-                .find(
+                if let Ok(mut cursor) = Session::find(
-                    doc! {
+                        &get_db(),
-                        "_id": {
+                        doc! {
-                            "$in": target_ids
+                            "_id": {
-                        },
+                                "$in": target_ids
-                        "sessions.subscription": {
+                            },
-                            "$exists": true
+                            "subscription": {
-                        }
+                                "$exists": true
-                    },
-                    FindOptions::builder()
-                        .projection(doc! { "sessions": 1 })
-                        .build(),
-                )
-                .await
-            {
-                let mut subscriptions = vec![];
-                while let Some(result) = cursor.next().await {
-                    if let Ok(doc) = result {
-                        if let Ok(sessions) = doc.get_array("sessions") {
-                            for session in sessions {
-                                if let Some(doc) = session.as_document() {
-                                    if let Ok(sub) = doc.get_document("subscription") {
-                                        let endpoint = sub.get_str("endpoint").unwrap().to_string();
-                                        let p256dh = sub.get_str("p256dh").unwrap().to_string();
-                                        let auth = sub.get_str("auth").unwrap().to_string();
-
-                                        subscriptions
-                                            .push(SubscriptionInfo::new(endpoint, p256dh, auth));
-                                    }
-                                }
                             }
-                        }
+                        },
-                    }
+                        None
-                }
+                    )
-
+                    .await {
-                if subscriptions.len() > 0 {
                     let enc = serde_json::to_string(&PushNotification::new(self, &c_clone).await).unwrap();
                     let client = WebPushClient::new();
                     let key =
                         base64::decode_config(VAPID_PRIVATE_KEY.clone(), base64::URL_SAFE).unwrap();
 
-                    for subscription in subscriptions {
+                    while let Some(Ok(session)) = cursor.next().await {
-                        let mut builder = WebPushMessageBuilder::new(&subscription).unwrap();
+                        if let Some(sub) = session.subscription {
-                        let sig_builder = VapidSignatureBuilder::from_pem(
+                            let subscription = SubscriptionInfo {
-                            std::io::Cursor::new(&key),
+                                endpoint: sub.endpoint,
-                            &subscription,
+                                keys: SubscriptionKeys {
-                        )
+                                    auth: sub.auth,
-                        .unwrap();
+                                    p256dh: sub.p256dh
-                        let signature = sig_builder.build().unwrap();
+                                }
-                        builder.set_vapid_signature(signature);
+                            };
-                        builder.set_payload(ContentEncoding::AesGcm, enc.as_bytes());
+
-                        let m = builder.build().unwrap();
+                            let mut builder = WebPushMessageBuilder::new(&subscription).unwrap();
-                        client.send(m).await.ok();
+                            let sig_builder = VapidSignatureBuilder::from_pem(
+                                std::io::Cursor::new(&key),
+                                &subscription,
+                            )
+                            .unwrap();
+                            let signature = sig_builder.build().unwrap();
+                            builder.set_vapid_signature(signature);
+                            builder.set_payload(ContentEncoding::AesGcm, enc.as_bytes());
+                            let m = builder.build().unwrap();
+                            client.send(m).await.ok();
+                        }
                     }
                 }
             }
@@ -482,7 +467,7 @@ impl Message {
             let attachment_ids: Vec<String> =
                 attachments.iter().map(|f| f.id.to_string()).collect();
             get_collection("attachments")
-                .update_one(
+                .update_many(
                     doc! {
                         "_id": {
                             "$in": attachment_ids
@@ -497,7 +482,7 @@ impl Message {
                 )
                 .await
                 .map_err(|_| Error::DatabaseError {
-                    operation: "update_one",
+                    operation: "update_many",
                     with: "attachment",
                 })?;
         }

src/database/guards/user.rs

@@ -1,7 +1,7 @@
 use crate::database::*;
 
 use mongodb::bson::{doc, from_document};
-use rauth::auth::Session;
+use rauth::entities::Session;
 use rocket::http::Status;
 use rocket::request::{self, FromRequest, Outcome, Request};
 
@@ -15,7 +15,7 @@ impl<'r> FromRequest<'r> for User {
             .get("x-bot-token")
             .next()
             .map(|x| x.to_string());
-        
+
         if let Some(bot_token) = header_bot_token {
             return if let Ok(result) = get_collection("bots")
                 .find_one(
@@ -40,7 +40,10 @@ impl<'r> FromRequest<'r> for User {
                         if let Some(doc) = result {
                             Outcome::Success(from_document(doc).unwrap())
                         } else {
-                            Outcome::Failure((Status::Forbidden, rauth::util::Error::InvalidSession))
+                            Outcome::Failure((
+                                Status::Forbidden,
+                                rauth::util::Error::InvalidSession,
+                            ))
                         }
                     } else {
                         Outcome::Failure((
@@ -62,32 +65,34 @@ impl<'r> FromRequest<'r> for User {
                         with: "bot",
                     },
                 ))
-            }
+            };
         } else {
-            let session: Session = request.guard::<Session>().await.unwrap();
+            if let Outcome::Success(session) = request.guard::<Session>().await {
-
+                if let Ok(result) = get_collection("users")
-            if let Ok(result) = get_collection("users")
+                    .find_one(
-                .find_one(
+                        doc! {
-                    doc! {
+                            "_id": &session.user_id
-                        "_id": &session.user_id
+                        },
-                    },
+                        None,
-                    None,
+                    )
-                )
+                    .await
-                .await
+                {
-            {
+                    if let Some(doc) = result {
-                if let Some(doc) = result {
+                        Outcome::Success(from_document(doc).unwrap())
-                    Outcome::Success(from_document(doc).unwrap())
+                    } else {
+                        Outcome::Failure((Status::Forbidden, rauth::util::Error::InvalidSession))
+                    }
                 } else {
-                    Outcome::Failure((Status::Forbidden, rauth::util::Error::InvalidSession))
+                    Outcome::Failure((
+                        Status::InternalServerError,
+                        rauth::util::Error::DatabaseError {
+                            operation: "find_one",
+                            with: "user",
+                        },
+                    ))
                 }
             } else {
-                Outcome::Failure((
+                Outcome::Failure((Status::Forbidden, rauth::util::Error::InvalidSession))
-                    Status::InternalServerError,
-                    rauth::util::Error::DatabaseError {
-                        operation: "find_one",
-                        with: "user",
-                    },
-                ))
             }
         }
     }

src/database/migrations/init.rs

@@ -71,39 +71,6 @@ pub async fn create_database() {
     .await
     .expect("Failed to create pubsub collection.");
 
-    db.run_command(
-        doc! {
-            "createIndexes": "accounts",
-            "indexes": [
-                {
-                    "key": {
-                        "email": 1
-                    },
-                    "name": "email",
-                    "unique": true,
-                    "collation": {
-                        "locale": "en",
-                        "strength": 2
-                    }
-                },
-                {
-                    "key": {
-                        "email_normalised": 1
-                    },
-                    "name": "email_normalised",
-                    "unique": true,
-                    "collation": {
-                        "locale": "en",
-                        "strength": 2
-                    }
-                }
-            ]
-        },
-        None,
-    )
-    .await
-    .expect("Failed to create account index.");
-
     db.run_command(
         doc! {
             "createIndexes": "users",

src/database/migrations/mod.rs

@@ -10,10 +10,13 @@ pub async fn run_migrations() {
         .list_database_names(None, None)
         .await
         .expect("Failed to fetch database names.");
-
+    
     if list.iter().position(|x| x == "revolt").is_none() {
         init::create_database().await;
     } else {
         scripts::migrate_database().await;
     }
+    
+    // panic!("https://pbs.twimg.com/media/EDTpB5JWwAUvyxd.jpg");
+    rauth::entities::sync_models(&super::get_db()).await;
 }

src/database/migrations/scripts.rs

@@ -2,7 +2,7 @@ use crate::database::{permissions, get_collection, get_db, PermissionTuple};
 
 use futures::StreamExt;
 use log::info;
-use mongodb::{bson::{doc, from_document, to_document}, options::FindOptions};
+use mongodb::{bson::{Document, doc, from_bson, from_document, to_document}, options::FindOptions};
 use serde::{Deserialize, Serialize};
 
 #[derive(Serialize, Deserialize)]
@@ -11,7 +11,7 @@ struct MigrationInfo {
     revision: i32,
 }
 
-pub const LATEST_REVISION: i32 = 8;
+pub const LATEST_REVISION: i32 = 9;
 
 pub async fn migrate_database() {
     let migrations = get_collection("migrations");
@@ -212,6 +212,80 @@ pub async fn run_migrations(revision: i32) -> i32 {
             .expect("Failed to create bots collection.");
     }
 
+    if revision <= 8 {
+        info!("Running migration [revision 8 / 2021-09-10]: Update to rAuth version 1.");
+
+        get_db()
+            .run_command(
+                doc! {
+                    "dropIndexes": "accounts",
+                    "index": ["email", "email_normalised"]
+                },
+                None,
+            )
+            .await
+            .expect("Failed to delete legacy account indexes.");
+
+        let col = get_collection("sessions");
+        let mut cursor = get_collection("accounts")
+            .find(doc! { }, None)
+            .await
+            .unwrap();
+        
+        while let Some(doc) = cursor.next().await {
+            if let Ok(mut account) = doc {
+                if let Some(sessions) = account.remove("sessions") {
+                    #[derive(Deserialize)]
+                    struct Session {
+                        id: String,
+                        token: String,
+                        friendly_name: String,
+                        subscription: Option<Document>,
+                    }
+
+                    let sessions = from_bson::<Vec<Session>>(sessions).unwrap();
+                    for session in sessions {
+                        info!("Converting session {} to new format.", &session.id);
+
+                        let mut doc = doc! {
+                            "_id": session.id,
+                            "token": session.token,
+                            "name": session.friendly_name,
+                        };
+
+                        if let Some(sub) = session.subscription {
+                            doc.insert("subscription", sub);
+                        }
+
+                        col.insert_one(doc, None).await.ok();
+                    }
+                } else {
+                    info!("Account doesn't have any sessions!");
+                }
+            }
+        }
+
+        get_collection("accounts")
+            .update_many(
+                doc! { },
+                doc! {
+                    "$unset": {
+                        "sessions": 1,
+                    },
+                    "$set": {
+                        "mfa": {
+                            "recovery_codes": []
+                        }
+                    }
+                },
+                None
+            )
+            .await
+            .unwrap();
+    }
+
+    // Need to migrate fields on attachments, change `user_id`, `object_id`, etc to `parent`.
+
     // Reminder to update LATEST_REVISION when adding new migrations.
     LATEST_REVISION
 }

src/main.rs

@@ -21,20 +21,18 @@ pub mod util;
 pub mod version;
 
 use async_std::task;
-use chrono::Duration;
 use futures::join;
 use log::info;
-use rauth::options::{EmailVerification, Options, SMTP};
 use rauth::{
-    auth::Auth,
+    config::{Captcha, Config, EmailVerification, SMTPSettings, Template, Templates},
-    options::{Template, Templates},
+    logic::Auth,
 };
-use std::str::FromStr;
-use rocket_cors::AllowedOrigins;
 use rocket::catchers;
+use rocket_cors::AllowedOrigins;
+use std::str::FromStr;
 use util::variables::{
-    APP_URL, HCAPTCHA_KEY, INVITE_ONLY, PUBLIC_URL, SMTP_FROM, SMTP_HOST, SMTP_PASSWORD,
+    APP_URL, HCAPTCHA_KEY, INVITE_ONLY, SMTP_FROM, SMTP_HOST, SMTP_PASSWORD, SMTP_USERNAME,
-    SMTP_USERNAME, USE_EMAIL, USE_HCAPTCHA,
+    USE_EMAIL, USE_HCAPTCHA,
 };
 
 #[async_std::main]
@@ -71,79 +69,67 @@ async fn launch_web() {
     let cors = rocket_cors::CorsOptions {
         allowed_origins: AllowedOrigins::All,
         allowed_methods: [
-            "Get",
+            "Get", "Put", "Post", "Delete", "Options", "Head", "Trace", "Connect", "Patch",
-            "Put",
-            "Post",
-            "Delete",
-            "Options",
-            "Head",
-            "Trace",
-            "Connect",
-            "Patch",
         ]
-            .iter()
+        .iter()
-            .map(|s| FromStr::from_str(s).unwrap())
+        .map(|s| FromStr::from_str(s).unwrap())
-            .collect(),
+        .collect(),
         ..Default::default()
     }
     .to_cors()
     .expect("Failed to create CORS.");
 
-    let mut options = Options::new()
+    let mut config = Config {
-        .base_url(format!("{}/auth", *PUBLIC_URL))
+        email_verification: if *USE_EMAIL {
-        .email_verification(if *USE_EMAIL {
             EmailVerification::Enabled {
-                success_redirect_uri: format!("{}/login", *APP_URL),
+                smtp: SMTPSettings {
-                welcome_redirect_uri: format!("{}/welcome", *APP_URL),
-                password_reset_url: Some(format!("{}/login/reset", *APP_URL)),
-
-                verification_expiry: Duration::days(1),
-                password_reset_expiry: Duration::hours(1),
-
-                templates: Templates {
-                    verify_email: Template {
-                        title: "Verify your Revolt account.",
-                        text: "You're almost there!
-If you did not perform this action you can safely ignore this email.
-
-Please verify your account here: {{url}}",
-                        html: None,
-                    },
-                    reset_password: Template {
-                        title: "Reset your Revolt password.",
-                        text: "You requested a password reset, if you did not perform this action you can safely ignore this email.
-
-Reset your password here: {{url}}",
-                        html: None,
-                    },
-                    welcome: None,
-                },
-
-                smtp: SMTP {
                     from: (*SMTP_FROM).to_string(),
                     host: (*SMTP_HOST).to_string(),
                     username: (*SMTP_USERNAME).to_string(),
                     password: (*SMTP_PASSWORD).to_string(),
+                    reply_to: Some("support@revolt.chat".into()),
+                    port: None,
+                    use_tls: None,
+                },
+                expiry: Default::default(),
+                templates: Templates {
+                    verify: Template {
+                        title: "Verify your Revolt account.".into(),
+                        text: include_str!("../assets/templates/verify.txt").into(),
+                        url: format!("{}/login/verify/", *APP_URL),
+                        html: None,
+                    },
+                    reset: Template {
+                        title: "Reset your Revolt password.".into(),
+                        text: include_str!("../assets/templates/reset.txt").into(),
+                        url: format!("{}/login/reset/", *APP_URL),
+                        html: None,
+                    },
+                    welcome: None,
                 },
             }
         } else {
             EmailVerification::Disabled
-        });
+        },
+        ..Default::default()
+    };
 
     if *INVITE_ONLY {
-        options = options.invite_only_collection(database::get_collection("invites"))
+        config.invite_only = true;
     }
 
     if *USE_HCAPTCHA {
-        options = options.hcaptcha_secret(HCAPTCHA_KEY.clone());
+        config.captcha = Captcha::HCaptcha {
+            secret: HCAPTCHA_KEY.clone(),
+        };
     }
 
-    let auth = Auth::new(database::get_collection("accounts"), options);
+    let auth = Auth::new(database::get_db(), config);
     let rocket = rocket::build();
-
     routes::mount(rocket)
         .mount("/", rocket_cors::catch_all_options_routes())
-        .mount("/auth", rauth::routes::routes())
+        .mount("/auth/account", rauth::web::account::routes())
+        .mount("/auth/session", rauth::web::session::routes())
         .manage(auth)
         .manage(cors.clone())
         .attach(cors)

src/notifications/events.rs

@@ -1,6 +1,5 @@
 use hive_pubsub::PubSub;
 use mongodb::bson::doc;
-use rauth::auth::Session;
 use rocket::serde::json::Value;
 use serde::{Deserialize, Serialize};
 
@@ -19,24 +18,24 @@ pub enum WebSocketError {
 }
 
 #[derive(Deserialize, Debug)]
-pub struct BotAuth {
+pub struct Auth {
-    pub token: String
+    pub token: String,
 }
 
-#[derive(Deserialize, Debug)]
+#[derive(Serialize, Deserialize, Debug, Clone)]
 #[serde(untagged)]
-pub enum AuthType {
+pub enum Ping {
-    User(Session),
+    Binary(Vec<u8>),
-    Bot(BotAuth)
+    Number(usize)
 }
 
 #[derive(Deserialize, Debug)]
 #[serde(tag = "type")]
 pub enum ServerboundNotification {
-    Authenticate(AuthType),
+    Authenticate(Auth),
     BeginTyping { channel: String },
     EndTyping { channel: String },
-    Ping { data: Vec<u8> }
+    Ping { data: Ping, responded: Option<()> },
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -50,7 +49,7 @@ pub enum RemoveUserField {
 #[derive(Serialize, Deserialize, Debug, Clone)]
 pub enum RemoveChannelField {
     Icon,
-    Description
+    Description,
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -85,8 +84,9 @@ pub enum ClientboundNotification {
         users: Vec<User>,
         servers: Vec<Server>,
         channels: Vec<Channel>,
-        members: Vec<Member>
+        members: Vec<Member>,
     },
+    Pong { data: Ping },
 
     Message(Message),
     MessageUpdate {
@@ -159,11 +159,11 @@ pub enum ClientboundNotification {
         role_id: String,
         data: Value,
         #[serde(skip_serializing_if = "Option::is_none")]
-        clear: Option<RemoveRoleField>
+        clear: Option<RemoveRoleField>,
     },
     ServerRoleDelete {
         id: String,
-        role_id: String
+        role_id: String,
     },
 
     UserUpdate {
@@ -222,8 +222,7 @@ pub async fn prehandle_hook(notification: &ClientboundNotification) -> Result<()
                         subscribe_if_exists(recipient.clone(), channel_id.to_string()).ok();
                     }
                 }
-                Channel::TextChannel { server, .. }
+                Channel::TextChannel { server, .. } | Channel::VoiceChannel { server, .. } => {
-                | Channel::VoiceChannel { server, .. } => {
                     // ! FIXME: write a better algorithm?
                     let members = Server::fetch_member_ids(server).await?;
                     for member in members {

src/notifications/websocket.rs

@@ -1,28 +1,28 @@
 use crate::database::*;
-use crate::notifications::events::{AuthType, BotAuth};
+use crate::notifications::events::Ping;
 use crate::util::variables::WS_HOST;
 
 use super::subscriptions;
 
 use async_std::net::{TcpListener, TcpStream};
 use async_std::task;
-use async_tungstenite::tungstenite::{Message, handshake::server};
+use async_tungstenite::tungstenite::{handshake::server, Message};
-use futures::channel::{oneshot, mpsc::{unbounded, UnboundedSender}};
+use futures::channel::{
+    mpsc::{unbounded, UnboundedSender},
+    oneshot,
+};
 use futures::stream::TryStreamExt;
 use futures::{pin_mut, prelude::*};
 use hive_pubsub::PubSub;
 use log::{debug, info};
 use many_to_many::ManyToMany;
 use mongodb::bson::doc;
-use rauth::{
+use rauth::entities::{Model, Session};
-    auth::{Auth},
-    options::Options,
-};
 use rmp_serde;
-use url::Url;
 use std::collections::HashMap;
 use std::net::SocketAddr;
 use std::sync::{Arc, Mutex, RwLock};
+use url::Url;
 
 use super::{
     events::{ClientboundNotification, ServerboundNotification, WebSocketError},
@@ -51,28 +51,43 @@ pub async fn launch_server() {
 #[derive(Debug)]
 enum MSGFormat {
     JSON,
-    MSGPACK
+    MSGPACK,
 }
 
 struct HeaderCallback {
-    sender: oneshot::Sender<MSGFormat>
+    sender: oneshot::Sender<MSGFormat>,
 }
 
 impl server::Callback for HeaderCallback {
-    fn on_request(self, request: &server::Request, response: server::Response) -> Result<server::Response, server::ErrorResponse> {
+    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 url = format!(
-        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
+            "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
+            _ => MSGFormat::JSON, // Fallback to JSON.
         };
 
-        self.sender.send(format).unwrap();  // TODO: not use unwrap
+        if self.sender.send(format).is_ok() {
-        Ok(response)
+            Ok(response)
+        } else {
+            Err(server::ErrorResponse::new(None))
+        }
     }
 }
 
@@ -81,12 +96,13 @@ async fn accept(stream: TcpStream) {
         .peer_addr()
         .expect("Connected streams should have a peer address.");
     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
+    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);
 
@@ -98,12 +114,12 @@ async fn accept(stream: TcpStream) {
         let res = match msg_format {
             MSGFormat::JSON => match serde_json::to_string(&notification) {
                 Ok(s) => Message::Text(s),
-                Err(_) => return
+                Err(_) => return,
-            }
+            },
             MSGFormat::MSGPACK => match rmp_serde::to_vec(&notification) {
                 Ok(v) => Message::Binary(v),
-                Err(_) => return
+                Err(_) => return,
-            }
+            },
         };
 
         if let Err(_) = tx.unbounded_send(res) {
@@ -118,21 +134,27 @@ async fn accept(stream: TcpStream) {
         let mutex = mutex_generator();
 
         let maybe_decoded = match msg {
-            Message::Text(text) => serde_json::from_str::<ServerboundNotification>(&text).map_err(|e| e.to_string()),
+            Message::Text(text) => {
-            Message::Binary(vec) => rmp_serde::decode::from_read::<&[u8], ServerboundNotification>(vec.as_slice()).map_err(|e| e.to_string()),
+                serde_json::from_str::<ServerboundNotification>(&text).map_err(|e| e.to_string())
-            Message::Ping(vec) => Ok(ServerboundNotification::Ping { data: vec }),
+            }
-            _ => return Ok(())
+            Message::Binary(vec) => {
+                rmp_serde::decode::from_read::<&[u8], ServerboundNotification>(vec.as_slice())
+                    .map_err(|e| e.to_string())
+            }
+            Message::Ping(vec) => Ok(ServerboundNotification::Ping { data: Ping::Binary(vec), responded: Some(()) }),
+            _ => return Ok(()),
         };
 
         let notification = match maybe_decoded {
             Err(why) => {
                 send(ClientboundNotification::Error(
                     WebSocketError::MalformedData {
-                        msg: why.to_string()
+                        msg: why.to_string(),
-                }));
+                    },
-                return Ok(())
+                ));
-            },
+                return Ok(());
-            Ok(n) => n
+            }
+            Ok(n) => n,
         };
 
         match notification {
@@ -147,34 +169,20 @@ async fn accept(stream: TcpStream) {
                     }
                 }
 
-                if let Some(id) = match auth {
+                let id = if let Ok(Some(session)) =
-                    AuthType::User(new_session) => {
+                    Session::find_one(&get_db(), doc! { "token": &auth.token }, None).await
-                        if let Ok(validated_session) =
+                {
-                        Auth::new(get_collection("accounts"), Options::new())
+                    Some(session.user_id)
-                            .verify_session(new_session)
+                } else if let Ok(Some(bot)) = get_collection("bots")
-                            .await
+                    .find_one(doc! { "token": auth.token }, None)
-                        {
+                    .await
-                            Some(validated_session.user_id.clone())
+                {
-                        } else {
+                    Some(bot.get_str("_id").unwrap().to_string())
-                            None
+                } else {
-                        }
+                    None
-                    }
+                };
-                    AuthType::Bot(BotAuth { token }) => {
+
-                        if let Ok(doc) = get_collection("bots")
+                if let Some(id) = id {
-                            .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 {
@@ -229,7 +237,7 @@ async fn accept(stream: TcpStream) {
                                         data: json!({
                                             "online": true
                                         }),
-                                        clear: None
+                                        clear: None,
                                     }
                                     .publish_as_user(id);
                                 }
@@ -297,8 +305,12 @@ async fn accept(stream: TcpStream) {
                     return Ok(());
                 }
             }
-            ServerboundNotification::Ping { data } => {
+            ServerboundNotification::Ping { data, responded } => {
-                info!("Ping received from User {}. Payload: {:?}", &addr, data);
+                debug!("Ping received from connection {}. Payload: {:?}", &addr, data);
+
+                if responded.is_none() {
+                    send(ClientboundNotification::Pong { data });
+                }
             }
         }
         Ok(())
@@ -329,7 +341,7 @@ async fn accept(stream: TcpStream) {
             data: json!({
                 "online": false
             }),
-            clear: None
+            clear: None,
         }
         .publish_as_user(id);
     }

src/routes/mod.rs

@@ -1,7 +1,8 @@
-pub use rocket::response::Redirect;
 pub use rocket::http::Status;
+pub use rocket::response::Redirect;
 use rocket::{Build, Rocket};
 
+mod bots;
 mod channels;
 mod invites;
 mod onboard;
@@ -10,7 +11,6 @@ mod root;
 mod servers;
 mod sync;
 mod users;
-mod bots;
 
 pub fn mount(rocket: Rocket<Build>) -> Rocket<Build> {
     rocket

src/routes/onboard/complete.rs

@@ -2,7 +2,7 @@ use crate::database::*;
 use crate::util::result::{Error, Result, EmptyResponse};
 
 use mongodb::bson::doc;
-use rauth::auth::Session;
+use rauth::entities::Session;
 use regex::Regex;
 use rocket::serde::json::Json;
 use serde::{Deserialize, Serialize};

src/routes/onboard/hello.rs

@@ -1,6 +1,6 @@
 use crate::database::*;
 
-use rauth::auth::Session;
+use rauth::entities::Session;
 use rocket::serde::json::Value;
 
 #[get("/hello")]

src/routes/push/subscribe.rs

@@ -1,37 +1,19 @@
 use crate::database::*;
 use crate::util::result::{EmptyResponse, Error, Result};
 
-use mongodb::bson::{doc, to_document};
+use mongodb::bson::doc;
-use rauth::auth::Session;
+use rauth::entities::{Model, Session, WebPushSubscription};
 use rocket::serde::json::Json;
-use serde::{Deserialize, Serialize};
-
-#[derive(Serialize, Deserialize)]
-pub struct Subscription {
-    endpoint: String,
-    p256dh: String,
-    auth: String,
-}
 
 #[post("/subscribe", data = "<data>")]
-pub async fn req(session: Session, data: Json<Subscription>) -> Result<EmptyResponse> {
+pub async fn req(mut session: Session, data: Json<WebPushSubscription>) -> Result<EmptyResponse> {
-    let data = data.into_inner();
+    session.subscription = Some(data.into_inner());
-    get_collection("accounts")
+    session
-        .update_one(
+        .save(&get_db(), None)
-            doc! {
-                "_id": session.user_id,
-                "sessions.id": session.id.unwrap()
-            },
-            doc! {
-                "$set": {
-                    "sessions.$.subscription": to_document(&data)
-                        .map_err(|_| Error::DatabaseError { operation: "to_document", with: "subscription" })?
-                }
-            },
-            None,
-        )
         .await
-        .map_err(|_| Error::DatabaseError { operation: "update_one", with: "account" })?;
+        .map(|_| EmptyResponse)
-
+        .map_err(|_| Error::DatabaseError {
-    Ok(EmptyResponse {})
+            operation: "save",
+            with: "session",
+        })
 }

src/routes/push/unsubscribe.rs

@@ -1,29 +1,18 @@
 use crate::database::*;
-use crate::util::result::{Error, Result, EmptyResponse};
+use crate::util::result::{EmptyResponse, Error, Result};
 
 use mongodb::bson::doc;
-use rauth::auth::Session;
+use rauth::entities::{Model, Session};
 
 #[post("/unsubscribe")]
-pub async fn req(session: Session) -> Result<EmptyResponse> {
+pub async fn req(mut session: Session) -> Result<EmptyResponse> {
-    get_collection("accounts")
+    session.subscription = None;
-        .update_one(
+    session
-            doc! {
+        .save(&get_db(), None)
-                "_id": session.user_id,
-                "sessions.id": session.id.unwrap()
-            },
-            doc! {
-                "$unset": {
-                    "sessions.$.subscription": 1
-                }
-            },
-            None,
-        )
         .await
+        .map(|_| EmptyResponse)
         .map_err(|_| Error::DatabaseError {
-            operation: "to_document",
+            operation: "save",
-            with: "subscription",
+            with: "session",
-        })?;
+        })
-
-    Ok(EmptyResponse {})
 }

src/routes/root.rs

@@ -1,7 +1,11 @@
-use crate::util::{ratelimit::RateLimitGuard, variables::{
+use crate::util::{
-    APP_URL, AUTUMN_URL, EXTERNAL_WS_URL, HCAPTCHA_SITEKEY, INVITE_ONLY, JANUARY_URL, USE_AUTUMN,
+    ratelimit::RateLimitGuard,
-    USE_EMAIL, USE_HCAPTCHA, USE_JANUARY, USE_VOSO, VAPID_PUBLIC_KEY, VOSO_URL, VOSO_WS_HOST,
+    variables::{
-}};
+        APP_URL, AUTUMN_URL, EXTERNAL_WS_URL, HCAPTCHA_SITEKEY, INVITE_ONLY, JANUARY_URL,
+        USE_AUTUMN, USE_EMAIL, USE_HCAPTCHA, USE_JANUARY, USE_VOSO, VAPID_PUBLIC_KEY, VOSO_URL,
+        VOSO_WS_HOST,
+    },
+};
 
 use mongodb::bson::doc;
 use rocket::{http::Status, serde::json::Value};

src/routes/sync/get_settings.rs

@@ -14,9 +14,9 @@ pub struct Options {
 #[post("/settings/fetch", data = "<options>")]
 pub async fn req(user: User, options: Json<Options>) -> Result<Value> {
     if user.bot.is_some() {
-        return Err(Error::IsBot)
+        return Err(Error::IsBot);
     }
-    
+
     let options = options.into_inner();
     let mut projection = doc! {
         "_id": 0,

src/routes/sync/get_unreads.rs

@@ -7,8 +7,8 @@ use rocket::serde::json::Value;
 #[get("/unreads")]
 pub async fn req(user: User) -> Result<Value> {
     if user.bot.is_some() {
-        return Err(Error::IsBot)
+        return Err(Error::IsBot);
     }
-    
+
     Ok(json!(User::fetch_unreads(&user.id).await?))
 }

src/routes/sync/set_settings.rs

@@ -1,6 +1,6 @@
 use crate::database::*;
 use crate::notifications::events::ClientboundNotification;
-use crate::util::result::{Error, Result, EmptyResponse};
+use crate::util::result::{EmptyResponse, Error, Result};
 
 use chrono::prelude::*;
 use mongodb::bson::{doc, to_bson};
@@ -20,7 +20,7 @@ pub struct Options {
 #[post("/settings/set?<options..>", data = "<data>")]
 pub async fn req(user: User, data: Json<Data>, options: Options) -> Result<EmptyResponse> {
     if user.bot.is_some() {
-        return Err(Error::IsBot)
+        return Err(Error::IsBot);
     }
 
     let data = data.into_inner();

src/routes/users/change_username.rs

@@ -3,9 +3,8 @@ use crate::notifications::events::ClientboundNotification;
 use crate::util::result::{Error, Result, EmptyResponse};
 
 use mongodb::bson::doc;
-use rauth::auth::{Auth, Session};
+use rauth::entities::Account;
 use regex::Regex;
-use rocket::State;
 use rocket::serde::json::Json;
 use serde::{Deserialize, Serialize};
 use validator::Validate;
@@ -26,8 +25,7 @@ pub struct Data {
 
 #[patch("/<_ignore_id>/username", data = "<data>")]
 pub async fn req(
-    auth: &State<Auth>,
+    account: Account,
-    session: Session,
     user: User,
     data: Json<Data>,
     _ignore_id: String,
@@ -39,8 +37,7 @@ pub async fn req(
     data.validate()
         .map_err(|error| Error::FailedValidation { error })?;
 
-    auth.verify_password(&session, data.password.clone())
+    account.verify_password(&data.password)
-        .await
         .map_err(|_| Error::InvalidCredentials)?;
 
     let mut set = doc! {};

src/version.rs

@@ -1 +1 @@
-pub const VERSION: &str = "0.5.2-alpha.2";
+pub const VERSION: &str = "0.5.3-alpha.0";