feat: allow connecting to ws via oauth2

This commit is contained in:
Zomatree
2025-07-13 23:51:09 +01:00
parent ef65223c89
commit f895ca7b23
18 changed files with 214 additions and 70 deletions

View File

@@ -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,
Err(err) => {
write
.send(config.encode(&EventV1::Error { data: err }))
Ok((user, session_id)) => {
db.update_session_last_seen(&session_id, Timestamp::now_utc())
.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();

View File

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

View File

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

View File

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

View File

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

View File

@@ -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()
.timestamp();
.unwrap();
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" => {
request.method() == Method::Get &&
segments.get(0) == Some("users") &&
segments.get(1) == Some("@me")
},
"full" => true,
_ => false
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: &str, request: &Request<'_>) -> bool {
println!("{scopes}");
for scope in scopes.split(' ') {
if scope_can_access_route(scope, request) {
return true
}
}
false
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<()> {

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@@ -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 authorized_bot.deauthorized_at.is_some() {
return Err(create_error!(NotAuthenticated));
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() {
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 => {

View File

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

View File

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

View File

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

View File

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

View File

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