From 9d4bcb5e3d4a079945075cce5fe34c633f0086c8 Mon Sep 17 00:00:00 2001 From: Zomatree Date: Mon, 21 Jul 2025 05:35:59 +0100 Subject: [PATCH] chore: rework oauth2 route scoping --- crates/core/database/src/models/bots/model.rs | 2 + crates/core/database/src/models/users/axum.rs | 22 +++++- .../core/database/src/models/users/rocket.rs | 17 +++-- crates/core/database/src/util/bridge/v0.rs | 2 + crates/core/database/src/util/oauth2/axum.rs | 17 +++++ .../src/util/{oauth2.rs => oauth2/mod.rs} | 67 ++++++------------- .../core/database/src/util/oauth2/rocket.rs | 53 +++++++++++++++ .../core/database/src/util/oauth2/scopes.rs | 27 ++++++++ crates/core/models/src/v0/oauth2.rs | 3 + crates/delta/src/routes/mod.rs | 13 +++- .../delta/src/routes/servers/server_fetch.rs | 3 +- crates/delta/src/routes/users/fetch_self.rs | 7 +- .../src/routes/users/fetch_self_servers.rs | 3 +- 13 files changed, 178 insertions(+), 58 deletions(-) create mode 100644 crates/core/database/src/util/oauth2/axum.rs rename crates/core/database/src/util/{oauth2.rs => oauth2/mod.rs} (64%) create mode 100644 crates/core/database/src/util/oauth2/rocket.rs create mode 100644 crates/core/database/src/util/oauth2/scopes.rs diff --git a/crates/core/database/src/models/bots/model.rs b/crates/core/database/src/models/bots/model.rs index d0822f08..5cb5a0af 100644 --- a/crates/core/database/src/models/bots/model.rs +++ b/crates/core/database/src/models/bots/model.rs @@ -54,6 +54,8 @@ auto_derived!( ReadIdentify, #[serde(rename = "read:servers")] ReadServers, + #[serde(rename = "write:files")] + WriteFiles, #[serde(rename = "events")] Events, #[serde(rename = "full")] diff --git a/crates/core/database/src/models/users/axum.rs b/crates/core/database/src/models/users/axum.rs index da83c5b1..60cbbfed 100644 --- a/crates/core/database/src/models/users/axum.rs +++ b/crates/core/database/src/models/users/axum.rs @@ -1,8 +1,10 @@ use axum::{extract::FromRequestParts, http::request::Parts}; +use revolt_config::config; +use revolt_models::v0; use revolt_result::{create_error, Error, Result}; -use crate::{Database, User}; +use crate::{util::oauth2, Database, OAuth2Scope, User}; #[async_trait::async_trait] impl FromRequestParts for User { @@ -17,6 +19,24 @@ impl FromRequestParts for User { { let session = db.fetch_session_by_token(session_token).await?; db.fetch_user(&session.user_id).await + } else if let Some(Ok(header_oauth_token)) = parts.headers.get("x-oauth2-token").map(|v| v.to_str()) { + let config = config().await; + + let claims = oauth2::decode_token( + &config.api.security.token_secret, + header_oauth_token, + ).map_err(|_| create_error!(NotAuthenticated))?; + + let required_scope: v0::OAuth2Scope = parts.extensions.get::() + .copied() + .ok_or_else(|| create_error!(NotAuthenticated))? + .into(); + + if claims.scopes.contains(&v0::OAuth2Scope::Full) || claims.scopes.contains(&required_scope) { + db.fetch_user(&claims.sub).await + } else { + Err(create_error!(MissingScope { scope: required_scope.to_string() })) + } } else { Err(create_error!(NotAuthenticated)) } diff --git a/crates/core/database/src/models/users/rocket.rs b/crates/core/database/src/models/users/rocket.rs index 81d27ee1..32f5b61f 100644 --- a/crates/core/database/src/models/users/rocket.rs +++ b/crates/core/database/src/models/users/rocket.rs @@ -1,5 +1,6 @@ use authifier::models::Session; use revolt_config::config; +use revolt_models::v0; use rocket::http::Status; use rocket::request::{self, FromRequest, Outcome, Request}; @@ -36,14 +37,22 @@ impl<'r> FromRequest<'r> for User { } else if let Some(header_oauth_token) = header_oauth_token { let config = config().await; - let claims = oauth2::decode_token(&config.api.security.token_secret, &header_oauth_token).ok()?; + let claims = oauth2::decode_token( + &config.api.security.token_secret, + &header_oauth_token, + ) + .ok()?; - if oauth2::scopes_can_access_route(&claims.scopes, request) { + let required_scope: v0::OAuth2Scope = request.local_cache(|| None::) + .as_ref() + .copied()? + .into(); + + if claims.scopes.contains(&v0::OAuth2Scope::Full) || claims.scopes.contains(&required_scope) { if let Ok(user) = db.fetch_user(&claims.sub).await { - return Some(user) + return Some(user); } } - } else if let Outcome::Success(session) = request.guard::().await { if let Ok(user) = db.fetch_user(&session.user_id).await { return Some(user); diff --git a/crates/core/database/src/util/bridge/v0.rs b/crates/core/database/src/util/bridge/v0.rs index 9e14a712..fc4290e8 100644 --- a/crates/core/database/src/util/bridge/v0.rs +++ b/crates/core/database/src/util/bridge/v0.rs @@ -94,6 +94,7 @@ impl From for OAuth2Scope { match value { crate::OAuth2Scope::ReadIdentify => OAuth2Scope::ReadIdentify, crate::OAuth2Scope::ReadServers => OAuth2Scope::ReadServers, + crate::OAuth2Scope::WriteFiles => OAuth2Scope::WriteFiles, crate::OAuth2Scope::Events => OAuth2Scope::Events, crate::OAuth2Scope::Full => OAuth2Scope::Full, } @@ -105,6 +106,7 @@ impl From for crate::OAuth2Scope { match value { OAuth2Scope::ReadIdentify => crate::OAuth2Scope::ReadIdentify, OAuth2Scope::ReadServers => crate::OAuth2Scope::ReadServers, + OAuth2Scope::WriteFiles => crate::OAuth2Scope::WriteFiles, OAuth2Scope::Events => crate::OAuth2Scope::Events, OAuth2Scope::Full => crate::OAuth2Scope::Full, } diff --git a/crates/core/database/src/util/oauth2/axum.rs b/crates/core/database/src/util/oauth2/axum.rs new file mode 100644 index 00000000..de46d5ae --- /dev/null +++ b/crates/core/database/src/util/oauth2/axum.rs @@ -0,0 +1,17 @@ +use std::marker::PhantomData; + +use axum::{extract::FromRequestParts, http::request::Parts}; +use revolt_result::{Error, Result}; + +use super::{OAuth2Scoped, scopes::OAuth2Scope}; + +#[rocket::async_trait] +impl FromRequestParts for OAuth2Scoped { + type Rejection = Error; + + async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result { + parts.extensions.insert(Scope::SCOPE); + + Ok(OAuth2Scoped { _scope: PhantomData }) + } +} \ No newline at end of file diff --git a/crates/core/database/src/util/oauth2.rs b/crates/core/database/src/util/oauth2/mod.rs similarity index 64% rename from crates/core/database/src/util/oauth2.rs rename to crates/core/database/src/util/oauth2/mod.rs index a7b28893..ce64fcd3 100644 --- a/crates/core/database/src/util/oauth2.rs +++ b/crates/core/database/src/util/oauth2/mod.rs @@ -1,12 +1,18 @@ use chrono::{TimeDelta, Utc}; use jsonwebtoken::{decode, encode, errors::Error, DecodingKey, EncodingKey, Header, Validation}; -use serde::{Serialize, Deserialize}; -use revolt_models::v0; use redis_kiss::AsyncCommands; +use revolt_models::v0; use revolt_result::Result; +use serde::{Deserialize, Serialize}; + +pub mod scopes; +pub use scopes::OAuth2Scoped; #[cfg(feature = "rocket")] -use rocket::Request; +pub mod rocket; + +#[cfg(feature = "axum")] +pub mod axum; pub use jsonwebtoken::errors::{Error as JWTError, ErrorKind as JWTErrorKind}; use ulid::Ulid; @@ -15,7 +21,7 @@ use ulid::Ulid; pub enum TokenType { Auth, Access, - Refresh + Refresh, } impl TokenType { @@ -54,9 +60,7 @@ pub fn encode_token( ) -> Result { let now = Utc::now(); - let exp = now - .checked_add_signed(token_type.lifetime()) - .unwrap(); + let exp = now.checked_add_signed(token_type.lifetime()).unwrap(); println!("{now:?} {exp:}"); @@ -69,7 +73,7 @@ pub fn encode_token( client_id, scopes, redirect_uri, - code_challange_method + code_challange_method, }; let encoding_key = EncodingKey::from_secret(token_secret.as_bytes()); @@ -80,48 +84,15 @@ pub fn encode_token( pub fn decode_token(token_secret: &str, code: &str) -> Result { let decoding_key = DecodingKey::from_secret(token_secret.as_bytes()); - let data = decode(code, &decoding_key, &Validation::new(jsonwebtoken::Algorithm::HS256))?; + let data = decode( + code, + &decoding_key, + &Validation::new(jsonwebtoken::Algorithm::HS256), + )?; Ok(data.claims) } -#[cfg(feature = "rocket")] -pub fn scope_can_access_route(scope: v0::OAuth2Scope, request: &Request<'_>) -> bool { - println!("{:?}", request.route()); - - // TODO: figure out why request.segments(0..) is skipping the first segment - let mut segments = request.uri().path().segments(); - - if segments.get(0) == Some("0.8") { - segments.next(); // skip first segment - }; - - // Extract the function name of the route - let Some(route_name) = request - .route() - .and_then(|route| route.name.as_ref()) - .map(|name| name.as_ref()) - else { - return false - }; - - match scope { - v0::OAuth2Scope::ReadIdentify => route_name == "fetch_self", - v0::OAuth2Scope::ReadServers => route_name == "fetch_self_servers" || route_name == "fetch_server", - // This only grants access to the events websocket and not any routes - v0::OAuth2Scope::Events => false, - // TODO: maybe disallow revoking other sessions - v0::OAuth2Scope::Full => true, - } -} - -#[cfg(feature = "rocket")] -pub fn scopes_can_access_route(scopes: &[v0::OAuth2Scope], request: &Request<'_>) -> bool { - scopes - .iter() - .any(|&scope| scope_can_access_route(scope, request)) -} - pub async fn add_code_challange(token: &str, code_challenge: &str) -> Result<()> { let mut conn = redis_kiss::get_connection() .await @@ -130,7 +101,7 @@ pub async fn add_code_challange(token: &str, code_challenge: &str) -> Result<()> conn.pset_ex::<_, _, ()>( format!("oauth2:{token}:code_challenge"), code_challenge, - TokenType::Access.lifetime().num_milliseconds() as usize + TokenType::Access.lifetime().num_milliseconds() as usize, ) .await .map_err(|_| create_error!(InternalError))?; @@ -146,4 +117,4 @@ pub async fn get_code_challange(token: &str) -> Result> { conn.get(format!("oauth2:{token}:code_challenge")) .await .map_err(|_| create_error!(InternalError)) -} \ No newline at end of file +} diff --git a/crates/core/database/src/util/oauth2/rocket.rs b/crates/core/database/src/util/oauth2/rocket.rs new file mode 100644 index 00000000..3f952389 --- /dev/null +++ b/crates/core/database/src/util/oauth2/rocket.rs @@ -0,0 +1,53 @@ +use std::{convert::Infallible, marker::PhantomData}; + +use revolt_okapi::openapi3::{OAuthFlows, SecurityScheme, SecuritySchemeData}; +use revolt_rocket_okapi::request::{OpenApiFromRequest, RequestHeaderInput}; +use rocket::{request::{self, FromRequest}, Request}; + +use super::{OAuth2Scoped, scopes::OAuth2Scope}; + + +#[rocket::async_trait] +impl<'r, Scope: OAuth2Scope> FromRequest<'r> for OAuth2Scoped { + type Error = Infallible; + + async fn from_request(request: &'r Request<'_>) -> request::Outcome { + request.local_cache(|| Some(Scope::SCOPE)); + + request::Outcome::Success(OAuth2Scoped { _scope: PhantomData }) + } +} + +impl<'r, Scope: OAuth2Scope> OpenApiFromRequest<'r> for OAuth2Scoped { + fn from_request_input( + _gen: &mut revolt_rocket_okapi::r#gen::OpenApiGenerator, + _name: String, + _required: bool, + ) -> revolt_rocket_okapi::Result { + Ok(RequestHeaderInput::Security( + "OAuth2".to_owned(), + SecurityScheme { + description: None, + extensions: Default::default(), + data: SecuritySchemeData::OAuth2 { + flows: OAuthFlows::AuthorizationCode { + authorization_url: "todo".to_string(), + token_url: "todo".to_string(), + refresh_url: Some("todo".to_string()), + scopes: revolt_okapi::map! { + "read:identify".to_string() => "".to_string(), + "read:servers".to_string() => "".to_string(), + "write:files".to_string() => "".to_string(), + "events".to_string() => "".to_string(), + "full".to_string() => "".to_string(), + }, + extensions: Default::default() + } + } + }, + revolt_okapi::map! { + "OAuth2".to_owned() => vec![Scope::NAME.to_owned()] + }, + )) + } +} diff --git a/crates/core/database/src/util/oauth2/scopes.rs b/crates/core/database/src/util/oauth2/scopes.rs new file mode 100644 index 00000000..1d93b247 --- /dev/null +++ b/crates/core/database/src/util/oauth2/scopes.rs @@ -0,0 +1,27 @@ +use std::marker::PhantomData; + +pub struct OAuth2Scoped { + pub(crate) _scope: PhantomData +} + +pub trait OAuth2Scope { + const NAME: &'static str; + const SCOPE: crate::OAuth2Scope; +} + +macro_rules! define_oauth2_scope { + ($struct_name:ident, $name:literal) => { + pub struct $struct_name; + + impl OAuth2Scope for $struct_name { + const NAME: &'static str = $name; + const SCOPE: crate::OAuth2Scope = crate::OAuth2Scope::$struct_name; + } + }; +} + +define_oauth2_scope!(ReadIdentify, "read:identify"); +define_oauth2_scope!(ReadServers, "read:servers"); +define_oauth2_scope!(WriteFiles, "write:files"); +define_oauth2_scope!(Events, "events"); +define_oauth2_scope!(Full, "full"); diff --git a/crates/core/models/src/v0/oauth2.rs b/crates/core/models/src/v0/oauth2.rs index 3d15c5b7..6903a57a 100644 --- a/crates/core/models/src/v0/oauth2.rs +++ b/crates/core/models/src/v0/oauth2.rs @@ -91,6 +91,8 @@ auto_derived!( ReadIdentify, #[cfg_attr(feature = "serde", serde(rename = "read:servers"))] ReadServers, + #[cfg_attr(feature = "serde", serde(rename = "write:files"))] + WriteFiles, #[cfg_attr(feature = "serde", serde(rename = "events"))] Events, #[cfg_attr(feature = "serde", serde(rename = "full"))] @@ -103,6 +105,7 @@ impl Display for OAuth2Scope { f.write_str(match self { OAuth2Scope::ReadIdentify => "read:identify", OAuth2Scope::ReadServers => "read:servers", + OAuth2Scope::WriteFiles => "write:files", OAuth2Scope::Events => "events", OAuth2Scope::Full => "full", }) diff --git a/crates/delta/src/routes/mod.rs b/crates/delta/src/routes/mod.rs index e7c22001..d7e48b36 100644 --- a/crates/delta/src/routes/mod.rs +++ b/crates/delta/src/routes/mod.rs @@ -203,7 +203,13 @@ fn custom_openapi_spec() -> OpenApi { "Sync", "Web Push" ] - } + }, + { + "name": "OAuth2", + "tags": [ + "OAuth2", + ] + }, ]), ); @@ -375,6 +381,11 @@ fn custom_openapi_spec() -> OpenApi { description: Some("Send messages from 3rd party services".to_owned()), ..Default::default() }, + Tag { + name: "OAuth2".to_owned(), + description: Some("Integrate Revolt into 3rd party applications".to_owned()), + ..Default::default() + }, ], ..Default::default() } diff --git a/crates/delta/src/routes/servers/server_fetch.rs b/crates/delta/src/routes/servers/server_fetch.rs index c56b8830..c7fb1f3a 100644 --- a/crates/delta/src/routes/servers/server_fetch.rs +++ b/crates/delta/src/routes/servers/server_fetch.rs @@ -1,5 +1,5 @@ use revolt_database::{ - util::{permissions::DatabasePermissionQuery, reference::Reference}, + util::{oauth2::{scopes, OAuth2Scoped}, permissions::DatabasePermissionQuery, reference::Reference}, Database, User, }; use revolt_models::v0; @@ -14,6 +14,7 @@ use rocket::{serde::json::Json, State}; #[get("/?")] pub async fn fetch_server( db: &State, + _oauth2_scope: OAuth2Scoped, user: User, target: Reference<'_>, options: v0::OptionsFetchServer, diff --git a/crates/delta/src/routes/users/fetch_self.rs b/crates/delta/src/routes/users/fetch_self.rs index b363f70b..b714b87e 100755 --- a/crates/delta/src/routes/users/fetch_self.rs +++ b/crates/delta/src/routes/users/fetch_self.rs @@ -1,4 +1,4 @@ -use revolt_database::User; +use revolt_database::{User, util::oauth2::{OAuth2Scoped, scopes}}; use revolt_models::v0; use revolt_result::Result; use rocket::serde::json::Json; @@ -8,6 +8,9 @@ use rocket::serde::json::Json; /// Retrieve your user information. #[openapi(tag = "User Information")] #[get("/@me")] -pub async fn fetch_self(user: User) -> Result> { +pub async fn fetch_self( + _oauth2_scope: OAuth2Scoped, + user: User +) -> Result> { Ok(Json(user.into_self(false).await)) } diff --git a/crates/delta/src/routes/users/fetch_self_servers.rs b/crates/delta/src/routes/users/fetch_self_servers.rs index 1f1a407b..70f09f6a 100644 --- a/crates/delta/src/routes/users/fetch_self_servers.rs +++ b/crates/delta/src/routes/users/fetch_self_servers.rs @@ -1,4 +1,4 @@ -use revolt_database::{util::permissions::DatabasePermissionQuery, Database, User}; +use revolt_database::{util::{oauth2::{scopes, OAuth2Scoped}, permissions::DatabasePermissionQuery}, Database, User}; use revolt_models::v0; use revolt_permissions::{calculate_channel_permissions, ChannelPermission}; use revolt_result::Result; @@ -11,6 +11,7 @@ use rocket::{serde::json::Json, State}; #[get("/@me/servers?")] pub async fn fetch_self_servers( db: &State, + _oauth2_scope: OAuth2Scoped, user: User, options: v0::OptionsFetchServer, ) -> Result>> {