feat: allow connecting to ws via oauth2

crates/bonfire/src/websocket.rs

@@ -17,9 +17,11 @@ use redis_kiss::{PayloadType, REDIS_PAYLOAD_TYPE, REDIS_URI};
 use revolt_config::report_internal_error;
 use revolt_database::{
     events::{client::EventV1, server::ClientMessage},
+    util::oauth2,
     iso8601_timestamp::Timestamp,
     Database, User, UserHint,
 };
+use revolt_models::v0;
 use revolt_presence::{create_session, delete_session};
 
 use async_std::{
@@ -88,23 +90,51 @@ pub async fn client(db: &'static Database, stream: TcpStream, addr: SocketAddr)
         return;
     };
 
+    // Presume the token is a proper token first
     let (user, session_id) = match User::from_token(db, token, UserHint::Any).await {
-        Ok(user) => user,
+        Ok((user, session_id)) => {
-        Err(err) => {
+            db.update_session_last_seen(&session_id, Timestamp::now_utc())
-            write
-                .send(config.encode(&EventV1::Error { data: err }))
                 .await
                 .ok();
-            return;
+
+            (user, session_id)
+        },
+        Err(err) => {
+            let revolt_config = revolt_config::config().await;
+
+            // If it fails to find the user from the token see if its an OAuth2 token
+            let res = match oauth2::decode_token(&revolt_config.api.security.token_secret, token) {
+                // Check if the OAuth2 token is allowed to establish an events websocket
+                Ok(claims) => if !claims.scopes.contains(&v0::OAuth2Scope::Events) {
+                    // TODO: maybe a last_seen system for OAuth2 as well
+                    db.fetch_user(&claims.sub).await.map(|user| (user, claims.jti))
+                } else {
+                    Err(create_error!(MissingScope { scope: v0::OAuth2Scope::Events.to_string() }))
+                },
+                // If its expired return an error
+                Err(e) => if e.into_kind() == oauth2::JWTErrorKind::ExpiredSignature {
+                    Err(create_error!(ExpiredToken))
+                } else {
+                    // Finally re-return the error from User::from_token if everything else fails to avoid a confusing error
+                    Err(err)
+                }
+            };
+
+            match res {
+                Ok(user) => user,
+                Err(err) => {
+                    write
+                        .send(config.encode(&EventV1::Error { data: err }))
+                        .await
+                        .ok();
+                    return;
+                }
+            }
         }
     };
 
     info!("User {addr:?} authenticated as @{}", user.username);
 
-    db.update_session_last_seen(&session_id, Timestamp::now_utc())
-        .await
-        .ok();
-
     // Create local state.
     let mut state = State::from(user, session_id);
     let user_id = state.cache.user_id.clone();

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

@@ -1,5 +1,7 @@
 use iso8601_timestamp::Timestamp;
 
+use crate::OAuth2Scope;
+
 auto_derived! (
     /// Unique id of the user and bot
     #[derive(Hash)]
@@ -23,6 +25,6 @@ auto_derived! (
         pub deauthorized_at: Option<Timestamp>,
 
         /// Scopes the bot has access to
-        pub scope: String
+        pub scope: Vec<OAuth2Scope>,
     }
 );

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

@@ -49,9 +49,14 @@ auto_derived_partial!(
 
 auto_derived!(
     #[derive(Copy, Hash)]
-    #[serde(rename = "lowercase")]
     pub enum OAuth2Scope {
-        Identify,
+        #[serde(rename = "read:identify")]
+        ReadIdentify,
+        #[serde(rename = "read:servers")]
+        ReadServers,
+        #[serde(rename = "events")]
+        Events,
+        #[serde(rename = "full")]
         Full,
     }
 

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

@@ -38,7 +38,7 @@ impl<'r> FromRequest<'r> for User {
 
                     let claims = oauth2::decode_token(&config.api.security.token_secret, &header_oauth_token).ok()?;
 
-                    if oauth2::scopes_can_access_route(&claims.scope, request) {
+                    if oauth2::scopes_can_access_route(&claims.scopes, request) {
                         if let Ok(user) = db.fetch_user(&claims.sub).await {
                             return Some(user)
                         }

crates/core/database/src/util/bridge/v0.rs

@@ -92,7 +92,9 @@ impl From<crate::BotOauth2> for BotOauth2 {
 impl From<crate::OAuth2Scope> for OAuth2Scope {
     fn from(value: crate::OAuth2Scope) -> Self {
         match value {
-            crate::OAuth2Scope::Identify => OAuth2Scope::Identify,
+            crate::OAuth2Scope::ReadIdentify => OAuth2Scope::ReadIdentify,
+            crate::OAuth2Scope::ReadServers => OAuth2Scope::ReadServers,
+            crate::OAuth2Scope::Events => OAuth2Scope::Events,
             crate::OAuth2Scope::Full => OAuth2Scope::Full,
         }
     }
@@ -101,7 +103,9 @@ impl From<crate::OAuth2Scope> for OAuth2Scope {
 impl From<OAuth2Scope> for crate::OAuth2Scope {
     fn from(value: OAuth2Scope) -> Self {
         match value {
-            OAuth2Scope::Identify => crate::OAuth2Scope::Identify,
+            OAuth2Scope::ReadIdentify => crate::OAuth2Scope::ReadIdentify,
+            OAuth2Scope::ReadServers => crate::OAuth2Scope::ReadServers,
+            OAuth2Scope::Events => crate::OAuth2Scope::Events,
             OAuth2Scope::Full => crate::OAuth2Scope::Full,
         }
     }
@@ -1479,7 +1483,7 @@ impl From<AuthorizedBot> for crate::AuthorizedBot {
             id: value.id.into(),
             created_at: value.created_at,
             deauthorized_at: value.deauthorized_at,
-            scope: value.scope
+            scope: value.scope.into_iter().map(|scope| scope.into()).collect(),
         }
     }
 }
@@ -1490,7 +1494,7 @@ impl From<crate::AuthorizedBot> for AuthorizedBot {
             id: value.id.into(),
             created_at: value.created_at,
             deauthorized_at: value.deauthorized_at,
-            scope: value.scope
+            scope: value.scope.into_iter().map(|scope| scope.into()).collect(),
         }
     }
 }

crates/core/database/src/util/oauth2.rs

@@ -6,7 +6,10 @@ use redis_kiss::AsyncCommands;
 use revolt_result::Result;
 
 #[cfg(feature = "rocket")]
-use rocket::{http::Method, Request};
+use rocket::Request;
+
+pub use jsonwebtoken::errors::{Error as JWTError, ErrorKind as JWTErrorKind};
+use ulid::Ulid;
 
 #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
 pub enum TokenType {
@@ -27,12 +30,13 @@ impl TokenType {
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct Claims {
+    pub jti: String,
     pub sub: String,
     pub exp: i64,
 
     pub r#type: TokenType,
     pub client_id: String,
-    pub scope: String,
+    pub scopes: Vec<v0::OAuth2Scope>,
     pub redirect_uri: String,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub code_challange_method: Option<v0::OAuth2CodeChallangeMethod>,
@@ -45,21 +49,25 @@ pub fn encode_token(
     user_id: String,
     client_id: String,
     redirect_uri: String,
-    scope: String,
+    scopes: Vec<v0::OAuth2Scope>,
     code_challange_method: Option<v0::OAuth2CodeChallangeMethod>,
 ) -> Result<String, Error> {
-    let exp = Utc::now()
+    let now = Utc::now();
+
+    let exp = now
         .checked_add_signed(token_type.lifetime())
-        .unwrap()
+        .unwrap();
-        .timestamp();
+
+    println!("{now:?} {exp:}");
 
     let claims = Claims {
+        jti: Ulid::new().to_string(),
         sub: user_id,
-        exp,
+        exp: exp.timestamp(),
 
         r#type: token_type,
         client_id,
-        scope,
+        scopes,
         redirect_uri,
         code_challange_method
     };
@@ -78,7 +86,9 @@ pub fn decode_token(token_secret: &str, code: &str) -> Result<Claims, Error> {
 }
 
 #[cfg(feature = "rocket")]
-pub fn scope_can_access_route(scope: &str, request: &Request<'_>) -> bool {
+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();
 
@@ -86,27 +96,30 @@ pub fn scope_can_access_route(scope: &str, request: &Request<'_>) -> bool {
         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 {
-        "identify" => {
+        v0::OAuth2Scope::ReadIdentify => route_name == "fetch_self",
-            request.method() == Method::Get &&
+        v0::OAuth2Scope::ReadServers => route_name == "fetch_self_servers" || route_name == "fetch_server",
-            segments.get(0) == Some("users") &&
+        // This only grants access to the events websocket and not any routes
-            segments.get(1) == Some("@me")
+        v0::OAuth2Scope::Events => false,
-        },
+        // TODO: maybe disallow revoking other sessions
-        "full" => true,
+        v0::OAuth2Scope::Full => true,
-        _ => false
     }
 }
 
 #[cfg(feature = "rocket")]
-pub fn scopes_can_access_route(scopes: &str, request: &Request<'_>) -> bool {
+pub fn scopes_can_access_route(scopes: &[v0::OAuth2Scope], request: &Request<'_>) -> bool {
-    println!("{scopes}");
+    scopes
-    for scope in scopes.split(' ') {
+        .iter()
-        if scope_can_access_route(scope, request) {
+        .any(|&scope| scope_can_access_route(scope, request))
-            return true
-        }
-    }
-
-    false
 }
 
 pub async fn add_code_challange(token: &str, code_challenge: &str) -> Result<()> {

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

@@ -1,6 +1,6 @@
 use iso8601_timestamp::Timestamp;
 
-use crate::v0::PublicBot;
+use crate::v0::{PublicBot, OAuth2Scope};
 
 auto_derived!(
     /// Unique id of the user and bot
@@ -24,11 +24,11 @@ auto_derived!(
         pub deauthorized_at: Option<Timestamp>,
 
         /// Scopes the bot has access to
-        pub scope: String
+        pub scope: Vec<OAuth2Scope>,
     }
 
     pub struct AuthorizedBotsResponse {
         pub public_bot: PublicBot,
-        pub authorized_bot: AuthorizedBot
+        pub authorized_bot: AuthorizedBot,
     }
 );

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

@@ -1,5 +1,5 @@
 use crate::v0::{PublicBot, User};
-use std::collections::HashMap;
+use std::{collections::HashMap, fmt::Display};
 
 auto_derived!(
     pub struct OAuth2AuthorizeAuthResponse {
@@ -85,18 +85,37 @@ auto_derived!(
     }
 
     #[derive(Copy, Hash)]
-    #[cfg_attr(feature = "serde", serde(rename = "lowercase"))]
+    #[cfg_attr(feature = "serde", serde(rename = "snake_case"))]
     pub enum OAuth2Scope {
-        Identify,
+        #[cfg_attr(feature = "serde", serde(rename = "read:identify"))]
+        ReadIdentify,
+        #[cfg_attr(feature = "serde", serde(rename = "read:servers"))]
+        ReadServers,
+        #[cfg_attr(feature = "serde", serde(rename = "events"))]
+        Events,
+        #[cfg_attr(feature = "serde", serde(rename = "full"))]
         Full,
     }
 );
 
+impl Display for OAuth2Scope {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str(match self {
+            OAuth2Scope::ReadIdentify => "read:identify",
+            OAuth2Scope::ReadServers => "read:servers",
+            OAuth2Scope::Events => "events",
+            OAuth2Scope::Full => "full",
+        })
+    }
+}
+
 impl OAuth2Scope {
     #[allow(clippy::should_implement_trait)]
     pub fn from_str(str: &str) -> Option<OAuth2Scope> {
         match str {
-            "identify" => Some(OAuth2Scope::Identify),
+            "read:identify" => Some(OAuth2Scope::ReadIdentify),
+            "read:servers" => Some(OAuth2Scope::ReadServers),
+            "events" => Some(OAuth2Scope::Events),
             "full" => Some(OAuth2Scope::Full),
             _ => None,
         }

crates/core/result/src/axum.rs

@@ -78,6 +78,9 @@ impl IntoResponse for Error {
             ErrorType::InvalidFlagValue => StatusCode::BAD_REQUEST,
             ErrorType::FeatureDisabled { .. } => StatusCode::BAD_REQUEST,
 
+            ErrorType::MissingScope { .. } => StatusCode::UNAUTHORIZED,
+            ErrorType::ExpiredToken => StatusCode::UNAUTHORIZED,
+
             ErrorType::ProxyError => StatusCode::BAD_REQUEST,
             ErrorType::FileTooSmall => StatusCode::UNPROCESSABLE_ENTITY,
             ErrorType::FileTooLarge { .. } => StatusCode::UNPROCESSABLE_ENTITY,

crates/core/result/src/lib.rs

@@ -168,6 +168,12 @@ pub enum ErrorType {
     ImageProcessingFailed,
     NoEmbedData,
 
+    // ? OAuth2 Errors
+    ExpiredToken,
+    MissingScope {
+        scope: String
+    },
+
     // ? Legacy errors
     VosoUnavailable,
 

crates/core/result/src/rocket.rs

@@ -84,6 +84,9 @@ impl<'r> Responder<'r, 'static> for Error {
             ErrorType::FailedValidation { .. } => Status::BadRequest,
             ErrorType::FeatureDisabled { .. } => Status::BadRequest,
 
+            ErrorType::MissingScope { .. } => Status::Unauthorized,
+            ErrorType::ExpiredToken => Status::Unauthorized,
+
             ErrorType::ProxyError => Status::BadRequest,
             ErrorType::FileTooSmall => Status::UnprocessableEntity,
             ErrorType::FileTooLarge { .. } => Status::UnprocessableEntity,

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

@@ -33,7 +33,7 @@ pub async fn auth(
         return Err(create_error!(InvalidOperation));
     };
 
-    if scopes.into_iter().any(|scope| !oauth2.allowed_scopes.contains_key(&scope.into()))
+    if scopes.iter().any(|&scope| !oauth2.allowed_scopes.contains_key(&scope.into()))
         || !oauth2.redirects.contains(&info.redirect_uri)
         || v0::OAuth2Scope::scopes_from_str(&info.scope).is_none()
     {
@@ -55,7 +55,7 @@ pub async fn auth(
                 user.id.clone(),
                 info.client_id.clone(),
                 info.redirect_uri.clone(),
-                info.scope.clone(),
+                scopes.clone(),
                 info.code_challenge_method,
             )
             .map_err(|_| create_error!(InternalError))?;
@@ -64,7 +64,7 @@ pub async fn auth(
                 id: AuthorizedBotId { bot: info.client_id.clone(), user: user.id.clone() },
                 created_at: Timestamp::now_utc(),
                 deauthorized_at: None,
-                scope: info.scope.clone()
+                scope: scopes.iter().map(|&scope| scope.into()).collect()
             }).await?;
 
             token
@@ -87,11 +87,11 @@ pub async fn auth(
 
             let token = oauth2::encode_token(
                 &config.api.security.token_secret,
-                oauth2::TokenType::Access,
+                oauth2::TokenType::Auth,
                 user.id.clone(),
                 info.client_id.clone(),
                 info.redirect_uri.clone(),
-                info.scope.clone(),
+                scopes.clone(),
                 info.code_challenge_method,
             )
             .map_err(|_| create_error!(InternalError))?;

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

@@ -45,10 +45,10 @@ pub async fn token(
         return Err(create_error!(NotAuthenticated))
     };
 
-    let authorized_bot = db.fetch_authorized_bot(&AuthorizedBotId { user: claims.sub.clone(), bot: claims.client_id.clone() }).await?;
+    if let Ok(authorized_bot) = db.fetch_authorized_bot(&AuthorizedBotId { user: claims.sub.clone(), bot: claims.client_id.clone() }).await {
-
+        if authorized_bot.deauthorized_at.is_some() {
-    if authorized_bot.deauthorized_at.is_some() {
+            return Err(create_error!(NotAuthenticated));
-        return Err(create_error!(NotAuthenticated));
+        }
     };
 
     // TODO: track used tokens and dont allow them to be used twice
@@ -88,7 +88,7 @@ pub async fn token(
                 claims.sub.clone(),
                 claims.client_id.clone(),
                 claims.redirect_uri.clone(),
-                claims.scope.clone(),
+                claims.scopes.clone(),
                 None,
             )
             .map_err(|_| create_error!(InternalError))?;
@@ -99,19 +99,23 @@ pub async fn token(
                 claims.sub.clone(),
                 claims.client_id.clone(),
                 claims.redirect_uri.clone(),
-                claims.scope.clone(),
+                claims.scopes.clone(),
                 None,
             )
             .map_err(|_| create_error!(InternalError))?;
 
             let authorized_bot_id = AuthorizedBotId { bot: claims.client_id.clone(), user: claims.sub.clone() };
+            let auth_bot = db.fetch_authorized_bot(&authorized_bot_id).await;
+            println!("{auth_bot:?}");
+
+            if auth_bot.is_err_and(|err| err.error_type == ErrorType::NotFound) {
+                println!("inserting");
 
-            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()
+                    scope: claims.scopes.iter().map(|&scope| scope.into()).collect()
                 }).await?;
             }
 
@@ -119,7 +123,7 @@ pub async fn token(
                 access_token: token,
                 refresh_token: Some(refresh_token),
                 token_type: "OAuth2".to_string(),
-                scope: claims.scope,
+                scope: claims.scopes.iter().map(|scope| scope.to_string()).collect::<Vec<_>>().join(" "),
             }))
         },
         v0::OAuth2GrantType::RefreshToken => {
@@ -133,7 +137,7 @@ pub async fn token(
                 claims.sub.clone(),
                 claims.client_id.clone(),
                 claims.redirect_uri.clone(),
-                claims.scope.clone(),
+                claims.scopes.clone(),
                 None,
             )
             .map_err(|_| create_error!(InternalError))?;
@@ -144,7 +148,7 @@ pub async fn token(
                 claims.sub.clone(),
                 claims.client_id.clone(),
                 claims.redirect_uri.clone(),
-                claims.scope.clone(),
+                claims.scopes.clone(),
                 None,
             )
             .map_err(|_| create_error!(InternalError))?;
@@ -153,7 +157,7 @@ pub async fn token(
                 access_token: token,
                 refresh_token: Some(refresh_token),
                 token_type: "OAuth2".to_string(),
-                scope: claims.scope,
+                scope: claims.scopes.iter().map(|scope| scope.to_string()).collect::<Vec<_>>().join(" "),
             }))
         }
         v0::OAuth2GrantType::Implicit => {

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

@@ -29,7 +29,7 @@ pub fn routes() -> (Vec<Route>, OpenApi) {
     openapi_get_routes_spec![
         server_create::create_server,
         server_delete::delete,
-        server_fetch::fetch,
+        server_fetch::fetch_server,
         server_edit::edit,
         server_ack::ack,
         channel_create::create_server_channel,

crates/delta/src/routes/servers/server_fetch.rs

@@ -12,7 +12,7 @@ use rocket::{serde::json::Json, State};
 /// Fetch a server by its id.
 #[openapi(tag = "Server Information")]
 #[get("/<target>?<options..>")]
-pub async fn fetch(
+pub async fn fetch_server(
     db: &State<Database>,
     user: User,
     target: Reference<'_>,

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

@@ -8,6 +8,6 @@ use rocket::serde::json::Json;
 /// Retrieve your user information.
 #[openapi(tag = "User Information")]
 #[get("/@me")]
-pub async fn fetch(user: User) -> Result<Json<v0::User>> {
+pub async fn fetch_self(user: User) -> Result<Json<v0::User>> {
     Ok(Json(user.into_self(false).await))
 }

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

@@ -0,0 +1,53 @@
+use revolt_database::{util::permissions::DatabasePermissionQuery, Database, User};
+use revolt_models::v0;
+use revolt_permissions::{calculate_channel_permissions, ChannelPermission};
+use revolt_result::Result;
+use rocket::{serde::json::Json, State};
+
+/// # Fetch current user's servers
+///
+/// Retrieve all servers the current user is in.
+#[openapi(tag = "User Information")]
+#[get("/@me/servers?<options..>")]
+pub async fn fetch_self_servers(
+    db: &State<Database>,
+    user: User,
+    options: v0::OptionsFetchServer,
+) -> Result<Json<Vec<v0::FetchServerResponse>>> {
+    let members = db.fetch_all_memberships(&user.id).await?;
+
+    let server_ids = members
+        .iter()
+        .map(|x| x.id.server.clone())
+        .collect::<Vec<_>>();
+
+    let mut servers: Vec<v0::FetchServerResponse> = Vec::new();
+
+    for server in db.fetch_servers(&server_ids).await? {
+        if let Some(true) = options.include_channels {
+            let query = DatabasePermissionQuery::new(db, &user).server(&server);
+
+            let all_channels = db.fetch_channels(&server.channels).await?;
+            let mut visible_channels: Vec<v0::Channel> = vec![];
+
+            for channel in all_channels {
+                let mut channel_query = query.clone().channel(&channel);
+                if calculate_channel_permissions(&mut channel_query)
+                    .await
+                    .has_channel_permission(ChannelPermission::ViewChannel)
+                {
+                    visible_channels.push(channel.into());
+                }
+            }
+
+            servers.push(v0::FetchServerResponse::ServerWithChannels {
+                server: server.into(),
+                channels: visible_channels,
+            });
+        } else {
+            servers.push(v0::FetchServerResponse::JustServer(server.into()))
+        }
+    }
+
+    Ok(Json(servers))
+}

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

@@ -7,6 +7,7 @@ mod change_username;
 mod edit_user;
 mod fetch_dms;
 mod fetch_profile;
+mod fetch_self_servers;
 mod fetch_self;
 mod fetch_user;
 mod fetch_user_flags;
@@ -20,7 +21,8 @@ mod unblock_user;
 pub fn routes() -> (Vec<Route>, OpenApi) {
     openapi_get_routes_spec![
         // User Information
-        fetch_self::fetch,
+        fetch_self::fetch_self,
+        fetch_self_servers::fetch_self_servers,
         fetch_user::fetch,
         fetch_user_flags::fetch_user_flags,
         edit_user::edit,