jmug/stoatchat
1
0
Fork 1
Code Issues Packages Projects Releases Wiki Activity

feat: initial work on refresh tokens

Browse Source
This commit is contained in:
Zomatree
2025-07-05 00:58:57 +01:00
parent 660e646b2b
commit ef65223c89
15 changed files with 210 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
crates/core/database/src/drivers/reference.rs
View File

@@ -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>>>,

1
crates/core/database/src/models/authorized_bots/model.rs
View File

@@ -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,

10
crates/core/database/src/models/authorized_bots/ops.rs
View File

@@ -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>>;
}

32
crates/core/database/src/models/authorized_bots/ops/mongodb.rs
View File

@@ -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<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(
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<Vec<AuthorizedBot>> {
query!(
self,

64
crates/core/database/src/models/authorized_bots/ops/reference.rs
View File

@@ -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<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()
)
}
}

7
crates/core/models/src/v0/authorized_bots.rs
View File

@@ -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
}
);

10
crates/core/models/src/v0/oauth2.rs
View File

@@ -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,
}

2
crates/core/result/src/lib.rs
View File

@@ -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,

2
crates/daemons/crond/src/tasks/prune_authorized_bots.rs
View File

@@ -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);

3
crates/delta/src/routes/oauth2/authorize_auth.rs
View File

@@ -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
///

2
crates/delta/src/routes/oauth2/authorize_info.rs
View File

@@ -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
///

21
crates/delta/src/routes/oauth2/authorized_bots.rs
View File

@@ -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
.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))
}

131
crates/delta/src/routes/oauth2/token.rs
View File

@@ -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))

1
crates/delta/src/routes/users/mod.rs
View File

@@ -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;

0
crates/delta/src/routes/users/oauth2_connections.rs
View File