From ef65223c89b2d2841876a5f31c80c4d6b81e3b44 Mon Sep 17 00:00:00 2001 From: Zomatree Date: Sat, 5 Jul 2025 00:58:57 +0100 Subject: [PATCH] feat: initial work on refresh tokens --- crates/core/database/src/drivers/reference.rs | 3 +- .../src/models/authorized_bots/model.rs | 1 + .../src/models/authorized_bots/ops.rs | 10 +- .../src/models/authorized_bots/ops/mongodb.rs | 32 +++-- .../models/authorized_bots/ops/reference.rs | 64 +++++++-- crates/core/models/src/v0/authorized_bots.rs | 7 + crates/core/models/src/v0/oauth2.rs | 10 +- crates/core/result/src/lib.rs | 2 +- .../crond/src/tasks/prune_authorized_bots.rs | 2 - .../delta/src/routes/oauth2/authorize_auth.rs | 3 +- .../delta/src/routes/oauth2/authorize_info.rs | 2 +- .../src/routes/oauth2/authorized_bots.rs | 21 ++- crates/delta/src/routes/oauth2/token.rs | 131 ++++++++++++++---- crates/delta/src/routes/users/mod.rs | 1 - .../src/routes/users/oauth2_connections.rs | 0 15 files changed, 210 insertions(+), 79 deletions(-) delete mode 100644 crates/delta/src/routes/users/oauth2_connections.rs diff --git a/crates/core/database/src/drivers/reference.rs b/crates/core/database/src/drivers/reference.rs index ff18f253..af5add99 100644 --- a/crates/core/database/src/drivers/reference.rs +++ b/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>>, pub bots: Arc>>, pub channels: Arc>>, pub channel_invites: Arc>>, diff --git a/crates/core/database/src/models/authorized_bots/model.rs b/crates/core/database/src/models/authorized_bots/model.rs index 6d341994..963c6c0a 100644 --- a/crates/core/database/src/models/authorized_bots/model.rs +++ b/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, diff --git a/crates/core/database/src/models/authorized_bots/ops.rs b/crates/core/database/src/models/authorized_bots/ops.rs index 3cea932f..d93f4f6f 100644 --- a/crates/core/database/src/models/authorized_bots/ops.rs +++ b/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; - /// 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>; /// 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; - /// Fetches all deauthorized bots + // Fetches all authorized bots which have been deauthorized async fn fetch_deauthorized_authorized_bots(&self) -> Result>; } diff --git a/crates/core/database/src/models/authorized_bots/ops/mongodb.rs b/crates/core/database/src/models/authorized_bots/ops/mongodb.rs index 876bd7d6..cb868c7e 100644 --- a/crates/core/database/src/models/authorized_bots/ops/mongodb.rs +++ b/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,44 +10,46 @@ static COL: &str = "authorized_bots"; #[async_trait] impl AbstractAuthorizedBots for MongoDb { - /// Insert emoji into database. - async fn insert_authorized_bot(&self, authorised_bot: &AuthorizedBot) -> Result<()> { - query!(self, insert_one, COL, &authorised_bot).map(|_| ()) + /// Insert an authorized bot into database. + async fn insert_authorized_bot(&self, authorized_bot: &AuthorizedBot) -> Result<()> { + 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 { 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> { 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 { self.col::(COL) .find_one_and_update( - doc! { - "id.user": &id.user, - "id.bot": &id.bot - }, - doc! { - "$set": { - "unauthorized_at": to_bson(&Timestamp::now_utc()).unwrap() + doc! { + "id.user": &id.user, + "id.bot": &id.bot + }, + doc! { + "$set": { + "deauthorized_at": to_bson(&Timestamp::now_utc()).unwrap() + } } - } ) .await .map_err(|_| create_database_error!("find_one_and_update", COL)) .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> { query!( self, diff --git a/crates/core/database/src/models/authorized_bots/ops/reference.rs b/crates/core/database/src/models/authorized_bots/ops/reference.rs index 46360032..a549177b 100644 --- a/crates/core/database/src/models/authorized_bots/ops/reference.rs +++ b/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. - async fn insert_authorized_bot(&self, authorised_bot: &AuthorizedBot) -> Result<()> { - todo!() + /// Insert an authorized bot into database. + async fn insert_authorized_bot(&self, authorized_bot: &AuthorizedBot) -> Result<()> { + 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 { - 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> { - 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 { - 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> { - todo!() + let authorized_bots = self.authorized_bots.lock().await; + + Ok(authorized_bots + .values() + .filter(|authorized_bot| authorized_bot.deauthorized_at.is_some()) + .cloned() + .collect() + ) } } \ No newline at end of file diff --git a/crates/core/models/src/v0/authorized_bots.rs b/crates/core/models/src/v0/authorized_bots.rs index 88b74f31..b22c7b0e 100644 --- a/crates/core/models/src/v0/authorized_bots.rs +++ b/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 + } ); \ No newline at end of file diff --git a/crates/core/models/src/v0/oauth2.rs b/crates/core/models/src/v0/oauth2.rs index 0048a558..05f789a7 100644 --- a/crates/core/models/src/v0/oauth2.rs +++ b/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, - pub code: String, + /// Authorization code + pub code: Option, + // Refresh token to generate new access token + pub refresh_token: Option, + pub code_verifier: Option, } pub struct OAuth2TokenExchangeResponse { pub access_token: String, + pub refresh_token: Option, pub token_type: String, pub scope: String, } diff --git a/crates/core/result/src/lib.rs b/crates/core/result/src/lib.rs index 10b56d5d..00273297 100644 --- a/crates/core/result/src/lib.rs +++ b/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, diff --git a/crates/daemons/crond/src/tasks/prune_authorized_bots.rs b/crates/daemons/crond/src/tasks/prune_authorized_bots.rs index e8774623..98a2dba6 100644 --- a/crates/daemons/crond/src/tasks/prune_authorized_bots.rs +++ b/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); diff --git a/crates/delta/src/routes/oauth2/authorize_auth.rs b/crates/delta/src/routes/oauth2/authorize_auth.rs index a83cc3c9..33ef4fdb 100644 --- a/crates/delta/src/routes/oauth2/authorize_auth.rs +++ b/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 redis_kiss::AsyncCommands; +use rocket::{serde::json::Json, State}; /// # OAuth2 Authorization Auth /// diff --git a/crates/delta/src/routes/oauth2/authorize_info.rs b/crates/delta/src/routes/oauth2/authorize_info.rs index ccd5a55a..41a92656 100644 --- a/crates/delta/src/routes/oauth2/authorize_info.rs +++ b/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 /// diff --git a/crates/delta/src/routes/oauth2/authorized_bots.rs b/crates/delta/src/routes/oauth2/authorized_bots.rs index de0198a5..f0e1241d 100644 --- a/crates/delta/src/routes/oauth2/authorized_bots.rs +++ b/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, user: User -) -> Result>> { +) -> Result>> { let authorized_bots = db.fetch_users_authorized_bots(&user.id).await?; - Ok(Json(authorized_bots - .into_iter() - .map(|bot| bot.into()) - .collect() - )) + let mut response = Vec::new(); + + for authorized_bot in authorized_bots { + 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)) } \ No newline at end of file diff --git a/crates/delta/src/routes/oauth2/token.rs b/crates/delta/src/routes/oauth2/token.rs index 1b735c24..a85b53c2 100644 --- a/crates/delta/src/routes/oauth2/token.rs +++ b/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,64 +26,136 @@ 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 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 { - return Err(create_error!(NotAuthenticated)) - }; - - let is_valid = match method { - v0::OAuth2CodeChallangeMethod::Plain => &server_code_challenge == code_verifier, - v0::OAuth2CodeChallangeMethod::S256 => pkce::code_challenge(code_verifier.as_bytes()) == server_code_challenge, - }; - - if !is_valid { - return Err(create_error!(NotAuthenticated)) - } - } else { + if claims.r#type != token_type || claims.client_id != info.client_id || claims.exp < Utc::now().timestamp() { return Err(create_error!(NotAuthenticated)) - } + }; - if claims.client_id != info.client_id || claims.exp < Utc::now().timestamp() { + 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(token).await? else { + return Err(create_error!(NotAuthenticated)) + }; + + let is_valid = match method { + v0::OAuth2CodeChallangeMethod::Plain => &server_code_challenge == code_verifier, + v0::OAuth2CodeChallangeMethod::S256 => pkce::code_challenge(code_verifier.as_bytes()) == server_code_challenge, + }; + + if !is_valid { + return Err(create_error!(NotAuthenticated)) + } + } else { + return Err(create_error!(NotAuthenticated)) + } + 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 { - id: AuthorizedBotId { bot: claims.client_id.clone(), user: claims.sub.clone() }, - created_at: Timestamp::now_utc(), - deauthorized_at: None, - scope: claims.scope.clone() - }).await?; + let refresh_token = oauth2::encode_token( + &config.api.security.token_secret, + oauth2::TokenType::Refresh, + 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.clone(), + claims.scope.clone(), + None, + ) + .map_err(|_| create_error!(InternalError))?; + + let refresh_token = oauth2::encode_token( + &config.api.security.token_secret, + oauth2::TokenType::Refresh, + claims.sub.clone(), + claims.client_id.clone(), + 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)) diff --git a/crates/delta/src/routes/users/mod.rs b/crates/delta/src/routes/users/mod.rs index 53eb8980..ba7445ac 100644 --- a/crates/delta/src/routes/users/mod.rs +++ b/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; diff --git a/crates/delta/src/routes/users/oauth2_connections.rs b/crates/delta/src/routes/users/oauth2_connections.rs deleted file mode 100644 index e69de29b..00000000