feat: initial work on refresh tokens

crates/core/database/src/drivers/reference.rs

@@ -5,13 +5,14 @@ use futures::lock::Mutex;
 use crate::{
     Bot, Channel, ChannelCompositeKey, ChannelUnread, Emoji, File, FileHash, Invite, Member,
     MemberCompositeKey, Message, PolicyChange, RatelimitEvent, Report, Server, ServerBan, Snapshot,
-    User, UserSettings, Webhook,
+    User, UserSettings, Webhook, AuthorizedBotId, AuthorizedBot,
 };
 
 database_derived!(
     /// Reference implementation
     #[derive(Default)]
     pub struct ReferenceDb {
+        pub authorized_bots: Arc<Mutex<HashMap<AuthorizedBotId, AuthorizedBot>>>,
         pub bots: Arc<Mutex<HashMap<String, Bot>>>,
         pub channels: Arc<Mutex<HashMap<String, Channel>>>,
         pub channel_invites: Arc<Mutex<HashMap<String, Invite>>>,

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

@@ -2,6 +2,7 @@ use iso8601_timestamp::Timestamp;
 
 auto_derived! (
     /// Unique id of the user and bot
+    #[derive(Hash)]
     pub struct AuthorizedBotId {
         /// User id
         pub user: String,

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

@@ -8,20 +8,20 @@ mod reference;
 #[async_trait]
 pub trait AbstractAuthorizedBots: Sync + Send {
     /// Insert emoji into database.
-    async fn insert_authorized_bot(&self, authorised_bot: &AuthorizedBot) -> Result<()>;
+    async fn insert_authorized_bot(&self, authorized_bot: &AuthorizedBot) -> Result<()>;
 
-    /// Fetch an emoji by its id
+    /// Fetch an authorized bot by its id
     async fn fetch_authorized_bot(&self, id: &AuthorizedBotId) -> Result<AuthorizedBot>;
 
-    /// Fetch all authorized bots for a user
+    /// Fetch a users authorized bot by its id
     async fn fetch_users_authorized_bots(&self, user_id: &str) -> Result<Vec<AuthorizedBot>>;
 
     /// Deletes an authorized bot
     async fn delete_authorized_bot(&self, id: &AuthorizedBotId) -> Result<()>;
 
-    /// Deauthorizes a bot access to a user's information
+    /// Deauthorizes an authorized bot
     async fn deauthorize_authorized_bot(&self, id: &AuthorizedBotId) -> Result<AuthorizedBot>;
 
-    /// Fetches all deauthorized bots
+    // Fetches all authorized bots which have been deauthorized
     async fn fetch_deauthorized_authorized_bots(&self) -> Result<Vec<AuthorizedBot>>;
 }

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

@@ -1,4 +1,4 @@
-use bson::{Document, to_bson};
+use bson::to_bson;
 use revolt_result::Result;
 use iso8601_timestamp::Timestamp;
 
@@ -10,26 +10,27 @@ static COL: &str = "authorized_bots";
 
 #[async_trait]
 impl AbstractAuthorizedBots for MongoDb {
-    /// Insert emoji into database.
+    /// Insert an authorized bot into database.
-    async fn insert_authorized_bot(&self, authorised_bot: &AuthorizedBot) -> Result<()> {
+    async fn insert_authorized_bot(&self, authorized_bot: &AuthorizedBot) -> Result<()> {
-        query!(self, insert_one, COL, &authorised_bot).map(|_| ())
+        query!(self, insert_one, COL, &authorized_bot).map(|_| ())
     }
 
-    /// Fetch an emoji by its id
+    /// Fetch an authorized bot by its id
     async fn fetch_authorized_bot(&self, id: &AuthorizedBotId) -> Result<AuthorizedBot> {
         query!(self, find_one, COL, doc! { "id.user": &id.user, "id.bot": &id.bot })?.ok_or_else(|| create_error!(NotFound))
     }
 
-    /// Fetch an emoji by its id
+    /// Fetch a users authorized bot by its id
     async fn fetch_users_authorized_bots(&self, user_id: &str) -> Result<Vec<AuthorizedBot>> {
         query!(self, find, COL, doc! { "id.user": &user_id })
     }
 
-    /// Deletes an authori
+    /// Deletes an authorized bot
     async fn delete_authorized_bot(&self, id: &AuthorizedBotId) -> Result<()> {
         query!(self, delete_one, COL, doc! { "id.user": &id.user, "id.bot": &id.bot }).map(|_| ())
     }
 
+    /// Deauthorizes an authorized bot
     async fn deauthorize_authorized_bot(&self, id: &AuthorizedBotId) -> Result<AuthorizedBot> {
         self.col::<AuthorizedBot>(COL)
             .find_one_and_update(
@@ -39,7 +40,7 @@ impl AbstractAuthorizedBots for MongoDb {
                 },
                 doc! {
                     "$set": {
-                    "unauthorized_at": to_bson(&Timestamp::now_utc()).unwrap()
+                        "deauthorized_at": to_bson(&Timestamp::now_utc()).unwrap()
                     }
                 }
             )
@@ -48,6 +49,7 @@ impl AbstractAuthorizedBots for MongoDb {
             .and_then(|opt| opt.ok_or_else(|| create_database_error!("find_one_and_update", COL)))
     }
 
+    // Fetches all authorized bots which have been deauthorized
     async fn fetch_deauthorized_authorized_bots(&self) -> Result<Vec<AuthorizedBot>> {
         query!(
             self,

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

@@ -1,38 +1,76 @@
-use bson::Document;
 use revolt_result::Result;
+use iso8601_timestamp::Timestamp;
 
 use crate::{ReferenceDb, AuthorizedBot, AuthorizedBotId};
 
 use super::AbstractAuthorizedBots;
 
-static COL: &str = "authorized_bots";
-
 #[async_trait]
 impl AbstractAuthorizedBots for ReferenceDb {
-        /// Insert emoji into database.
+    /// Insert an authorized bot into database.
-    async fn insert_authorized_bot(&self, authorised_bot: &AuthorizedBot) -> Result<()> {
+    async fn insert_authorized_bot(&self, authorized_bot: &AuthorizedBot) -> Result<()> {
-        todo!()
+        let mut authorized_bots = self.authorized_bots.lock().await;
+
+        if authorized_bots.contains_key(&authorized_bot.id) {
+            Err(create_database_error!("insert", "authorized_bots"))
+        } else {
+            authorized_bots.insert(authorized_bot.id.clone(), authorized_bot.clone());
+            Ok(())
+        }
     }
 
-    /// Fetch an emoji by its id
+    /// Fetch an authorized bot by its id
     async fn fetch_authorized_bot(&self, id: &AuthorizedBotId) -> Result<AuthorizedBot> {
-        todo!()
+        let authorized_bots = self.authorized_bots.lock().await;
+
+        authorized_bots.get(id).cloned().ok_or_else(|| create_error!(NotFound))
     }
 
+    /// Fetch a users authorized bot by its id
     async fn fetch_users_authorized_bots(&self, user_id: &str) -> Result<Vec<AuthorizedBot>> {
-        todo!()
+        let authorized_bots = self.authorized_bots.lock().await;
+
+        Ok(authorized_bots
+            .values()
+            .filter(|authorized_bot| authorized_bot.id.user == user_id)
+            .cloned()
+            .collect()
+        )
     }
 
-    /// Deletes an authori
+    /// Deletes an authorized bot
     async fn delete_authorized_bot(&self, id: &AuthorizedBotId) -> Result<()> {
-        todo!()
+        let mut authorized_bots = self.authorized_bots.lock().await;
+
+        if authorized_bots.remove(id).is_some() {
+            Ok(())
+        } else {
+            Err(create_error!(NotFound))
+        }
     }
 
+    /// Deauthorizes an authorized bot
     async fn deauthorize_authorized_bot(&self, id: &AuthorizedBotId) -> Result<AuthorizedBot> {
-        todo!()
+        let mut authorized_bots = self.authorized_bots.lock().await;
+
+        if let Some(authorized_bot) = authorized_bots.get_mut(id) {
+            authorized_bot.deauthorized_at = Some(Timestamp::now_utc());
+
+            Ok(authorized_bot.clone())
+        } else {
+            Err(create_error!(NotFound))
+        }
     }
 
+    // Fetches all authorized bots which have been deauthorized
     async fn fetch_deauthorized_authorized_bots(&self) -> Result<Vec<AuthorizedBot>> {
-        todo!()
+        let authorized_bots = self.authorized_bots.lock().await;
+
+        Ok(authorized_bots
+            .values()
+            .filter(|authorized_bot| authorized_bot.deauthorized_at.is_some())
+            .cloned()
+            .collect()
+        )
     }
 }

crates/core/models/src/v0/authorized_bots.rs

@@ -1,5 +1,7 @@
 use iso8601_timestamp::Timestamp;
 
+use crate::v0::PublicBot;
+
 auto_derived!(
     /// Unique id of the user and bot
     pub struct AuthorizedBotId {
@@ -24,4 +26,9 @@ auto_derived!(
         /// Scopes the bot has access to
         pub scope: String
     }
+
+    pub struct AuthorizedBotsResponse {
+        pub public_bot: PublicBot,
+        pub authorized_bot: AuthorizedBot
+    }
 );

crates/core/models/src/v0/oauth2.rs

@@ -37,6 +37,9 @@ auto_derived!(
         #[cfg_attr(feature = "rocket", field(value = "implicit"))]
         #[cfg_attr(feature = "serde", serde(rename = "implicit"))]
         Implicit,
+        #[cfg_attr(feature = "rocket", field(value = "refresh_token"))]
+        #[cfg_attr(feature = "serde", serde(rename = "refresh_token"))]
+        RefreshToken
     }
 
     #[derive(Copy)]
@@ -66,12 +69,17 @@ auto_derived!(
         pub client_id: String,
         pub client_secret: Option<String>,
 
-        pub code: String,
+        /// Authorization code
+        pub code: Option<String>,
+        // Refresh token to generate new access token
+        pub refresh_token: Option<String>,
+
         pub code_verifier: Option<String>,
     }
 
     pub struct OAuth2TokenExchangeResponse {
         pub access_token: String,
+        pub refresh_token: Option<String>,
         pub token_type: String,
         pub scope: String,
     }

crates/core/result/src/lib.rs

@@ -51,7 +51,7 @@ impl std::error::Error for Error {}
 #[cfg_attr(feature = "serde", serde(tag = "type"))]
 #[cfg_attr(feature = "schemas", derive(JsonSchema))]
 #[cfg_attr(feature = "utoipa", derive(ToSchema))]
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Eq, PartialEq)]
 pub enum ErrorType {
     /// This error was not labeled :(
     LabelMe,

crates/daemons/crond/src/tasks/prune_authorized_bots.rs

@@ -2,8 +2,6 @@ use revolt_database::{iso8601_timestamp::{Duration, Timestamp}, util::oauth2::To
 use revolt_result::Result;
 use tokio::time::sleep;
 
-use log::info;
-
 pub async fn task(db: Database) -> Result<()> {
     let lifetime = Duration::new(TokenType::Access.lifetime().num_seconds(), 0);
 

crates/delta/src/routes/oauth2/authorize_auth.rs

@@ -5,8 +5,7 @@ use revolt_database::{
 };
 use revolt_models::v0;
 use revolt_result::{create_error, Result};
-use rocket::{form::Form, serde::json::Json, State};
+use rocket::{serde::json::Json, State};
-use redis_kiss::AsyncCommands;
 
 /// # OAuth2 Authorization Auth
 ///

crates/delta/src/routes/oauth2/authorize_info.rs

@@ -1,7 +1,7 @@
 use revolt_database::{util::reference::Reference, Database, User};
 use revolt_models::v0;
 use revolt_result::{create_error, Result};
-use rocket::{form::Form, serde::json::Json, State};
+use rocket::{serde::json::Json, State};
 
 /// # Authorize OAuth Information
 ///

crates/delta/src/routes/oauth2/authorized_bots.rs

@@ -3,18 +3,25 @@ use rocket::{serde::json::Json, State};
 use revolt_database::{Database, User};
 use revolt_result::Result;
 
-
 #[openapi(tag = "OAuth2")]
 #[get("/authorized_bots")]
 pub async fn authorized_bots(
     db: &State<Database>,
     user: User
-) -> Result<Json<Vec<v0::AuthorizedBot>>> {
+) -> Result<Json<Vec<v0::AuthorizedBotsResponse>>> {
     let authorized_bots = db.fetch_users_authorized_bots(&user.id).await?;
 
-    Ok(Json(authorized_bots
+    let mut response = Vec::new();
-        .into_iter()
+
-        .map(|bot| bot.into())
+    for authorized_bot in authorized_bots {
-        .collect()
+        let bot = db.fetch_bot(&authorized_bot.id.bot).await?;
-    ))
+        let bot_user = db.fetch_user(&authorized_bot.id.bot).await?;
+
+        response.push(v0::AuthorizedBotsResponse {
+            authorized_bot: authorized_bot.into(),
+            public_bot: bot.into_public_bot(bot_user)
+        });
+    };
+
+    Ok(Json(response))
 }

crates/delta/src/routes/oauth2/token.rs

@@ -8,9 +8,8 @@ use revolt_database::{
     }, AuthorizedBot, AuthorizedBotId, Database
 };
 use revolt_models::v0;
-use revolt_result::{create_error, Result};
+use revolt_result::{create_error, ErrorType, Result};
 use rocket::{form::Form, serde::json::Json, State};
-use redis_kiss::AsyncCommands;
 
 /// # OAuth Token Exchange
 ///
@@ -27,15 +26,47 @@ pub async fn token(
 
     let config = config().await;
 
-    let claims = oauth2::decode_token(&config.api.security.token_secret, &info.code)
+    let (token, token_type) = match (&info.code, &info.refresh_token) {
+        (Some(code), None) => {
+            (code, oauth2::TokenType::Auth)
+        },
+        (None, Some(refresh_token)) => {
+            (refresh_token, oauth2::TokenType::Refresh)
+        },
+        (Some(_), Some(_)) | (None, None) => {
+            return Err(create_error!(InvalidOperation))
+        }
+    };
+
+    let claims = oauth2::decode_token(&config.api.security.token_secret, token)
         .map_err(|_| create_error!(NotAuthenticated))?;
 
+    if claims.r#type != token_type || claims.client_id != info.client_id || claims.exp < Utc::now().timestamp() {
+        return Err(create_error!(NotAuthenticated))
+    };
+
+    let authorized_bot = db.fetch_authorized_bot(&AuthorizedBotId { user: claims.sub.clone(), bot: claims.client_id.clone() }).await?;
+
+    if authorized_bot.deauthorized_at.is_some() {
+        return Err(create_error!(NotAuthenticated));
+    };
+
+    // TODO: track used tokens and dont allow them to be used twice
+    //  - refresh token
+    //  - auth tokens (only last 5 mins but still should be considered)
+
+    match info.grant_type {
+        v0::OAuth2GrantType::AuthorizationCode => {
+            if claims.r#type != oauth2::TokenType::Auth {
+                return Err(create_error!(InvalidOperation))
+            }
+
             if let Some((client_secret, bot_secret)) = info.client_secret.as_ref().zip(bot.oauth2.as_ref().and_then(|oauth2| oauth2.secret.as_ref())) {
                 if client_secret != bot_secret {
                     return Err(create_error!(NotAuthenticated))
                 }
             } else if let Some((code_verifier, method)) = info.code_verifier.as_ref().zip(claims.code_challange_method) {
-        let Some(server_code_challenge) = oauth2::get_code_challange(&info.code).await? else {
+                let Some(server_code_challenge) = oauth2::get_code_challange(token).await? else {
                     return Err(create_error!(NotAuthenticated))
                 };
 
@@ -51,40 +82,80 @@ pub async fn token(
                 return Err(create_error!(NotAuthenticated))
             }
 
-    if claims.client_id != info.client_id || claims.exp < Utc::now().timestamp() {
+            let token = oauth2::encode_token(
-        return Err(create_error!(NotAuthenticated));
+                &config.api.security.token_secret,
-    };
+                oauth2::TokenType::Access,
+                claims.sub.clone(),
+                claims.client_id.clone(),
+                claims.redirect_uri.clone(),
+                claims.scope.clone(),
+                None,
+            )
+            .map_err(|_| create_error!(InternalError))?;
 
-    match info.grant_type {
+            let refresh_token = oauth2::encode_token(
-        v0::OAuth2GrantType::AuthorizationCode => {
+                &config.api.security.token_secret,
-            if claims.r#type != oauth2::TokenType::Auth {
+                oauth2::TokenType::Refresh,
-                return Err(create_error!(InvalidOperation))
+                claims.sub.clone(),
+                claims.client_id.clone(),
+                claims.redirect_uri.clone(),
+                claims.scope.clone(),
+                None,
+            )
+            .map_err(|_| create_error!(InternalError))?;
+
+            let authorized_bot_id = AuthorizedBotId { bot: claims.client_id.clone(), user: claims.sub.clone() };
+
+            if db.fetch_authorized_bot(&authorized_bot_id).await.is_err_and(|err| err.error_type == ErrorType::NotFound) {
+                db.insert_authorized_bot(&AuthorizedBot {
+                    id: authorized_bot_id,
+                    created_at: Timestamp::now_utc(),
+                    deauthorized_at: None,
+                    scope: claims.scope.clone()
+                }).await?;
             }
 
+            Ok(Json(v0::OAuth2TokenExchangeResponse {
+                access_token: token,
+                refresh_token: Some(refresh_token),
+                token_type: "OAuth2".to_string(),
+                scope: claims.scope,
+            }))
+        },
+        v0::OAuth2GrantType::RefreshToken => {
+            if claims.r#type != oauth2::TokenType::Refresh {
+                return Err(create_error!(InvalidOperation))
+            };
+
             let token = oauth2::encode_token(
                 &config.api.security.token_secret,
                 oauth2::TokenType::Access,
                 claims.sub.clone(),
                 claims.client_id.clone(),
-                claims.redirect_uri,
+                claims.redirect_uri.clone(),
                 claims.scope.clone(),
                 None,
             )
             .map_err(|_| create_error!(InternalError))?;
 
-            db.insert_authorized_bot(&AuthorizedBot {
+            let refresh_token = oauth2::encode_token(
-                id: AuthorizedBotId { bot: claims.client_id.clone(), user: claims.sub.clone() },
+                &config.api.security.token_secret,
-                created_at: Timestamp::now_utc(),
+                oauth2::TokenType::Refresh,
-                deauthorized_at: None,
+                claims.sub.clone(),
-                scope: claims.scope.clone()
+                claims.client_id.clone(),
-            }).await?;
+                claims.redirect_uri.clone(),
+                claims.scope.clone(),
+                None,
+            )
+            .map_err(|_| create_error!(InternalError))?;
 
             Ok(Json(v0::OAuth2TokenExchangeResponse {
                 access_token: token,
+                refresh_token: Some(refresh_token),
                 token_type: "OAuth2".to_string(),
                 scope: claims.scope,
             }))
-        },
+        }
         v0::OAuth2GrantType::Implicit => {
             // token is already an access token so this endpoint does not need to be called - in theory this should be unreachable
             Err(create_error!(InvalidOperation))

crates/delta/src/routes/users/mod.rs

@@ -12,7 +12,6 @@ mod fetch_user;
 mod fetch_user_flags;
 mod find_mutual;
 mod get_default_avatar;
-mod oauth2_connections;
 mod open_dm;
 mod remove_friend;
 mod send_friend_request;