matrix/authentication-service
1
0
Fork 0
mirror of https://github.com/matrix-org/matrix-authentication-service.git synced 2025-07-31 09:24:31 +03:00
Code Activity

storage: remaining oauth2 repositories

Browse Source 
- authorization grants
 - access tokens
 - refresh tokens
This commit is contained in:
Quentin Gliech
2023-01-12 18:26:04 +01:00
parent 36396c0b45
commit 488a666a8d
17 changed files with 1700 additions and 1366 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
crates/axum-utils/src/user_authorization.rs
View File

@ -29,7 +29,7 @@ use headers::{authorization::Bearer, Authorization, Header, HeaderMapExt, Header
use http::{header::WWW_AUTHENTICATE, HeaderMap, HeaderValue, Request, StatusCode};
use mas_data_model::Session;
use mas_storage::{
oauth2::{access_token::find_access_token, OAuth2SessionRepository},
oauth2::{OAuth2AccessTokenRepository, OAuth2SessionRepository},
DatabaseError, Repository,
};
use serde::{de::DeserializeOwned, Deserialize};
@ -62,7 +62,9 @@ impl AccessToken {
AccessToken::None => return Err(AuthorizationVerificationError::MissingToken),
};
let token = find_access_token(conn, token.as_str())
let token = conn
.oauth2_access_token()
.find_by_token(token.as_str())
.await?
.ok_or(AuthorizationVerificationError::InvalidToken)?;

26
crates/data-model/src/oauth2/authorization_grant.rs
View File

@ -78,7 +78,7 @@ impl AuthorizationGrantStage {
Self::Pending
}
pub fn fulfill(
fn fulfill(
self,
fulfilled_at: DateTime<Utc>,
session: &Session,
@ -92,7 +92,7 @@ impl AuthorizationGrantStage {
}
}
pub fn exchange(self, exchanged_at: DateTime<Utc>) -> Result<Self, InvalidTransitionError> {
fn exchange(self, exchanged_at: DateTime<Utc>) -> Result<Self, InvalidTransitionError> {
match self {
Self::Fulfilled {
fulfilled_at,
@ -106,7 +106,7 @@ impl AuthorizationGrantStage {
}
}
pub fn cancel(self, cancelled_at: DateTime<Utc>) -> Result<Self, InvalidTransitionError> {
fn cancel(self, cancelled_at: DateTime<Utc>) -> Result<Self, InvalidTransitionError> {
match self {
Self::Pending => Ok(Self::Cancelled { cancelled_at }),
_ => Err(InvalidTransitionError),
@ -146,4 +146,24 @@ impl AuthorizationGrant {
let max_age: Option<i64> = self.max_age.map(|x| x.get().into());
self.created_at - Duration::seconds(max_age.unwrap_or(3600 * 24 * 365))
}
pub fn exchange(mut self, exchanged_at: DateTime<Utc>) -> Result<Self, InvalidTransitionError> {
self.stage = self.stage.exchange(exchanged_at)?;
Ok(self)
}
pub fn fulfill(
mut self,
fulfilled_at: DateTime<Utc>,
session: &Session,
) -> Result<Self, InvalidTransitionError> {
self.stage = self.stage.fulfill(fulfilled_at, session)?;
Ok(self)
}
// TODO: this is not used?
pub fn cancel(mut self, canceld_at: DateTime<Utc>) -> Result<Self, InvalidTransitionError> {
self.stage = self.stage.cancel(canceld_at)?;
Ok(self)
}
}

20
crates/handlers/src/oauth2/authorization/complete.rs
View File

@ -26,11 +26,7 @@ use mas_keystore::Encrypter;
use mas_policy::PolicyFactory;
use mas_router::{PostAuthAction, Route};
use mas_storage::{
oauth2::{
authorization_grant::{fulfill_grant, get_grant_by_id},
consent::fetch_client_consent,
OAuth2ClientRepository, OAuth2SessionRepository,
},
oauth2::{OAuth2AuthorizationGrantRepository, OAuth2ClientRepository, OAuth2SessionRepository},
Repository,
};
use mas_templates::Templates;
@ -94,7 +90,9 @@ pub(crate) async fn get(
let maybe_session = session_info.load_session(&mut txn).await?;
let grant = get_grant_by_id(&mut txn, grant_id)
let grant = txn
.oauth2_authorization_grant()
.lookup(grant_id)
.await?
.ok_or(RouteError::NotFound)?;
@ -192,7 +190,10 @@ pub(crate) async fn complete(
.await?
.ok_or(GrantCompletionError::NoSuchClient)?;
let current_consent = fetch_client_consent(&mut txn, &browser_session.user, &client).await?;
let current_consent = txn
.oauth2_client()
.get_consent_for_user(&client, &browser_session.user)
.await?;
let lacks_consent = grant
.scope
@ -211,7 +212,10 @@ pub(crate) async fn complete(
.create_from_grant(&mut rng, &clock, &grant, &browser_session)
.await?;
let grant = fulfill_grant(&mut txn, grant, session.clone()).await?;
let grant = txn
.oauth2_authorization_grant()
.fulfill(&clock, &session, grant)
.await?;
// Yep! Let's complete the auth now
let mut params = AuthorizationResponse::default();

10
crates/handlers/src/oauth2/authorization/mod.rs
View File

@ -26,7 +26,7 @@ use mas_keystore::Encrypter;
use mas_policy::PolicyFactory;
use mas_router::{PostAuthAction, Route};
use mas_storage::{
oauth2::{authorization_grant::new_authorization_grant, OAuth2ClientRepository},
oauth2::{OAuth2AuthorizationGrantRepository, OAuth2ClientRepository},
Repository,
};
use mas_templates::Templates;
@ -275,18 +275,18 @@ pub(crate) async fn get(
let requires_consent = prompt.contains(&Prompt::Consent);
let grant = new_authorization_grant(
&mut txn,
let grant = txn
.oauth2_authorization_grant()
.add(
&mut rng,
&clock,
client,
&client,
redirect_uri.clone(),
params.auth.scope,
code,
params.auth.state.clone(),
params.auth.nonce,
params.auth.max_age,
None,
response_mode,
response_type.has_id_token(),
requires_consent,

22
crates/handlers/src/oauth2/consent.rs
View File

@ -29,11 +29,7 @@ use mas_keystore::Encrypter;
use mas_policy::PolicyFactory;
use mas_router::{PostAuthAction, Route};
use mas_storage::{
oauth2::{
authorization_grant::{get_grant_by_id, give_consent_to_grant},
consent::insert_client_consent,
OAuth2ClientRepository,
},
oauth2::{OAuth2AuthorizationGrantRepository, OAuth2ClientRepository},
Repository,
};
use mas_templates::{ConsentContext, PolicyViolationContext, TemplateContext, Templates};
@ -91,7 +87,9 @@ pub(crate) async fn get(
let maybe_session = session_info.load_session(&mut conn).await?;
let grant = get_grant_by_id(&mut conn, grant_id)
let grant = conn
.oauth2_authorization_grant()
.lookup(grant_id)
.await?
.ok_or(RouteError::GrantNotFound)?;
@ -146,7 +144,9 @@ pub(crate) async fn post(
let maybe_session = session_info.load_session(&mut txn).await?;
let grant = get_grant_by_id(&mut txn, grant_id)
let grant = txn
.oauth2_authorization_grant()
.lookup(grant_id)
.await?
.ok_or(RouteError::GrantNotFound)?;
let next = PostAuthAction::continue_grant(grant_id);
@ -180,17 +180,17 @@ pub(crate) async fn post(
.filter(|s| !s.starts_with("urn:matrix:org.matrix.msc2967.client:device:"))
.cloned()
.collect();
insert_client_consent(
&mut txn,
txn.oauth2_client()
.give_consent_for_user(
&mut rng,
&clock,
&session.user,
&client,
&session.user,
&scope_without_device,
)
.await?;
let _grant = give_consent_to_grant(&mut txn, grant).await?;
txn.oauth2_authorization_grant().give_consent(grant).await?;
txn.commit().await?;

13
crates/handlers/src/oauth2/introspection.rs
View File

@ -23,10 +23,7 @@ use mas_iana::oauth::{OAuthClientAuthenticationMethod, OAuthTokenTypeHint};
use mas_keystore::Encrypter;
use mas_storage::{
compat::{CompatAccessTokenRepository, CompatRefreshTokenRepository, CompatSessionRepository},
oauth2::{
access_token::find_access_token, refresh_token::lookup_refresh_token,
OAuth2SessionRepository,
},
oauth2::{OAuth2AccessTokenRepository, OAuth2RefreshTokenRepository, OAuth2SessionRepository},
user::{BrowserSessionRepository, UserRepository},
Clock, Repository,
};
@ -169,7 +166,9 @@ pub(crate) async fn post(
let reply = match token_type {
TokenType::AccessToken => {
let token = find_access_token(&mut conn, token)
let token = conn
.oauth2_access_token()
.find_by_token(token)
.await?
.filter(|t| t.is_valid(clock.now()))
.ok_or(RouteError::UnknownToken)?;
@ -206,7 +205,9 @@ pub(crate) async fn post(
}
TokenType::RefreshToken => {
let token = lookup_refresh_token(&mut conn, token)
let token = conn
.oauth2_refresh_token()
.find_by_token(token)
.await?
.filter(|t| t.is_valid())
.ok_or(RouteError::UnknownToken)?;

66
crates/handlers/src/oauth2/token.rs
View File

@ -1,4 +1,4 @@
// Copyright 2021, 2022 The Matrix.org Foundation C.I.C.
// Copyright 2021-2023 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
@ -33,10 +33,8 @@ use mas_keystore::{Encrypter, Keystore};
use mas_router::UrlBuilder;
use mas_storage::{
oauth2::{
access_token::{add_access_token, lookup_access_token, revoke_access_token},
authorization_grant::{exchange_grant, lookup_grant_by_code},
refresh_token::{add_refresh_token, consume_refresh_token, lookup_refresh_token},
OAuth2SessionRepository,
OAuth2AccessTokenRepository, OAuth2AuthorizationGrantRepository,
OAuth2RefreshTokenRepository, OAuth2SessionRepository,
},
user::BrowserSessionRepository,
Repository,
@ -217,9 +215,9 @@ async fn authorization_code_grant(
) -> Result<AccessTokenResponse, RouteError> {
let (clock, mut rng) = crate::clock_and_rng();
// TODO: there is a bunch of unnecessary cloning here
// TODO: handle "not found" cases
let authz_grant = lookup_grant_by_code(&mut txn, &grant.code)
let authz_grant = txn
.oauth2_authorization_grant()
.find_by_code(&grant.code)
.await?
.ok_or(RouteError::GrantNotFound)?;
@ -301,17 +299,14 @@ async fn authorization_code_grant(
let access_token_str = TokenType::AccessToken.generate(&mut rng);
let refresh_token_str = TokenType::RefreshToken.generate(&mut rng);
let access_token =
add_access_token(&mut txn, &mut rng, &clock, &session, access_token_str, ttl).await?;
let access_token = txn
.oauth2_access_token()
.add(&mut rng, &clock, &session, access_token_str, ttl)
.await?;
let refresh_token = add_refresh_token(
&mut txn,
&mut rng,
&clock,
&session,
&access_token,
refresh_token_str,
)
let refresh_token = txn
.oauth2_refresh_token()
.add(&mut rng, &clock, &session, &access_token, refresh_token_str)
.await?;
let id_token = if session.scope.contains(&scope::OPENID) {
@ -360,7 +355,9 @@ async fn authorization_code_grant(
params = params.with_id_token(id_token);
}
exchange_grant(&mut txn, &clock, authz_grant).await?;
txn.oauth2_authorization_grant()
.exchange(&clock, authz_grant)
.await?;
txn.commit().await?;
@ -374,7 +371,9 @@ async fn refresh_token_grant(
) -> Result<AccessTokenResponse, RouteError> {
let (clock, mut rng) = crate::clock_and_rng();
let refresh_token = lookup_refresh_token(&mut txn, &grant.refresh_token)
let refresh_token = txn
.oauth2_refresh_token()
.find_by_token(&grant.refresh_token)
.await?
.ok_or(RouteError::InvalidGrant)?;
@ -397,18 +396,14 @@ async fn refresh_token_grant(
let access_token_str = TokenType::AccessToken.generate(&mut rng);
let refresh_token_str = TokenType::RefreshToken.generate(&mut rng);
let new_access_token = add_access_token(
&mut txn,
&mut rng,
&clock,
&session,
access_token_str.clone(),
ttl,
)
let new_access_token = txn
.oauth2_access_token()
.add(&mut rng, &clock, &session, access_token_str.clone(), ttl)
.await?;
let new_refresh_token = add_refresh_token(
&mut txn,
let new_refresh_token = txn
.oauth2_refresh_token()
.add(
&mut rng,
&clock,
&session,
@ -417,11 +412,16 @@ async fn refresh_token_grant(
)
.await?;
let refresh_token = consume_refresh_token(&mut txn, &clock, refresh_token).await?;
let refresh_token = txn
.oauth2_refresh_token()
.consume(&clock, refresh_token)
.await?;
if let Some(access_token_id) = refresh_token.access_token_id {
if let Some(access_token) = lookup_access_token(&mut txn, access_token_id).await? {
revoke_access_token(&mut txn, &clock, access_token).await?;
if let Some(access_token) = txn.oauth2_access_token().lookup(access_token_id).await? {
txn.oauth2_access_token()
.revoke(&clock, access_token)
.await?;
}
}

6
crates/handlers/src/views/shared.rs
View File

@ -15,7 +15,7 @@
use anyhow::Context;
use mas_router::{PostAuthAction, Route};
use mas_storage::{
compat::CompatSsoLoginRepository, oauth2::authorization_grant::get_grant_by_id,
compat::CompatSsoLoginRepository, oauth2::OAuth2AuthorizationGrantRepository,
upstream_oauth2::UpstreamOAuthProviderRepository, Repository, UpstreamOAuthLinkRepository,
};
use mas_templates::{PostAuthContext, PostAuthContextInner};
@ -46,7 +46,9 @@ impl OptionalPostAuthAction {
let Some(action) = self.post_auth_action.clone() else { return Ok(None) };
let ctx = match action {
PostAuthAction::ContinueAuthorizationGrant { id } => {
let grant = get_grant_by_id(conn, id)
let grant = conn
.oauth2_authorization_grant()
.lookup(id)
.await?
.context("Failed to load authorization grant")?;
let grant = Box::new(grant);

1075
crates/storage/sqlx-data.json
View File

File diff suppressed because it is too large Load Diff

250
crates/storage/src/oauth2/access_token.rs
View File

@ -1,4 +1,4 @@
// Copyright 2021 The Matrix.org Foundation C.I.C.
// Copyright 2021-2023 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
@ -12,67 +12,61 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use async_trait::async_trait;
use chrono::{DateTime, Duration, Utc};
use mas_data_model::{AccessToken, AccessTokenState, Session};
use rand::Rng;
use sqlx::{PgConnection, PgExecutor};
use rand::RngCore;
use sqlx::PgConnection;
use ulid::Ulid;
use uuid::Uuid;
use crate::{Clock, DatabaseError, LookupResultExt};
use crate::{tracing::ExecuteExt, Clock, DatabaseError, LookupResultExt};
#[tracing::instrument(
skip_all,
fields(
%session.id,
user_session.id = %session.user_session_id,
client.id = %session.client_id,
access_token.id,
),
err,
)]
pub async fn add_access_token(
executor: impl PgExecutor<'_>,
mut rng: impl Rng + Send,
#[async_trait]
pub trait OAuth2AccessTokenRepository: Send + Sync {
type Error;
/// Lookup an access token by its ID
async fn lookup(&mut self, id: Ulid) -> Result<Option<AccessToken>, Self::Error>;
/// Find an access token by its token
async fn find_by_token(
&mut self,
access_token: &str,
) -> Result<Option<AccessToken>, Self::Error>;
/// Add a new access token to the database
async fn add(
&mut self,
rng: &mut (dyn RngCore + Send),
clock: &Clock,
session: &Session,
access_token: String,
expires_after: Duration,
) -> Result<AccessToken, sqlx::Error> {
let created_at = clock.now();
let expires_at = created_at + expires_after;
let id = Ulid::from_datetime_with_source(created_at.into(), &mut rng);
) -> Result<AccessToken, Self::Error>;
tracing::Span::current().record("access_token.id", tracing::field::display(id));
/// Revoke an access token
async fn revoke(
&mut self,
clock: &Clock,
access_token: AccessToken,
) -> Result<AccessToken, Self::Error>;
sqlx::query!(
r#"
INSERT INTO oauth2_access_tokens
(oauth2_access_token_id, oauth2_session_id, access_token, created_at, expires_at)
VALUES
($1, $2, $3, $4, $5)
"#,
Uuid::from(id),
Uuid::from(session.id),
&access_token,
created_at,
expires_at,
)
.execute(executor)
.await?;
Ok(AccessToken {
id,
state: AccessTokenState::default(),
access_token,
session_id: session.id,
created_at,
expires_at,
})
/// Cleanup expired access tokens
async fn cleanup_expired(&mut self, clock: &Clock) -> Result<usize, Self::Error>;
}
#[derive(Debug)]
pub struct OAuth2AccessTokenLookup {
pub struct PgOAuth2AccessTokenRepository<'c> {
conn: &'c mut PgConnection,
}
impl<'c> PgOAuth2AccessTokenRepository<'c> {
pub fn new(conn: &'c mut PgConnection) -> Self {
Self { conn }
}
}
struct OAuth2AccessTokenLookup {
oauth2_access_token_id: Uuid,
oauth2_session_id: Uuid,
access_token: String,
@ -99,11 +93,48 @@ impl From<OAuth2AccessTokenLookup> for AccessToken {
}
}
#[tracing::instrument(skip_all, err)]
pub async fn find_access_token(
conn: &mut PgConnection,
token: &str,
) -> Result<Option<AccessToken>, DatabaseError> {
#[async_trait]
impl<'c> OAuth2AccessTokenRepository for PgOAuth2AccessTokenRepository<'c> {
type Error = DatabaseError;
async fn lookup(&mut self, id: Ulid) -> Result<Option<AccessToken>, Self::Error> {
let res = sqlx::query_as!(
OAuth2AccessTokenLookup,
r#"
SELECT oauth2_access_token_id
, access_token
, created_at
, expires_at
, revoked_at
, oauth2_session_id
FROM oauth2_access_tokens
WHERE oauth2_access_token_id = $1
"#,
Uuid::from(id),
)
.fetch_one(&mut *self.conn)
.await
.to_option()?;
let Some(res) = res else { return Ok(None) };
Ok(Some(res.into()))
}
#[tracing::instrument(
name = "db.oauth2_access_token.find_by_token",
skip_all,
fields(
db.statement,
),
err,
)]
async fn find_by_token(
&mut self,
access_token: &str,
) -> Result<Option<AccessToken>, Self::Error> {
let res = sqlx::query_as!(
OAuth2AccessTokenLookup,
r#"
@ -118,64 +149,75 @@ pub async fn find_access_token(
WHERE access_token = $1
"#,
token,
access_token,
)
.fetch_one(&mut *conn)
.fetch_one(&mut *self.conn)
.await
.to_option()?;
let Some(res) = res else { return Ok(None) };
Ok(Some(res.into()))
}
}
#[tracing::instrument(
skip_all,
fields(access_token.id = %access_token_id),
err,
)]
pub async fn lookup_access_token(
conn: &mut PgConnection,
access_token_id: Ulid,
) -> Result<Option<AccessToken>, DatabaseError> {
let res = sqlx::query_as!(
OAuth2AccessTokenLookup,
r#"
SELECT oauth2_access_token_id
, access_token
, created_at
, expires_at
, revoked_at
, oauth2_session_id
FROM oauth2_access_tokens
WHERE oauth2_access_token_id = $1
"#,
Uuid::from(access_token_id),
)
.fetch_one(&mut *conn)
.await
.to_option()?;
let Some(res) = res else { return Ok(None) };
Ok(Some(res.into()))
}
#[tracing::instrument(
#[tracing::instrument(
name = "db.oauth2_access_token.add",
skip_all,
fields(
%access_token.id,
session.id = %access_token.session_id,
db.statement,
%session.id,
user_session.id = %session.user_session_id,
client.id = %session.client_id,
access_token.id,
),
err,
)]
pub async fn revoke_access_token(
executor: impl PgExecutor<'_>,
)]
async fn add(
&mut self,
rng: &mut (dyn RngCore + Send),
clock: &Clock,
session: &Session,
access_token: String,
expires_after: Duration,
) -> Result<AccessToken, Self::Error> {
let created_at = clock.now();
let expires_at = created_at + expires_after;
let id = Ulid::from_datetime_with_source(created_at.into(), rng);
tracing::Span::current().record("access_token.id", tracing::field::display(id));
sqlx::query!(
r#"
INSERT INTO oauth2_access_tokens
(oauth2_access_token_id, oauth2_session_id, access_token, created_at, expires_at)
VALUES
($1, $2, $3, $4, $5)
"#,
Uuid::from(id),
Uuid::from(session.id),
&access_token,
created_at,
expires_at,
)
.traced()
.execute(&mut *self.conn)
.await?;
Ok(AccessToken {
id,
state: AccessTokenState::default(),
access_token,
session_id: session.id,
created_at,
expires_at,
})
}
async fn revoke(
&mut self,
clock: &Clock,
access_token: AccessToken,
) -> Result<AccessToken, DatabaseError> {
) -> Result<AccessToken, Self::Error> {
let revoked_at = clock.now();
let res = sqlx::query!(
r#"
@ -186,7 +228,7 @@ pub async fn revoke_access_token(
Uuid::from(access_token.id),
revoked_at,
)
.execute(executor)
.execute(&mut *self.conn)
.await?;
DatabaseError::ensure_affected_rows(&res, 1)?;
@ -194,12 +236,9 @@ pub async fn revoke_access_token(
access_token
.revoke(revoked_at)
.map_err(DatabaseError::to_invalid_operation)
}
}
pub async fn cleanup_expired(
executor: impl PgExecutor<'_>,
clock: &Clock,
) -> Result<u64, sqlx::Error> {
async fn cleanup_expired(&mut self, clock: &Clock) -> Result<usize, Self::Error> {
// Cleanup token which expired more than 15 minutes ago
let threshold = clock.now() - Duration::minutes(15);
let res = sqlx::query!(
@ -209,8 +248,9 @@ pub async fn cleanup_expired(
"#,
threshold,
)
.execute(executor)
.execute(&mut *self.conn)
.await?;
Ok(res.rows_affected())
Ok(res.rows_affected().try_into().unwrap_or(usize::MAX))
}
}

667
crates/storage/src/oauth2/authorization_grant.rs
View File

@ -14,45 +14,261 @@
use std::num::NonZeroU32;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use mas_data_model::{
AuthorizationCode, AuthorizationGrant, AuthorizationGrantStage, Client, Pkce, Session,
};
use mas_iana::oauth::PkceCodeChallengeMethod;
use oauth2_types::{requests::ResponseMode, scope::Scope};
use rand::Rng;
use sqlx::{PgConnection, PgExecutor};
use rand::RngCore;
use sqlx::PgConnection;
use ulid::Ulid;
use url::Url;
use uuid::Uuid;
use crate::{Clock, DatabaseError, DatabaseInconsistencyError, LookupResultExt};
use crate::{
tracing::ExecuteExt, Clock, DatabaseError, DatabaseInconsistencyError, LookupResultExt,
};
#[tracing::instrument(
skip_all,
fields(
%client.id,
grant.id,
),
err,
)]
#[allow(clippy::too_many_arguments)]
pub async fn new_authorization_grant(
executor: impl PgExecutor<'_>,
mut rng: impl Rng + Send,
#[async_trait]
pub trait OAuth2AuthorizationGrantRepository {
type Error;
#[allow(clippy::too_many_arguments)]
async fn add(
&mut self,
rng: &mut (dyn RngCore + Send),
clock: &Clock,
client: Client,
client: &Client,
redirect_uri: Url,
scope: Scope,
code: Option<AuthorizationCode>,
state: Option<String>,
nonce: Option<String>,
max_age: Option<NonZeroU32>,
_acr_values: Option<String>,
response_mode: ResponseMode,
response_type_id_token: bool,
requires_consent: bool,
) -> Result<AuthorizationGrant, sqlx::Error> {
) -> Result<AuthorizationGrant, Self::Error>;
async fn lookup(&mut self, id: Ulid) -> Result<Option<AuthorizationGrant>, Self::Error>;
async fn find_by_code(&mut self, code: &str)
-> Result<Option<AuthorizationGrant>, Self::Error>;
async fn fulfill(
&mut self,
clock: &Clock,
session: &Session,
authorization_grant: AuthorizationGrant,
) -> Result<AuthorizationGrant, Self::Error>;
async fn exchange(
&mut self,
clock: &Clock,
authorization_grant: AuthorizationGrant,
) -> Result<AuthorizationGrant, Self::Error>;
async fn give_consent(
&mut self,
authorization_grant: AuthorizationGrant,
) -> Result<AuthorizationGrant, Self::Error>;
}
pub struct PgOAuth2AuthorizationGrantRepository<'c> {
conn: &'c mut PgConnection,
}
impl<'c> PgOAuth2AuthorizationGrantRepository<'c> {
pub fn new(conn: &'c mut PgConnection) -> Self {
Self { conn }
}
}
#[allow(clippy::struct_excessive_bools)]
struct GrantLookup {
oauth2_authorization_grant_id: Uuid,
created_at: DateTime<Utc>,
cancelled_at: Option<DateTime<Utc>>,
fulfilled_at: Option<DateTime<Utc>>,
exchanged_at: Option<DateTime<Utc>>,
scope: String,
state: Option<String>,
nonce: Option<String>,
redirect_uri: String,
response_mode: String,
max_age: Option<i32>,
response_type_code: bool,
response_type_id_token: bool,
authorization_code: Option<String>,
code_challenge: Option<String>,
code_challenge_method: Option<String>,
requires_consent: bool,
oauth2_client_id: Uuid,
oauth2_session_id: Option<Uuid>,
}
impl TryFrom<GrantLookup> for AuthorizationGrant {
type Error = DatabaseInconsistencyError;
#[allow(clippy::too_many_lines)]
fn try_from(value: GrantLookup) -> Result<Self, Self::Error> {
let id = value.oauth2_authorization_grant_id.into();
let scope: Scope = value.scope.parse().map_err(|e| {
DatabaseInconsistencyError::on("oauth2_authorization_grants")
.column("scope")
.row(id)
.source(e)
})?;
let stage = match (
value.fulfilled_at,
value.exchanged_at,
value.cancelled_at,
value.oauth2_session_id,
) {
(None, None, None, None) => AuthorizationGrantStage::Pending,
(Some(fulfilled_at), None, None, Some(session_id)) => {
AuthorizationGrantStage::Fulfilled {
session_id: session_id.into(),
fulfilled_at,
}
}
(Some(fulfilled_at), Some(exchanged_at), None, Some(session_id)) => {
AuthorizationGrantStage::Exchanged {
session_id: session_id.into(),
fulfilled_at,
exchanged_at,
}
}
(None, None, Some(cancelled_at), None) => {
AuthorizationGrantStage::Cancelled { cancelled_at }
}
_ => {
return Err(
DatabaseInconsistencyError::on("oauth2_authorization_grants")
.column("stage")
.row(id),
);
}
};
let pkce = match (value.code_challenge, value.code_challenge_method) {
(Some(challenge), Some(challenge_method)) if challenge_method == "plain" => {
Some(Pkce {
challenge_method: PkceCodeChallengeMethod::Plain,
challenge,
})
}
(Some(challenge), Some(challenge_method)) if challenge_method == "S256" => Some(Pkce {
challenge_method: PkceCodeChallengeMethod::S256,
challenge,
}),
(None, None) => None,
_ => {
return Err(
DatabaseInconsistencyError::on("oauth2_authorization_grants")
.column("code_challenge_method")
.row(id),
);
}
};
let code: Option<AuthorizationCode> =
match (value.response_type_code, value.authorization_code, pkce) {
(false, None, None) => None,
(true, Some(code), pkce) => Some(AuthorizationCode { code, pkce }),
_ => {
return Err(
DatabaseInconsistencyError::on("oauth2_authorization_grants")
.column("authorization_code")
.row(id),
);
}
};
let redirect_uri = value.redirect_uri.parse().map_err(|e| {
DatabaseInconsistencyError::on("oauth2_authorization_grants")
.column("redirect_uri")
.row(id)
.source(e)
})?;
let response_mode = value.response_mode.parse().map_err(|e| {
DatabaseInconsistencyError::on("oauth2_authorization_grants")
.column("response_mode")
.row(id)
.source(e)
})?;
let max_age = value
.max_age
.map(u32::try_from)
.transpose()
.map_err(|e| {
DatabaseInconsistencyError::on("oauth2_authorization_grants")
.column("max_age")
.row(id)
.source(e)
})?
.map(NonZeroU32::try_from)
.transpose()
.map_err(|e| {
DatabaseInconsistencyError::on("oauth2_authorization_grants")
.column("max_age")
.row(id)
.source(e)
})?;
Ok(AuthorizationGrant {
id,
stage,
client_id: value.oauth2_client_id.into(),
code,
scope,
state: value.state,
nonce: value.nonce,
max_age,
response_mode,
redirect_uri,
created_at: value.created_at,
response_type_id_token: value.response_type_id_token,
requires_consent: value.requires_consent,
})
}
}
#[async_trait]
impl<'c> OAuth2AuthorizationGrantRepository for PgOAuth2AuthorizationGrantRepository<'c> {
type Error = DatabaseError;
#[tracing::instrument(
name = "db.oauth2_authorization_grant.add",
skip_all,
fields(
db.statement,
grant.id,
grant.scope = %scope,
%client.id,
),
err,
)]
async fn add(
&mut self,
rng: &mut (dyn RngCore + Send),
clock: &Clock,
client: &Client,
redirect_uri: Url,
scope: Scope,
code: Option<AuthorizationCode>,
state: Option<String>,
nonce: Option<String>,
max_age: Option<NonZeroU32>,
response_mode: ResponseMode,
response_type_id_token: bool,
requires_consent: bool,
) -> Result<AuthorizationGrant, Self::Error> {
let code_challenge = code
.as_ref()
.and_then(|c| c.pkce.as_ref())
@ -66,7 +282,7 @@ pub async fn new_authorization_grant(
let code_str = code.as_ref().map(|c| &c.code);
let created_at = clock.now();
let id = Ulid::from_datetime_with_source(created_at.into(), &mut rng);
let id = Ulid::from_datetime_with_source(created_at.into(), rng);
tracing::Span::current().record("grant.id", tracing::field::display(id));
sqlx::query!(
@ -107,7 +323,7 @@ pub async fn new_authorization_grant(
requires_consent,
created_at,
)
.execute(executor)
.execute(&mut *self.conn)
.await?;
Ok(AuthorizationGrant {
@ -125,207 +341,40 @@ pub async fn new_authorization_grant(
response_type_id_token,
requires_consent,
})
}
#[allow(clippy::struct_excessive_bools)]
struct GrantLookup {
oauth2_authorization_grant_id: Uuid,
oauth2_authorization_grant_created_at: DateTime<Utc>,
oauth2_authorization_grant_cancelled_at: Option<DateTime<Utc>>,
oauth2_authorization_grant_fulfilled_at: Option<DateTime<Utc>>,
oauth2_authorization_grant_exchanged_at: Option<DateTime<Utc>>,
oauth2_authorization_grant_scope: String,
oauth2_authorization_grant_state: Option<String>,
oauth2_authorization_grant_nonce: Option<String>,
oauth2_authorization_grant_redirect_uri: String,
oauth2_authorization_grant_response_mode: String,
oauth2_authorization_grant_max_age: Option<i32>,
oauth2_authorization_grant_response_type_code: bool,
oauth2_authorization_grant_response_type_id_token: bool,
oauth2_authorization_grant_code: Option<String>,
oauth2_authorization_grant_code_challenge: Option<String>,
oauth2_authorization_grant_code_challenge_method: Option<String>,
oauth2_authorization_grant_requires_consent: bool,
oauth2_client_id: Uuid,
oauth2_session_id: Option<Uuid>,
}
impl TryFrom<GrantLookup> for AuthorizationGrant {
type Error = DatabaseInconsistencyError;
#[allow(clippy::too_many_lines)]
fn try_from(value: GrantLookup) -> Result<Self, Self::Error> {
let id = value.oauth2_authorization_grant_id.into();
let scope: Scope = value
.oauth2_authorization_grant_scope
.parse()
.map_err(|e| {
DatabaseInconsistencyError::on("oauth2_authorization_grants")
.column("scope")
.row(id)
.source(e)
})?;
let stage = match (
value.oauth2_authorization_grant_fulfilled_at,
value.oauth2_authorization_grant_exchanged_at,
value.oauth2_authorization_grant_cancelled_at,
value.oauth2_session_id,
) {
(None, None, None, None) => AuthorizationGrantStage::Pending,
(Some(fulfilled_at), None, None, Some(session_id)) => {
AuthorizationGrantStage::Fulfilled {
session_id: session_id.into(),
fulfilled_at,
}
}
(Some(fulfilled_at), Some(exchanged_at), None, Some(session_id)) => {
AuthorizationGrantStage::Exchanged {
session_id: session_id.into(),
fulfilled_at,
exchanged_at,
}
}
(None, None, Some(cancelled_at), None) => {
AuthorizationGrantStage::Cancelled { cancelled_at }
}
_ => {
return Err(
DatabaseInconsistencyError::on("oauth2_authorization_grants")
.column("stage")
.row(id),
);
}
};
let pkce = match (
value.oauth2_authorization_grant_code_challenge,
value.oauth2_authorization_grant_code_challenge_method,
) {
(Some(challenge), Some(challenge_method)) if challenge_method == "plain" => {
Some(Pkce {
challenge_method: PkceCodeChallengeMethod::Plain,
challenge,
})
}
(Some(challenge), Some(challenge_method)) if challenge_method == "S256" => Some(Pkce {
challenge_method: PkceCodeChallengeMethod::S256,
challenge,
}),
(None, None) => None,
_ => {
return Err(
DatabaseInconsistencyError::on("oauth2_authorization_grants")
.column("code_challenge_method")
.row(id),
);
}
};
let code: Option<AuthorizationCode> = match (
value.oauth2_authorization_grant_response_type_code,
value.oauth2_authorization_grant_code,
pkce,
) {
(false, None, None) => None,
(true, Some(code), pkce) => Some(AuthorizationCode { code, pkce }),
_ => {
return Err(
DatabaseInconsistencyError::on("oauth2_authorization_grants")
.column("authorization_code")
.row(id),
);
}
};
let redirect_uri = value
.oauth2_authorization_grant_redirect_uri
.parse()
.map_err(|e| {
DatabaseInconsistencyError::on("oauth2_authorization_grants")
.column("redirect_uri")
.row(id)
.source(e)
})?;
let response_mode = value
.oauth2_authorization_grant_response_mode
.parse()
.map_err(|e| {
DatabaseInconsistencyError::on("oauth2_authorization_grants")
.column("response_mode")
.row(id)
.source(e)
})?;
let max_age = value
.oauth2_authorization_grant_max_age
.map(u32::try_from)
.transpose()
.map_err(|e| {
DatabaseInconsistencyError::on("oauth2_authorization_grants")
.column("max_age")
.row(id)
.source(e)
})?
.map(NonZeroU32::try_from)
.transpose()
.map_err(|e| {
DatabaseInconsistencyError::on("oauth2_authorization_grants")
.column("max_age")
.row(id)
.source(e)
})?;
Ok(AuthorizationGrant {
id,
stage,
client_id: value.oauth2_client_id.into(),
code,
scope,
state: value.oauth2_authorization_grant_state,
nonce: value.oauth2_authorization_grant_nonce,
max_age,
response_mode,
redirect_uri,
created_at: value.oauth2_authorization_grant_created_at,
response_type_id_token: value.oauth2_authorization_grant_response_type_id_token,
requires_consent: value.oauth2_authorization_grant_requires_consent,
})
}
}
#[tracing::instrument(
#[tracing::instrument(
name = "db.oauth2_authorization_grant.lookup",
skip_all,
fields(grant.id = %id),
fields(
db.statement,
grant.id = %id,
),
err,
)]
pub async fn get_grant_by_id(
conn: &mut PgConnection,
id: Ulid,
) -> Result<Option<AuthorizationGrant>, DatabaseError> {
)]
async fn lookup(&mut self, id: Ulid) -> Result<Option<AuthorizationGrant>, Self::Error> {
let res = sqlx::query_as!(
GrantLookup,
r#"
SELECT oauth2_authorization_grant_id
, created_at AS oauth2_authorization_grant_created_at
, cancelled_at AS oauth2_authorization_grant_cancelled_at
, fulfilled_at AS oauth2_authorization_grant_fulfilled_at
, exchanged_at AS oauth2_authorization_grant_exchanged_at
, scope AS oauth2_authorization_grant_scope
, state AS oauth2_authorization_grant_state
, redirect_uri AS oauth2_authorization_grant_redirect_uri
, response_mode AS oauth2_authorization_grant_response_mode
, nonce AS oauth2_authorization_grant_nonce
, max_age AS oauth2_authorization_grant_max_age
, oauth2_client_id AS oauth2_client_id
, authorization_code AS oauth2_authorization_grant_code
, response_type_code AS oauth2_authorization_grant_response_type_code
, response_type_id_token AS oauth2_authorization_grant_response_type_id_token
, code_challenge AS oauth2_authorization_grant_code_challenge
, code_challenge_method AS oauth2_authorization_grant_code_challenge_method
, requires_consent AS oauth2_authorization_grant_requires_consent
, oauth2_session_id AS "oauth2_session_id?"
, created_at
, cancelled_at
, fulfilled_at
, exchanged_at
, scope
, state
, redirect_uri
, response_mode
, nonce
, max_age
, oauth2_client_id
, authorization_code
, response_type_code
, response_type_id_token
, code_challenge
, code_challenge_method
, requires_consent
, oauth2_session_id
FROM
oauth2_authorization_grants
@ -333,42 +382,49 @@ pub async fn get_grant_by_id(
"#,
Uuid::from(id),
)
.fetch_one(&mut *conn)
.fetch_one(&mut *self.conn)
.await
.to_option()?;
let Some(res) = res else { return Ok(None) };
Ok(Some(res.try_into()?))
}
}
#[tracing::instrument(skip_all, err)]
pub async fn lookup_grant_by_code(
conn: &mut PgConnection,
#[tracing::instrument(
name = "db.oauth2_authorization_grant.find_by_code",
skip_all,
fields(
db.statement,
),
err,
)]
async fn find_by_code(
&mut self,
code: &str,
) -> Result<Option<AuthorizationGrant>, DatabaseError> {
) -> Result<Option<AuthorizationGrant>, Self::Error> {
let res = sqlx::query_as!(
GrantLookup,
r#"
SELECT oauth2_authorization_grant_id
, created_at AS oauth2_authorization_grant_created_at
, cancelled_at AS oauth2_authorization_grant_cancelled_at
, fulfilled_at AS oauth2_authorization_grant_fulfilled_at
, exchanged_at AS oauth2_authorization_grant_exchanged_at
, scope AS oauth2_authorization_grant_scope
, state AS oauth2_authorization_grant_state
, redirect_uri AS oauth2_authorization_grant_redirect_uri
, response_mode AS oauth2_authorization_grant_response_mode
, nonce AS oauth2_authorization_grant_nonce
, max_age AS oauth2_authorization_grant_max_age
, oauth2_client_id AS oauth2_client_id
, authorization_code AS oauth2_authorization_grant_code
, response_type_code AS oauth2_authorization_grant_response_type_code
, response_type_id_token AS oauth2_authorization_grant_response_type_id_token
, code_challenge AS oauth2_authorization_grant_code_challenge
, code_challenge_method AS oauth2_authorization_grant_code_challenge_method
, requires_consent AS oauth2_authorization_grant_requires_consent
, oauth2_session_id AS "oauth2_session_id?"
, created_at
, cancelled_at
, fulfilled_at
, exchanged_at
, scope
, state
, redirect_uri
, response_mode
, nonce
, max_age
, oauth2_client_id
, authorization_code
, response_type_code
, response_type_id_token
, code_challenge
, code_challenge_method
, requires_consent
, oauth2_session_id
FROM
oauth2_authorization_grants
@ -376,68 +432,110 @@ pub async fn lookup_grant_by_code(
"#,
code,
)
.fetch_one(&mut *conn)
.traced()
.fetch_one(&mut *self.conn)
.await
.to_option()?;
let Some(res) = res else { return Ok(None) };
Ok(Some(res.try_into()?))
}
}
#[tracing::instrument(
#[tracing::instrument(
name = "db.oauth2_authorization_grant.fulfill",
skip_all,
fields(
db.statement,
%grant.id,
client.id = %grant.client_id,
%session.id,
user_session.id = %session.user_session_id,
),
err,
)]
pub async fn fulfill_grant(
executor: impl PgExecutor<'_>,
mut grant: AuthorizationGrant,
session: Session,
) -> Result<AuthorizationGrant, DatabaseError> {
let fulfilled_at = sqlx::query_scalar!(
)]
async fn fulfill(
&mut self,
clock: &Clock,
session: &Session,
grant: AuthorizationGrant,
) -> Result<AuthorizationGrant, Self::Error> {
let fulfilled_at = clock.now();
let res = sqlx::query!(
r#"
UPDATE oauth2_authorization_grants AS og
SET
oauth2_session_id = os.oauth2_session_id,
fulfilled_at = os.created_at
FROM oauth2_sessions os
WHERE
og.oauth2_authorization_grant_id = $1
AND os.oauth2_session_id = $2
RETURNING fulfilled_at AS "fulfilled_at!: DateTime<Utc>"
UPDATE oauth2_authorization_grants
SET fulfilled_at = $2
, oauth2_session_id = $3
WHERE oauth2_authorization_grant_id = $1
"#,
Uuid::from(grant.id),
fulfilled_at,
Uuid::from(session.id),
)
.fetch_one(executor)
.execute(&mut *self.conn)
.await?;
grant.stage = grant
.stage
.fulfill(fulfilled_at, &session)
DatabaseError::ensure_affected_rows(&res, 1)?;
// XXX: check affected rows & new methods
let grant = grant
.fulfill(fulfilled_at, session)
.map_err(DatabaseError::to_invalid_operation)?;
Ok(grant)
}
}
#[tracing::instrument(
#[tracing::instrument(
name = "db.oauth2_authorization_grant.exchange",
skip_all,
fields(
db.statement,
%grant.id,
client.id = %grant.client_id,
),
err,
)]
pub async fn give_consent_to_grant(
executor: impl PgExecutor<'_>,
)]
async fn exchange(
&mut self,
clock: &Clock,
grant: AuthorizationGrant,
) -> Result<AuthorizationGrant, Self::Error> {
let exchanged_at = clock.now();
let res = sqlx::query!(
r#"
UPDATE oauth2_authorization_grants
SET exchanged_at = $2
WHERE oauth2_authorization_grant_id = $1
"#,
Uuid::from(grant.id),
exchanged_at,
)
.execute(&mut *self.conn)
.await?;
DatabaseError::ensure_affected_rows(&res, 1)?;
let grant = grant
.exchange(exchanged_at)
.map_err(DatabaseError::to_invalid_operation)?;
Ok(grant)
}
#[tracing::instrument(
name = "db.oauth2_authorization_grant.give_consent",
skip_all,
fields(
db.statement,
%grant.id,
client.id = %grant.client_id,
),
err,
)]
async fn give_consent(
&mut self,
mut grant: AuthorizationGrant,
) -> Result<AuthorizationGrant, sqlx::Error> {
) -> Result<AuthorizationGrant, Self::Error> {
sqlx::query!(
r#"
UPDATE oauth2_authorization_grants AS og
@ -448,44 +546,11 @@ pub async fn give_consent_to_grant(
"#,
Uuid::from(grant.id),
)
.execute(executor)
.execute(&mut *self.conn)
.await?;
grant.requires_consent = false;
Ok(grant)
}
#[tracing::instrument(
skip_all,
fields(
%grant.id,
client.id = %grant.client_id,
),
err,
)]
pub async fn exchange_grant(
executor: impl PgExecutor<'_>,
clock: &Clock,
mut grant: AuthorizationGrant,
) -> Result<AuthorizationGrant, DatabaseError> {
let exchanged_at = clock.now();
sqlx::query!(
r#"
UPDATE oauth2_authorization_grants
SET exchanged_at = $2
WHERE oauth2_authorization_grant_id = $1
"#,
Uuid::from(grant.id),
exchanged_at,
)
.execute(executor)
.await?;
grant.stage = grant
.stage
.exchange(exchanged_at)
.map_err(DatabaseError::to_invalid_operation)?;
Ok(grant)
}
}

113
crates/storage/src/oauth2/client.rs
View File

@ -14,17 +14,21 @@
use std::{
collections::{BTreeMap, BTreeSet},
str::FromStr,
string::ToString,
};
use async_trait::async_trait;
use mas_data_model::{Client, JwksOrJwksUri};
use mas_data_model::{Client, JwksOrJwksUri, User};
use mas_iana::{
jose::JsonWebSignatureAlg,
oauth::{OAuthAuthorizationEndpointResponseType, OAuthClientAuthenticationMethod},
};
use mas_jose::jwk::PublicJsonWebKeySet;
use oauth2_types::requests::GrantType;
use oauth2_types::{
requests::GrantType,
scope::{Scope, ScopeToken},
};
use rand::{Rng, RngCore};
use sqlx::PgConnection;
use tracing::{info_span, Instrument};
@ -87,6 +91,21 @@ pub trait OAuth2ClientRepository: Send + Sync {
jwks_uri: Option<Url>,
redirect_uris: Vec<Url>,
) -> Result<Client, Self::Error>;
async fn get_consent_for_user(
&mut self,
client: &Client,
user: &User,
) -> Result<Scope, Self::Error>;
async fn give_consent_for_user(
&mut self,
rng: &mut (dyn RngCore + Send),
clock: &Clock,
client: &Client,
user: &User,
scope: &Scope,
) -> Result<(), Self::Error>;
}
pub struct PgOAuth2ClientRepository<'c> {
@ -702,4 +721,94 @@ impl<'c> OAuth2ClientRepository for PgOAuth2ClientRepository<'c> {
initiate_login_uri: None,
})
}
#[tracing::instrument(
name = "db.oauth2_client.get_consent_for_user",
skip_all,
fields(
db.statement,
%user.id,
%client.id,
),
err,
)]
async fn get_consent_for_user(
&mut self,
client: &Client,
user: &User,
) -> Result<Scope, Self::Error> {
let scope_tokens: Vec<String> = sqlx::query_scalar!(
r#"
SELECT scope_token
FROM oauth2_consents
WHERE user_id = $1 AND oauth2_client_id = $2
"#,
Uuid::from(user.id),
Uuid::from(client.id),
)
.fetch_all(&mut *self.conn)
.await?;
let scope: Result<Scope, _> = scope_tokens
.into_iter()
.map(|s| ScopeToken::from_str(&s))
.collect();
let scope = scope.map_err(|e| {
DatabaseInconsistencyError::on("oauth2_consents")
.column("scope_token")
.source(e)
})?;
Ok(scope)
}
#[tracing::instrument(
skip_all,
fields(
db.statement,
%user.id,
%client.id,
%scope,
),
err,
)]
async fn give_consent_for_user(
&mut self,
rng: &mut (dyn RngCore + Send),
clock: &Clock,
client: &Client,
user: &User,
scope: &Scope,
) -> Result<(), Self::Error> {
let now = clock.now();
let (tokens, ids): (Vec<String>, Vec<Uuid>) = scope
.iter()
.map(|token| {
(
token.to_string(),
Uuid::from(Ulid::from_datetime_with_source(now.into(), rng)),
)
})
.unzip();
sqlx::query!(
r#"
INSERT INTO oauth2_consents
(oauth2_consent_id, user_id, oauth2_client_id, scope_token, created_at)
SELECT id, $2, $3, scope_token, $5 FROM UNNEST($1::uuid[], $4::text[]) u(id, scope_token)
ON CONFLICT (user_id, oauth2_client_id, scope_token) DO UPDATE SET refreshed_at = $5
"#,
&ids,
Uuid::from(user.id),
Uuid::from(client.id),
&tokens,
now,
)
.traced()
.execute(&mut *self.conn)
.await?;
Ok(())
}
}

110
crates/storage/src/oauth2/consent.rs
View File

@ -1,110 +0,0 @@
// Copyright 2022 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::str::FromStr;
use mas_data_model::{Client, User};
use oauth2_types::scope::{Scope, ScopeToken};
use rand::Rng;
use sqlx::PgExecutor;
use ulid::Ulid;
use uuid::Uuid;
use crate::{Clock, DatabaseError, DatabaseInconsistencyError};
#[tracing::instrument(
skip_all,
fields(
%user.id,
%client.id,
),
err,
)]
pub async fn fetch_client_consent(
executor: impl PgExecutor<'_>,
user: &User,
client: &Client,
) -> Result<Scope, DatabaseError> {
let scope_tokens: Vec<String> = sqlx::query_scalar!(
r#"
SELECT scope_token
FROM oauth2_consents
WHERE user_id = $1 AND oauth2_client_id = $2
"#,
Uuid::from(user.id),
Uuid::from(client.id),
)
.fetch_all(executor)
.await?;
let scope: Result<Scope, _> = scope_tokens
.into_iter()
.map(|s| ScopeToken::from_str(&s))
.collect();
let scope = scope.map_err(|e| {
DatabaseInconsistencyError::on("oauth2_consents")
.column("scope_token")
.source(e)
})?;
Ok(scope)
}
#[tracing::instrument(
skip_all,
fields(
%user.id,
%client.id,
%scope,
),
err,
)]
pub async fn insert_client_consent(
executor: impl PgExecutor<'_>,
mut rng: impl Rng + Send,
clock: &Clock,
user: &User,
client: &Client,
scope: &Scope,
) -> Result<(), sqlx::Error> {
let now = clock.now();
let (tokens, ids): (Vec<String>, Vec<Uuid>) = scope
.iter()
.map(|token| {
(
token.to_string(),
Uuid::from(Ulid::from_datetime_with_source(now.into(), &mut rng)),
)
})
.unzip();
sqlx::query!(
r#"
INSERT INTO oauth2_consents
(oauth2_consent_id, user_id, oauth2_client_id, scope_token, created_at)
SELECT id, $2, $3, scope_token, $5 FROM UNNEST($1::uuid[], $4::text[]) u(id, scope_token)
ON CONFLICT (user_id, oauth2_client_id, scope_token) DO UPDATE SET refreshed_at = $5
"#,
&ids,
Uuid::from(user.id),
Uuid::from(client.id),
&tokens,
now,
)
.execute(executor)
.await?;
Ok(())
}

10
crates/storage/src/oauth2/mod.rs
View File

@ -12,14 +12,18 @@
// See the License for the specific language governing permissions and
// limitations under the License.
pub mod access_token;
mod access_token;
pub mod authorization_grant;
mod client;
pub mod consent;
pub mod refresh_token;
mod refresh_token;
mod session;
pub use self::{
access_token::{OAuth2AccessTokenRepository, PgOAuth2AccessTokenRepository},
authorization_grant::{
OAuth2AuthorizationGrantRepository, PgOAuth2AuthorizationGrantRepository,
},
client::{OAuth2ClientRepository, PgOAuth2ClientRepository},
refresh_token::{OAuth2RefreshTokenRepository, PgOAuth2RefreshTokenRepository},
session::{OAuth2SessionRepository, PgOAuth2SessionRepository},
};

237
crates/storage/src/oauth2/refresh_token.rs
View File

@ -12,35 +12,181 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use mas_data_model::{AccessToken, RefreshToken, RefreshTokenState, Session};
use rand::Rng;
use sqlx::{PgConnection, PgExecutor};
use rand::RngCore;
use sqlx::PgConnection;
use ulid::Ulid;
use uuid::Uuid;
use crate::{Clock, DatabaseError};
use crate::{tracing::ExecuteExt, Clock, DatabaseError, LookupResultExt};
#[tracing::instrument(
#[async_trait]
pub trait OAuth2RefreshTokenRepository: Send + Sync {
type Error;
/// Lookup a refresh token by its ID
async fn lookup(&mut self, id: Ulid) -> Result<Option<RefreshToken>, Self::Error>;
/// Find a refresh token by its token
async fn find_by_token(
&mut self,
refresh_token: &str,
) -> Result<Option<RefreshToken>, Self::Error>;
/// Add a new refresh token to the database
async fn add(
&mut self,
rng: &mut (dyn RngCore + Send),
clock: &Clock,
session: &Session,
access_token: &AccessToken,
refresh_token: String,
) -> Result<RefreshToken, Self::Error>;
/// Consume a refresh token
async fn consume(
&mut self,
clock: &Clock,
refresh_token: RefreshToken,
) -> Result<RefreshToken, Self::Error>;
}
pub struct PgOAuth2RefreshTokenRepository<'c> {
conn: &'c mut PgConnection,
}
impl<'c> PgOAuth2RefreshTokenRepository<'c> {
pub fn new(conn: &'c mut PgConnection) -> Self {
Self { conn }
}
}
struct OAuth2RefreshTokenLookup {
oauth2_refresh_token_id: Uuid,
refresh_token: String,
created_at: DateTime<Utc>,
consumed_at: Option<DateTime<Utc>>,
oauth2_access_token_id: Option<Uuid>,
oauth2_session_id: Uuid,
}
impl From<OAuth2RefreshTokenLookup> for RefreshToken {
fn from(value: OAuth2RefreshTokenLookup) -> Self {
let state = match value.consumed_at {
None => RefreshTokenState::Valid,
Some(consumed_at) => RefreshTokenState::Consumed { consumed_at },
};
RefreshToken {
id: value.oauth2_refresh_token_id.into(),
state,
session_id: value.oauth2_session_id.into(),
refresh_token: value.refresh_token,
created_at: value.created_at,
access_token_id: value.oauth2_access_token_id.map(Ulid::from),
}
}
}
#[async_trait]
impl<'c> OAuth2RefreshTokenRepository for PgOAuth2RefreshTokenRepository<'c> {
type Error = DatabaseError;
#[tracing::instrument(
name = "db.oauth2_refresh_token.lookup",
skip_all,
fields(
db.statement,
refresh_token.id = %id,
),
err,
)]
async fn lookup(&mut self, id: Ulid) -> Result<Option<RefreshToken>, Self::Error> {
let res = sqlx::query_as!(
OAuth2RefreshTokenLookup,
r#"
SELECT oauth2_refresh_token_id
, refresh_token
, created_at
, consumed_at
, oauth2_access_token_id
, oauth2_session_id
FROM oauth2_refresh_tokens
WHERE oauth2_refresh_token_id = $1
"#,
Uuid::from(id),
)
.fetch_one(&mut *self.conn)
.await
.to_option()?;
let Some(res) = res else { return Ok(None) };
Ok(Some(res.into()))
}
#[tracing::instrument(
name = "db.oauth2_refresh_token.find_by_token",
skip_all,
fields(
db.statement,
),
err,
)]
async fn find_by_token(
&mut self,
refresh_token: &str,
) -> Result<Option<RefreshToken>, Self::Error> {
let res = sqlx::query_as!(
OAuth2RefreshTokenLookup,
r#"
SELECT oauth2_refresh_token_id
, refresh_token
, created_at
, consumed_at
, oauth2_access_token_id
, oauth2_session_id
FROM oauth2_refresh_tokens
WHERE refresh_token = $1
"#,
refresh_token,
)
.traced()
.fetch_one(&mut *self.conn)
.await
.to_option()?;
let Some(res) = res else { return Ok(None) };
Ok(Some(res.into()))
}
#[tracing::instrument(
name = "db.oauth2_refresh_token.add",
skip_all,
fields(
db.statement,
%session.id,
user_session.id = %session.user_session_id,
client.id = %session.client_id,
refresh_token.id,
),
err,
)]
pub async fn add_refresh_token(
executor: impl PgExecutor<'_>,
mut rng: impl Rng + Send,
)]
async fn add(
&mut self,
rng: &mut (dyn RngCore + Send),
clock: &Clock,
session: &Session,
access_token: &AccessToken,
refresh_token: String,
) -> Result<RefreshToken, sqlx::Error> {
) -> Result<RefreshToken, Self::Error> {
let created_at = clock.now();
let id = Ulid::from_datetime_with_source(created_at.into(), &mut rng);
let id = Ulid::from_datetime_with_source(created_at.into(), rng);
tracing::Span::current().record("refresh_token.id", tracing::field::display(id));
sqlx::query!(
@ -57,7 +203,8 @@ pub async fn add_refresh_token(
refresh_token,
created_at,
)
.execute(executor)
.traced()
.execute(&mut *self.conn)
.await?;
Ok(RefreshToken {
@ -68,70 +215,23 @@ pub async fn add_refresh_token(
access_token_id: Some(access_token.id),
created_at,
})
}
}
struct OAuth2RefreshTokenLookup {
oauth2_refresh_token_id: Uuid,
refresh_token: String,
created_at: DateTime<Utc>,
consumed_at: Option<DateTime<Utc>>,
oauth2_access_token_id: Option<Uuid>,
oauth2_session_id: Uuid,
}
#[tracing::instrument(skip_all, err)]
#[allow(clippy::too_many_lines)]
pub async fn lookup_refresh_token(
conn: &mut PgConnection,
token: &str,
) -> Result<Option<RefreshToken>, DatabaseError> {
let res = sqlx::query_as!(
OAuth2RefreshTokenLookup,
r#"
SELECT oauth2_refresh_token_id
, refresh_token
, created_at
, consumed_at
, oauth2_access_token_id
, oauth2_session_id
FROM oauth2_refresh_tokens
WHERE refresh_token = $1
"#,
token,
)
.fetch_one(&mut *conn)
.await?;
let state = match res.consumed_at {
None => RefreshTokenState::Valid,
Some(consumed_at) => RefreshTokenState::Consumed { consumed_at },
};
let refresh_token = RefreshToken {
id: res.oauth2_refresh_token_id.into(),
state,
session_id: res.oauth2_session_id.into(),
refresh_token: res.refresh_token,
created_at: res.created_at,
access_token_id: res.oauth2_access_token_id.map(Ulid::from),
};
Ok(Some(refresh_token))
}
#[tracing::instrument(
#[tracing::instrument(
name = "db.oauth2_refresh_token.consume",
skip_all,
fields(
db.statement,
%refresh_token.id,
session.id = %refresh_token.session_id,
),
err,
)]
pub async fn consume_refresh_token(
executor: impl PgExecutor<'_>,
)]
async fn consume(
&mut self,
clock: &Clock,
refresh_token: RefreshToken,
) -> Result<RefreshToken, DatabaseError> {
) -> Result<RefreshToken, Self::Error> {
let consumed_at = clock.now();
let res = sqlx::query!(
r#"
@ -142,7 +242,7 @@ pub async fn consume_refresh_token(
Uuid::from(refresh_token.id),
consumed_at,
)
.execute(executor)
.execute(&mut *self.conn)
.await?;
DatabaseError::ensure_affected_rows(&res, 1)?;
@ -150,4 +250,5 @@ pub async fn consume_refresh_token(
refresh_token
.consume(consumed_at)
.map_err(DatabaseError::to_invalid_operation)
}
}

50
crates/storage/src/repository.rs
View File

@ -19,7 +19,10 @@ use crate::{
PgCompatAccessTokenRepository, PgCompatRefreshTokenRepository, PgCompatSessionRepository,
PgCompatSsoLoginRepository,
},
oauth2::{PgOAuth2ClientRepository, PgOAuth2SessionRepository},
oauth2::{
PgOAuth2AccessTokenRepository, PgOAuth2AuthorizationGrantRepository,
PgOAuth2ClientRepository, PgOAuth2RefreshTokenRepository, PgOAuth2SessionRepository,
},
upstream_oauth2::{
PgUpstreamOAuthLinkRepository, PgUpstreamOAuthProviderRepository,
PgUpstreamOAuthSessionRepository,
@ -63,10 +66,22 @@ pub trait Repository {
where
Self: 'c;
type OAuth2AuthorizationGrantRepository<'c>
where
Self: 'c;
type OAuth2SessionRepository<'c>
where
Self: 'c;
type OAuth2AccessTokenRepository<'c>
where
Self: 'c;
type OAuth2RefreshTokenRepository<'c>
where
Self: 'c;
type CompatSessionRepository<'c>
where
Self: 'c;
@ -91,7 +106,10 @@ pub trait Repository {
fn user_password(&mut self) -> Self::UserPasswordRepository<'_>;
fn browser_session(&mut self) -> Self::BrowserSessionRepository<'_>;
fn oauth2_client(&mut self) -> Self::OAuth2ClientRepository<'_>;
fn oauth2_authorization_grant(&mut self) -> Self::OAuth2AuthorizationGrantRepository<'_>;
fn oauth2_session(&mut self) -> Self::OAuth2SessionRepository<'_>;
fn oauth2_access_token(&mut self) -> Self::OAuth2AccessTokenRepository<'_>;
fn oauth2_refresh_token(&mut self) -> Self::OAuth2RefreshTokenRepository<'_>;
fn compat_session(&mut self) -> Self::CompatSessionRepository<'_>;
fn compat_sso_login(&mut self) -> Self::CompatSsoLoginRepository<'_>;
fn compat_access_token(&mut self) -> Self::CompatAccessTokenRepository<'_>;
@ -107,7 +125,10 @@ impl Repository for PgConnection {
type UserPasswordRepository<'c> = PgUserPasswordRepository<'c> where Self: 'c;
type BrowserSessionRepository<'c> = PgBrowserSessionRepository<'c> where Self: 'c;
type OAuth2ClientRepository<'c> = PgOAuth2ClientRepository<'c> where Self: 'c;
type OAuth2AuthorizationGrantRepository<'c> = PgOAuth2AuthorizationGrantRepository<'c> where Self: 'c;
type OAuth2SessionRepository<'c> = PgOAuth2SessionRepository<'c> where Self: 'c;
type OAuth2AccessTokenRepository<'c> = PgOAuth2AccessTokenRepository<'c> where Self: 'c;
type OAuth2RefreshTokenRepository<'c> = PgOAuth2RefreshTokenRepository<'c> where Self: 'c;
type CompatSessionRepository<'c> = PgCompatSessionRepository<'c> where Self: 'c;
type CompatSsoLoginRepository<'c> = PgCompatSsoLoginRepository<'c> where Self: 'c;
type CompatAccessTokenRepository<'c> = PgCompatAccessTokenRepository<'c> where Self: 'c;
@ -145,10 +166,22 @@ impl Repository for PgConnection {
PgOAuth2ClientRepository::new(self)
}
fn oauth2_authorization_grant(&mut self) -> Self::OAuth2AuthorizationGrantRepository<'_> {
PgOAuth2AuthorizationGrantRepository::new(self)
}
fn oauth2_session(&mut self) -> Self::OAuth2SessionRepository<'_> {
PgOAuth2SessionRepository::new(self)
}
fn oauth2_access_token(&mut self) -> Self::OAuth2AccessTokenRepository<'_> {
PgOAuth2AccessTokenRepository::new(self)
}
fn oauth2_refresh_token(&mut self) -> Self::OAuth2RefreshTokenRepository<'_> {
PgOAuth2RefreshTokenRepository::new(self)
}
fn compat_session(&mut self) -> Self::CompatSessionRepository<'_> {
PgCompatSessionRepository::new(self)
}
@ -175,7 +208,10 @@ impl<'t> Repository for Transaction<'t, Postgres> {
type UserPasswordRepository<'c> = PgUserPasswordRepository<'c> where Self: 'c;
type BrowserSessionRepository<'c> = PgBrowserSessionRepository<'c> where Self: 'c;
type OAuth2ClientRepository<'c> = PgOAuth2ClientRepository<'c> where Self: 'c;
type OAuth2AuthorizationGrantRepository<'c> = PgOAuth2AuthorizationGrantRepository<'c> where Self: 'c;
type OAuth2SessionRepository<'c> = PgOAuth2SessionRepository<'c> where Self: 'c;
type OAuth2AccessTokenRepository<'c> = PgOAuth2AccessTokenRepository<'c> where Self: 'c;
type OAuth2RefreshTokenRepository<'c> = PgOAuth2RefreshTokenRepository<'c> where Self: 'c;
type CompatSessionRepository<'c> = PgCompatSessionRepository<'c> where Self: 'c;
type CompatSsoLoginRepository<'c> = PgCompatSsoLoginRepository<'c> where Self: 'c;
type CompatAccessTokenRepository<'c> = PgCompatAccessTokenRepository<'c> where Self: 'c;
@ -213,10 +249,22 @@ impl<'t> Repository for Transaction<'t, Postgres> {
PgOAuth2ClientRepository::new(self)
}
fn oauth2_authorization_grant(&mut self) -> Self::OAuth2AuthorizationGrantRepository<'_> {
PgOAuth2AuthorizationGrantRepository::new(self)
}
fn oauth2_session(&mut self) -> Self::OAuth2SessionRepository<'_> {
PgOAuth2SessionRepository::new(self)
}
fn oauth2_access_token(&mut self) -> Self::OAuth2AccessTokenRepository<'_> {
PgOAuth2AccessTokenRepository::new(self)
}
fn oauth2_refresh_token(&mut self) -> Self::OAuth2RefreshTokenRepository<'_> {
PgOAuth2RefreshTokenRepository::new(self)
}
fn compat_session(&mut self) -> Self::CompatSessionRepository<'_> {
PgCompatSessionRepository::new(self)
}

9
crates/tasks/src/database.rs
View File

@ -14,7 +14,7 @@
//! Database-related tasks
use mas_storage::Clock;
use mas_storage::{oauth2::OAuth2AccessTokenRepository, Clock, Repository};
use sqlx::{Pool, Postgres};
use tracing::{debug, error, info};
@ -32,7 +32,12 @@ impl std::fmt::Debug for CleanupExpired {
#[async_trait::async_trait]
impl Task for CleanupExpired {
async fn run(&self) {
let res = mas_storage::oauth2::access_token::cleanup_expired(&self.0, &self.1).await;
let res = async move {
let mut conn = self.0.acquire().await?;
conn.oauth2_access_token().cleanup_expired(&self.1).await
}
.await;
match res {
Ok(0) => {
debug!("no token to clean up");