Box the repository everywhere

crates/handlers/src/oauth2/authorization/complete.rs

@@ -27,9 +27,8 @@ use mas_policy::PolicyFactory;
 use mas_router::{PostAuthAction, Route};
 use mas_storage::{
     oauth2::{OAuth2AuthorizationGrantRepository, OAuth2ClientRepository, OAuth2SessionRepository},
-    BoxClock, BoxRng, Repository,
+    BoxClock, BoxRepository, BoxRng,
 };
-use mas_storage_pg::PgRepository;
 use mas_templates::Templates;
 use oauth2_types::requests::{AccessTokenResponse, AuthorizationResponse};
 use thiserror::Error;
@@ -69,7 +68,7 @@ impl IntoResponse for RouteError {
     }
 }
 
-impl_from_error_for_route!(mas_storage_pg::DatabaseError);
+impl_from_error_for_route!(mas_storage::RepositoryError);
 impl_from_error_for_route!(mas_policy::LoadError);
 impl_from_error_for_route!(mas_policy::InstanciateError);
 impl_from_error_for_route!(mas_policy::EvaluationError);
@@ -81,13 +80,13 @@ pub(crate) async fn get(
     clock: BoxClock,
     State(policy_factory): State<Arc<PolicyFactory>>,
     State(templates): State<Templates>,
-    mut repo: PgRepository,
+    mut repo: BoxRepository,
     cookie_jar: PrivateCookieJar<Encrypter>,
     Path(grant_id): Path<Ulid>,
 ) -> Result<Response, RouteError> {
     let (session_info, cookie_jar) = cookie_jar.session_info();
 
-    let maybe_session = session_info.load_session(&mut repo).await?;
+    let maybe_session = session_info.load_session(&mut *repo).await?;
 
     let grant = repo
         .oauth2_authorization_grant()
@@ -147,7 +146,7 @@ pub enum GrantCompletionError {
     NoSuchClient,
 }
 
-impl_from_error_for_route!(GrantCompletionError: mas_storage_pg::DatabaseError);
+impl_from_error_for_route!(GrantCompletionError: mas_storage::RepositoryError);
 impl_from_error_for_route!(GrantCompletionError: super::callback::IntoCallbackDestinationError);
 impl_from_error_for_route!(GrantCompletionError: mas_policy::LoadError);
 impl_from_error_for_route!(GrantCompletionError: mas_policy::InstanciateError);
@@ -159,7 +158,7 @@ pub(crate) async fn complete(
     grant: AuthorizationGrant,
     browser_session: BrowserSession,
     policy_factory: &PolicyFactory,
-    mut repo: PgRepository,
+    mut repo: BoxRepository,
 ) -> Result<AuthorizationResponse<Option<AccessTokenResponse>>, GrantCompletionError> {
     // Verify that the grant is in a pending stage
     if !grant.stage.is_pending() {

crates/handlers/src/oauth2/authorization/mod.rs

@@ -27,9 +27,8 @@ use mas_policy::PolicyFactory;
 use mas_router::{PostAuthAction, Route};
 use mas_storage::{
     oauth2::{OAuth2AuthorizationGrantRepository, OAuth2ClientRepository},
-    BoxClock, BoxRng, Repository,
+    BoxClock, BoxRepository, BoxRng,
 };
-use mas_storage_pg::PgRepository;
 use mas_templates::Templates;
 use oauth2_types::{
     errors::{ClientError, ClientErrorCode},
@@ -90,7 +89,7 @@ impl IntoResponse for RouteError {
     }
 }
 
-impl_from_error_for_route!(mas_storage_pg::DatabaseError);
+impl_from_error_for_route!(mas_storage::RepositoryError);
 impl_from_error_for_route!(self::callback::CallbackDestinationError);
 impl_from_error_for_route!(mas_policy::LoadError);
 impl_from_error_for_route!(mas_policy::InstanciateError);
@@ -135,7 +134,7 @@ pub(crate) async fn get(
     clock: BoxClock,
     State(policy_factory): State<Arc<PolicyFactory>>,
     State(templates): State<Templates>,
-    mut repo: PgRepository,
+    mut repo: BoxRepository,
     cookie_jar: PrivateCookieJar<Encrypter>,
     Form(params): Form<Params>,
 ) -> Result<Response, RouteError> {
@@ -168,7 +167,7 @@ pub(crate) async fn get(
         let templates = templates.clone();
         let callback_destination = callback_destination.clone();
         async move {
-            let maybe_session = session_info.load_session(&mut repo).await?;
+            let maybe_session = session_info.load_session(&mut *repo).await?;
             let prompt = params.auth.prompt.as_deref().unwrap_or_default();
 
             // Check if the request/request_uri/registration params are used. If so, reply

crates/handlers/src/oauth2/consent.rs

@@ -30,9 +30,8 @@ use mas_policy::PolicyFactory;
 use mas_router::{PostAuthAction, Route};
 use mas_storage::{
     oauth2::{OAuth2AuthorizationGrantRepository, OAuth2ClientRepository},
-    BoxClock, BoxRng, Repository,
+    BoxClock, BoxRepository, BoxRng,
 };
-use mas_storage_pg::PgRepository;
 use mas_templates::{ConsentContext, PolicyViolationContext, TemplateContext, Templates};
 use thiserror::Error;
 use ulid::Ulid;
@@ -61,7 +60,7 @@ pub enum RouteError {
 }
 
 impl_from_error_for_route!(mas_templates::TemplateError);
-impl_from_error_for_route!(mas_storage_pg::DatabaseError);
+impl_from_error_for_route!(mas_storage::RepositoryError);
 impl_from_error_for_route!(mas_policy::LoadError);
 impl_from_error_for_route!(mas_policy::InstanciateError);
 impl_from_error_for_route!(mas_policy::EvaluationError);
@@ -77,13 +76,13 @@ pub(crate) async fn get(
     clock: BoxClock,
     State(policy_factory): State<Arc<PolicyFactory>>,
     State(templates): State<Templates>,
-    mut repo: PgRepository,
+    mut repo: BoxRepository,
     cookie_jar: PrivateCookieJar<Encrypter>,
     Path(grant_id): Path<Ulid>,
 ) -> Result<Response, RouteError> {
     let (session_info, cookie_jar) = cookie_jar.session_info();
 
-    let maybe_session = session_info.load_session(&mut repo).await?;
+    let maybe_session = session_info.load_session(&mut *repo).await?;
 
     let grant = repo
         .oauth2_authorization_grant()
@@ -130,7 +129,7 @@ pub(crate) async fn post(
     mut rng: BoxRng,
     clock: BoxClock,
     State(policy_factory): State<Arc<PolicyFactory>>,
-    mut repo: PgRepository,
+    mut repo: BoxRepository,
     cookie_jar: PrivateCookieJar<Encrypter>,
     Path(grant_id): Path<Ulid>,
     Form(form): Form<ProtectedForm<()>>,
@@ -139,7 +138,7 @@ pub(crate) async fn post(
 
     let (session_info, cookie_jar) = cookie_jar.session_info();
 
-    let maybe_session = session_info.load_session(&mut repo).await?;
+    let maybe_session = session_info.load_session(&mut *repo).await?;
 
     let grant = repo
         .oauth2_authorization_grant()

crates/handlers/src/oauth2/introspection.rs

@@ -25,9 +25,8 @@ use mas_storage::{
     compat::{CompatAccessTokenRepository, CompatRefreshTokenRepository, CompatSessionRepository},
     oauth2::{OAuth2AccessTokenRepository, OAuth2RefreshTokenRepository, OAuth2SessionRepository},
     user::{BrowserSessionRepository, UserRepository},
-    BoxClock, Clock, Repository,
+    BoxClock, BoxRepository, Clock,
 };
-use mas_storage_pg::PgRepository;
 use oauth2_types::{
     errors::{ClientError, ClientErrorCode},
     requests::{IntrospectionRequest, IntrospectionResponse},
@@ -96,7 +95,7 @@ impl IntoResponse for RouteError {
     }
 }
 
-impl_from_error_for_route!(mas_storage_pg::DatabaseError);
+impl_from_error_for_route!(mas_storage::RepositoryError);
 
 impl From<TokenFormatError> for RouteError {
     fn from(_e: TokenFormatError) -> Self {
@@ -125,13 +124,13 @@ const API_SCOPE: ScopeToken = ScopeToken::from_static("urn:matrix:org.matrix.msc
 pub(crate) async fn post(
     clock: BoxClock,
     State(http_client_factory): State<HttpClientFactory>,
-    mut repo: PgRepository,
+    mut repo: BoxRepository,
     State(encrypter): State<Encrypter>,
     client_authorization: ClientAuthorization<IntrospectionRequest>,
 ) -> Result<impl IntoResponse, RouteError> {
     let client = client_authorization
         .credentials
-        .fetch(&mut repo)
+        .fetch(&mut *repo)
         .await
         .unwrap()
         .ok_or(RouteError::ClientNotFound)?;

crates/handlers/src/oauth2/registration.rs

@@ -19,8 +19,7 @@ use hyper::StatusCode;
 use mas_iana::oauth::OAuthClientAuthenticationMethod;
 use mas_keystore::Encrypter;
 use mas_policy::{PolicyFactory, Violation};
-use mas_storage::{oauth2::OAuth2ClientRepository, BoxClock, BoxRng, Repository};
-use mas_storage_pg::PgRepository;
+use mas_storage::{oauth2::OAuth2ClientRepository, BoxClock, BoxRepository, BoxRng};
 use oauth2_types::{
     errors::{ClientError, ClientErrorCode},
     registration::{
@@ -48,7 +47,7 @@ pub(crate) enum RouteError {
     PolicyDenied(Vec<Violation>),
 }
 
-impl_from_error_for_route!(mas_storage_pg::DatabaseError);
+impl_from_error_for_route!(mas_storage::RepositoryError);
 impl_from_error_for_route!(mas_policy::LoadError);
 impl_from_error_for_route!(mas_policy::InstanciateError);
 impl_from_error_for_route!(mas_policy::EvaluationError);
@@ -108,7 +107,7 @@ impl IntoResponse for RouteError {
 pub(crate) async fn post(
     mut rng: BoxRng,
     clock: BoxClock,
-    mut repo: PgRepository,
+    mut repo: BoxRepository,
     State(policy_factory): State<Arc<PolicyFactory>>,
     State(encrypter): State<Encrypter>,
     Json(body): Json<ClientMetadata>,

crates/handlers/src/oauth2/token.rs

@@ -37,9 +37,8 @@ use mas_storage::{
         OAuth2RefreshTokenRepository, OAuth2SessionRepository,
     },
     user::BrowserSessionRepository,
-    BoxClock, BoxRng, Clock, Repository,
+    BoxClock, BoxRepository, BoxRng, Clock,
 };
-use mas_storage_pg::PgRepository;
 use oauth2_types::{
     errors::{ClientError, ClientErrorCode},
     pkce::CodeChallengeError,
@@ -150,7 +149,7 @@ impl IntoResponse for RouteError {
     }
 }
 
-impl_from_error_for_route!(mas_storage_pg::DatabaseError);
+impl_from_error_for_route!(mas_storage::RepositoryError);
 impl_from_error_for_route!(mas_keystore::WrongAlgorithmError);
 impl_from_error_for_route!(mas_jose::claims::ClaimError);
 impl_from_error_for_route!(mas_jose::claims::TokenHashError);
@@ -163,13 +162,13 @@ pub(crate) async fn post(
     State(http_client_factory): State<HttpClientFactory>,
     State(key_store): State<Keystore>,
     State(url_builder): State<UrlBuilder>,
-    mut repo: PgRepository,
+    mut repo: BoxRepository,
     State(encrypter): State<Encrypter>,
     client_authorization: ClientAuthorization<AccessTokenRequest>,
 ) -> Result<impl IntoResponse, RouteError> {
     let client = client_authorization
         .credentials
-        .fetch(&mut repo)
+        .fetch(&mut *repo)
         .await?
         .ok_or(RouteError::ClientNotFound)?;
 
@@ -185,7 +184,7 @@ pub(crate) async fn post(
 
     let form = client_authorization.form.ok_or(RouteError::BadRequest)?;
 
-    let reply = match form {
+    let (reply, repo) = match form {
         AccessTokenRequest::AuthorizationCode(grant) => {
             authorization_code_grant(
                 &mut rng,
@@ -206,6 +205,8 @@ pub(crate) async fn post(
         }
     };
 
+    repo.save().await?;
+
     let mut headers = HeaderMap::new();
     headers.typed_insert(CacheControl::new().with_no_store());
     headers.typed_insert(Pragma::no_cache());
@@ -221,8 +222,8 @@ async fn authorization_code_grant(
     client: &Client,
     key_store: &Keystore,
     url_builder: &UrlBuilder,
-    mut repo: PgRepository,
-) -> Result<AccessTokenResponse, RouteError> {
+    mut repo: BoxRepository,
+) -> Result<(AccessTokenResponse, BoxRepository), RouteError> {
     let authz_grant = repo
         .oauth2_authorization_grant()
         .find_by_code(&grant.code)
@@ -367,9 +368,7 @@ async fn authorization_code_grant(
         .exchange(clock, authz_grant)
         .await?;
 
-    repo.save().await?;
-
-    Ok(params)
+    Ok((params, repo))
 }
 
 async fn refresh_token_grant(
@@ -377,8 +376,8 @@ async fn refresh_token_grant(
     clock: &impl Clock,
     grant: &RefreshTokenGrant,
     client: &Client,
-    mut repo: PgRepository,
-) -> Result<AccessTokenResponse, RouteError> {
+    mut repo: BoxRepository,
+) -> Result<(AccessTokenResponse, BoxRepository), RouteError> {
     let refresh_token = repo
         .oauth2_refresh_token()
         .find_by_token(&grant.refresh_token)
@@ -439,7 +438,5 @@ async fn refresh_token_grant(
         .with_refresh_token(new_refresh_token.refresh_token)
         .with_scope(session.scope);
 
-    repo.save().await?;
-
-    Ok(params)
+    Ok((params, repo))
 }

crates/handlers/src/oauth2/userinfo.rs

@@ -31,9 +31,8 @@ use mas_router::UrlBuilder;
 use mas_storage::{
     oauth2::OAuth2ClientRepository,
     user::{BrowserSessionRepository, UserEmailRepository},
-    BoxClock, BoxRng, Repository,
+    BoxClock, BoxRepository, BoxRng,
 };
-use mas_storage_pg::PgRepository;
 use oauth2_types::scope;
 use serde::Serialize;
 use serde_with::skip_serializing_none;
@@ -65,7 +64,7 @@ pub enum RouteError {
 
     #[error("failed to authenticate")]
     AuthorizationVerificationError(
-        #[from] AuthorizationVerificationError<mas_storage_pg::DatabaseError>,
+        #[from] AuthorizationVerificationError<mas_storage::RepositoryError>,
     ),
 
     #[error("no suitable key found for signing")]
@@ -78,7 +77,7 @@ pub enum RouteError {
     NoSuchBrowserSession,
 }
 
-impl_from_error_for_route!(mas_storage_pg::DatabaseError);
+impl_from_error_for_route!(mas_storage::RepositoryError);
 impl_from_error_for_route!(mas_keystore::WrongAlgorithmError);
 impl_from_error_for_route!(mas_jose::jwt::JwtSignatureError);
 
@@ -100,11 +99,11 @@ pub async fn get(
     mut rng: BoxRng,
     clock: BoxClock,
     State(url_builder): State<UrlBuilder>,
-    mut repo: PgRepository,
+    mut repo: BoxRepository,
     State(key_store): State<Keystore>,
     user_authorization: UserAuthorization,
 ) -> Result<Response, RouteError> {
-    let session = user_authorization.protected(&mut repo, &clock).await?;
+    let session = user_authorization.protected(&mut *repo, &clock).await?;
 
     let browser_session = repo
         .browser_session()