feat(bonfire): add disconnection mechanism

feat(bonfire): handle session deletion and logout events
feat(core): trigger logout on bot token reset

Cargo.lock

@@ -3412,6 +3412,7 @@ version = "0.7.3"
 dependencies = [
  "async-std",
  "async-tungstenite",
+ "authifier",
  "bincode",
  "fred",
  "futures",

crates/bonfire/Cargo.toml

@@ -34,6 +34,7 @@ async-std = { version = "1.8.0", features = [
 ] }
 
 # core
+authifier = { version = "1.0.8" }
 revolt-result = { path = "../core/result" }
 revolt-models = { path = "../core/models" }
 revolt-config = { path = "../core/config" }

crates/bonfire/src/events/state.rs

@@ -58,6 +58,7 @@ impl Default for Cache {
 pub struct State {
     pub cache: Cache,
 
+    pub session_id: String,
     pub private_topic: String,
     subscribed: HashSet<String>,
     state: SubscriptionStateChange,
@@ -65,7 +66,7 @@ pub struct State {
 
 impl State {
     /// Create state from User
-    pub fn from(user: User) -> State {
+    pub fn from(user: User, session_id: String) -> State {
         let mut subscribed = HashSet::new();
         let private_topic = format!("{}!", user.id);
         subscribed.insert(private_topic.clone());
@@ -81,6 +82,7 @@ impl State {
         State {
             cache,
             subscribed,
+            session_id,
             private_topic,
             state: SubscriptionStateChange::Reset,
         }

crates/bonfire/src/websocket.rs

@@ -1,6 +1,7 @@
 use std::net::SocketAddr;
 
 use async_tungstenite::WebSocketStream;
+use authifier::AuthifierEvent;
 use fred::{
     interfaces::{ClientLike, EventInterface, PubsubInterface},
     types::RedisConfig,
@@ -73,17 +74,19 @@ pub async fn client(db: &'static Database, stream: TcpStream, addr: SocketAddr)
         write.send(config.encode(&create_error!(InvalidSession))).await.ok();
         return;
     };
-    let user = match User::from_token(db, token, UserHint::Any).await {
+
+    let (user, session_id) = match User::from_token(db, token, UserHint::Any).await {
         Ok(user) => user,
         Err(err) => {
             write.send(config.encode(&err)).await.ok();
             return;
         }
     };
+
     info!("User {addr:?} authenticated as @{}", user.username);
 
     // Create local state.
-    let mut state = State::from(user);
+    let mut state = State::from(user, session_id);
     let user_id = state.cache.user_id.clone();
 
     // Notify socket we have authenticated.
@@ -189,6 +192,7 @@ async fn listener(
         }) else {
             return;
         };
+
         let event = match *REDIS_PAYLOAD_TYPE {
             PayloadType::Json => message
                 .value
@@ -203,13 +207,34 @@ async fn listener(
                 .as_bytes()
                 .and_then(|b| bincode::deserialize::<EventV1>(b).ok()),
         };
+
         let Some(mut event) = event else {
             warn!("Failed to deserialise an event for {}!", message.channel);
             return;
         };
-        let should_send = state.handle_incoming_event_v1(db, &mut event).await;
-        if !should_send {
-            continue;
+
+        if let EventV1::Auth(auth) = &event {
+            if let AuthifierEvent::DeleteSession { session_id, .. } = auth {
+                if &state.session_id == session_id {
+                    event = EventV1::Logout;
+                }
+            } else if let AuthifierEvent::DeleteAllSessions {
+                exclude_session_id, ..
+            } = auth
+            {
+                if let Some(excluded) = exclude_session_id {
+                    if &state.session_id != excluded {
+                        event = EventV1::Logout;
+                    }
+                } else {
+                    event = EventV1::Logout;
+                }
+            }
+        } else {
+            let should_send = state.handle_incoming_event_v1(db, &mut event).await;
+            if !should_send {
+                continue;
+            }
         }
 
         let result = write.lock().await.send(config.encode(&event)).await;
@@ -218,6 +243,11 @@ async fn listener(
             if !matches!(e, Error::AlreadyClosed | Error::ConnectionClosed) {
                 warn!("Error while sending an event to {addr:?}: {e:?}");
             }
+
+            return;
+        }
+
+        if let EventV1::Logout = event {
             return;
         }
     }

crates/core/database/src/events/client.rs

@@ -48,6 +48,8 @@ pub enum EventV1 {
 
     /// Successfully authenticated
     Authenticated,
+    /// Logged out
+    Logout,
     /// Basic data to cache
     Ready {
         users: Vec<User>,

crates/core/database/src/models/bots/model.rs

@@ -2,7 +2,7 @@ use revolt_config::config;
 use revolt_result::Result;
 use ulid::Ulid;
 
-use crate::{BotInformation, Database, PartialUser, User};
+use crate::{events::client::EventV1, BotInformation, Database, PartialUser, User};
 
 auto_derived_partial!(
     /// Bot
@@ -142,6 +142,10 @@ impl Bot {
 
         db.update_bot(&self.id, &partial, remove).await?;
 
+        if partial.token.is_some() {
+            EventV1::Logout.private(self.id.clone()).await;
+        }
+
         self.apply_options(partial);
         Ok(())
     }

crates/core/database/src/models/users/model.rs

@@ -276,23 +276,27 @@ impl User {
         Ok(username)
     }
 
-    /// Find a user from a given token and hint
+    /// Find a user and session ID from a given token and hint
     #[async_recursion]
-    pub async fn from_token(db: &Database, token: &str, hint: UserHint) -> Result<User> {
+    pub async fn from_token(db: &Database, token: &str, hint: UserHint) -> Result<(User, String)> {
         match hint {
-            UserHint::Bot => {
+            UserHint::Bot => Ok((
                 db.fetch_user(
                     &db.fetch_bot_by_token(token)
                         .await
                         .map_err(|_| create_error!(InvalidSession))?
                         .id,
                 )
-                .await
+                .await?,
+                String::new(),
+            )),
+            UserHint::User => {
+                let session = db.fetch_session_by_token(token).await?;
+                Ok((db.fetch_user(&session.user_id).await?, session.id))
             }
-            UserHint::User => db.fetch_user_by_token(token).await,
             UserHint::Any => {
-                if let Ok(user) = User::from_token(db, token, UserHint::User).await {
-                    Ok(user)
+                if let Ok(result) = User::from_token(db, token, UserHint::User).await {
+                    Ok(result)
                 } else {
                     User::from_token(db, token, UserHint::Bot).await
                 }

crates/core/database/src/models/users/ops.rs

@@ -1,3 +1,4 @@
+use authifier::models::Session;
 use revolt_result::Result;
 
 use crate::{FieldsUser, PartialUser, RelationshipStatus, User};
@@ -16,8 +17,8 @@ pub trait AbstractUsers: Sync + Send {
     /// Fetch a user from the database by their username
     async fn fetch_user_by_username(&self, username: &str, discriminator: &str) -> Result<User>;
 
-    /// Fetch a user from the database by their session token
-    async fn fetch_user_by_token(&self, token: &str) -> Result<User>;
+    /// Fetch a session from the database by token
+    async fn fetch_session_by_token(&self, token: &str) -> Result<Session>;
 
     /// Fetch multiple users by their ids
     async fn fetch_users<'a>(&self, ids: &'a [String]) -> Result<Vec<User>>;

crates/core/database/src/models/users/ops/mongodb.rs

@@ -46,10 +46,9 @@ impl AbstractUsers for MongoDb {
         .ok_or_else(|| create_error!(NotFound))
     }
 
-    /// Fetch a user from the database by their session token
-    async fn fetch_user_by_token(&self, token: &str) -> Result<User> {
-        let session = self
-            .col::<Session>("sessions")
+    /// Fetch a session from the database by token
+    async fn fetch_session_by_token(&self, token: &str) -> Result<Session> {
+        self.col::<Session>("sessions")
             .find_one(
                 doc! {
                     "token": token
@@ -58,9 +57,7 @@ impl AbstractUsers for MongoDb {
             )
             .await
             .map_err(|_| create_database_error!("find_one", "sessions"))?
-            .ok_or_else(|| create_error!(InvalidSession))?;
-
-        self.fetch_user(&session.user_id).await
+            .ok_or_else(|| create_error!(InvalidSession))
     }
 
     /// Fetch multiple users by their ids

crates/core/database/src/models/users/ops/reference.rs

@@ -1,3 +1,4 @@
+use authifier::models::Session;
 use revolt_result::Result;
 
 use crate::{FieldsUser, PartialUser, RelationshipStatus, User};
@@ -40,8 +41,8 @@ impl AbstractUsers for ReferenceDb {
             .ok_or_else(|| create_error!(NotFound))
     }
 
-    /// Fetch a user from the database by their session token
-    async fn fetch_user_by_token(&self, _token: &str) -> Result<User> {
+    /// Fetch a session from the database by token
+    async fn fetch_session_by_token(&self, _token: &str) -> Result<Session> {
         todo!()
     }