storage: unify most oauth2 related errors

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

@@ -26,12 +26,8 @@ use mas_data_model::{AuthorizationCode, Pkce};
 use mas_keystore::Encrypter;
 use mas_policy::PolicyFactory;
 use mas_router::{PostAuthAction, Route};
-use mas_storage::{
-    oauth2::{
-        authorization_grant::new_authorization_grant,
-        client::{lookup_client_by_client_id, ClientFetchError},
-    },
-    LookupError,
+use mas_storage::oauth2::{
+    authorization_grant::new_authorization_grant, client::lookup_client_by_client_id,
 };
 use mas_templates::Templates;
 use oauth2_types::{
@@ -46,6 +42,7 @@ use sqlx::PgPool;
 use thiserror::Error;
 
 use self::{callback::CallbackDestination, complete::GrantCompletionError};
+use crate::impl_from_error_for_route;
 
 mod callback;
 pub mod complete;
@@ -56,7 +53,7 @@ pub enum RouteError {
     Internal(Box<dyn std::error::Error + Send + Sync + 'static>),
 
     #[error(transparent)]
-    Anyhow(anyhow::Error),
+    Anyhow(#[from] anyhow::Error),
 
     #[error("could not find client")]
     ClientNotFound,
@@ -93,33 +90,9 @@ impl IntoResponse for RouteError {
     }
 }
 
-impl From<sqlx::Error> for RouteError {
-    fn from(e: sqlx::Error) -> Self {
-        Self::Internal(Box::new(e))
-    }
-}
-
-impl From<self::callback::CallbackDestinationError> for RouteError {
-    fn from(e: self::callback::CallbackDestinationError) -> Self {
-        Self::Internal(Box::new(e))
-    }
-}
-
-impl From<ClientFetchError> for RouteError {
-    fn from(e: ClientFetchError) -> Self {
-        if e.not_found() {
-            Self::ClientNotFound
-        } else {
-            Self::Internal(Box::new(e))
-        }
-    }
-}
-
-impl From<anyhow::Error> for RouteError {
-    fn from(e: anyhow::Error) -> Self {
-        Self::Anyhow(e)
-    }
-}
+impl_from_error_for_route!(sqlx::Error);
+impl_from_error_for_route!(mas_storage::DatabaseError);
+impl_from_error_for_route!(self::callback::CallbackDestinationError);
 
 #[derive(Deserialize)]
 pub(crate) struct Params {
@@ -166,7 +139,9 @@ pub(crate) async fn get(
     let mut txn = pool.begin().await?;
 
     // First, figure out what client it is
-    let client = lookup_client_by_client_id(&mut txn, &params.auth.client_id).await?;
+    let client = lookup_client_by_client_id(&mut txn, &params.auth.client_id)
+        .await?
+        .ok_or(RouteError::ClientNotFound)?;
 
     // And resolve the redirect_uri and response_mode
     let redirect_uri = client

crates/handlers/src/oauth2/introspection.rs

@@ -24,11 +24,9 @@ use mas_keystore::Encrypter;
 use mas_storage::{
     compat::{lookup_active_compat_access_token, lookup_active_compat_refresh_token},
     oauth2::{
-        access_token::{lookup_active_access_token, AccessTokenLookupError},
-        client::ClientFetchError,
-        refresh_token::{lookup_active_refresh_token, RefreshTokenLookupError},
+        access_token::lookup_active_access_token, refresh_token::lookup_active_refresh_token,
     },
-    Clock, LookupError,
+    Clock,
 };
 use oauth2_types::requests::{IntrospectionRequest, IntrospectionResponse};
 use sqlx::PgPool;
@@ -87,36 +85,6 @@ impl From<TokenFormatError> for RouteError {
     }
 }
 
-impl From<ClientFetchError> for RouteError {
-    fn from(e: ClientFetchError) -> Self {
-        if e.not_found() {
-            Self::ClientNotFound
-        } else {
-            Self::Internal(Box::new(e))
-        }
-    }
-}
-
-impl From<AccessTokenLookupError> for RouteError {
-    fn from(e: AccessTokenLookupError) -> Self {
-        if e.not_found() {
-            Self::UnknownToken
-        } else {
-            Self::Internal(Box::new(e))
-        }
-    }
-}
-
-impl From<RefreshTokenLookupError> for RouteError {
-    fn from(e: RefreshTokenLookupError) -> Self {
-        if e.not_found() {
-            Self::UnknownToken
-        } else {
-            Self::Internal(Box::new(e))
-        }
-    }
-}
-
 const INACTIVE: IntrospectionResponse = IntrospectionResponse {
     active: false,
     scope: None,
@@ -142,7 +110,11 @@ pub(crate) async fn post(
     let clock = Clock::default();
     let mut conn = pool.acquire().await?;
 
-    let client = client_authorization.credentials.fetch(&mut conn).await?;
+    let client = client_authorization
+        .credentials
+        .fetch(&mut conn)
+        .await?
+        .ok_or(RouteError::ClientNotFound)?;
 
     let method = match &client.token_endpoint_auth_method {
         None | Some(OAuthClientAuthenticationMethod::None) => {
@@ -172,7 +144,9 @@ pub(crate) async fn post(
 
     let reply = match token_type {
         TokenType::AccessToken => {
-            let (token, session) = lookup_active_access_token(&mut conn, token).await?;
+            let (token, session) = lookup_active_access_token(&mut conn, token)
+                .await?
+                .ok_or(RouteError::UnknownToken)?;
 
             IntrospectionResponse {
                 active: true,
@@ -190,7 +164,9 @@ pub(crate) async fn post(
             }
         }
         TokenType::RefreshToken => {
-            let (token, session) = lookup_active_refresh_token(&mut conn, token).await?;
+            let (token, session) = lookup_active_refresh_token(&mut conn, token)
+                .await?
+                .ok_or(RouteError::UnknownToken)?;
 
             IntrospectionResponse {
                 active: true,

crates/handlers/src/oauth2/token.rs

@@ -26,9 +26,9 @@ use mas_axum_utils::{
 use mas_data_model::{AuthorizationGrantStage, Client, TokenType};
 use mas_iana::jose::JsonWebSignatureAlg;
 use mas_jose::{
-    claims::{self, hash_token, ClaimError, TokenHashError},
+    claims::{self, hash_token},
     constraints::Constrainable,
-    jwt::{JsonWebSignatureHeader, Jwt, JwtSignatureError},
+    jwt::{JsonWebSignatureHeader, Jwt},
 };
 use mas_keystore::{Encrypter, Keystore};
 use mas_router::UrlBuilder;
@@ -36,14 +36,10 @@ use mas_storage::{
     oauth2::{
         access_token::{add_access_token, revoke_access_token},
         authorization_grant::{exchange_grant, lookup_grant_by_code},
-        client::ClientFetchError,
         end_oauth_session,
-        refresh_token::{
-            add_refresh_token, consume_refresh_token, lookup_active_refresh_token,
-            RefreshTokenLookupError,
-        },
+        refresh_token::{add_refresh_token, consume_refresh_token, lookup_active_refresh_token},
     },
-    DatabaseInconsistencyError, LookupError,
+    DatabaseInconsistencyError,
 };
 use oauth2_types::{
     errors::{ClientError, ClientErrorCode},
@@ -60,6 +56,8 @@ use thiserror::Error;
 use tracing::debug;
 use url::Url;
 
+use crate::impl_from_error_for_route;
+
 #[serde_as]
 #[skip_serializing_none]
 #[derive(Serialize, Debug)]
@@ -107,26 +105,6 @@ pub(crate) enum RouteError {
     UnauthorizedClient,
 }
 
-impl From<ClientFetchError> for RouteError {
-    fn from(e: ClientFetchError) -> Self {
-        if e.not_found() {
-            Self::ClientNotFound
-        } else {
-            Self::Internal(Box::new(e))
-        }
-    }
-}
-
-impl From<RefreshTokenLookupError> for RouteError {
-    fn from(e: RefreshTokenLookupError) -> Self {
-        if e.not_found() {
-            Self::InvalidGrant
-        } else {
-            Self::Internal(Box::new(e))
-        }
-    }
-}
-
 impl IntoResponse for RouteError {
     fn into_response(self) -> axum::response::Response {
         match self {
@@ -162,35 +140,12 @@ impl IntoResponse for RouteError {
     }
 }
 
-impl From<mas_keystore::WrongAlgorithmError> for RouteError {
-    fn from(e: mas_keystore::WrongAlgorithmError) -> Self {
-        Self::Internal(Box::new(e))
-    }
-}
-
-impl From<sqlx::Error> for RouteError {
-    fn from(e: sqlx::Error) -> Self {
-        Self::Internal(Box::new(e))
-    }
-}
-
-impl From<ClaimError> for RouteError {
-    fn from(e: ClaimError) -> Self {
-        Self::Internal(Box::new(e))
-    }
-}
-
-impl From<TokenHashError> for RouteError {
-    fn from(e: TokenHashError) -> Self {
-        Self::Internal(Box::new(e))
-    }
-}
-
-impl From<JwtSignatureError> for RouteError {
-    fn from(e: JwtSignatureError) -> Self {
-        Self::Internal(Box::new(e))
-    }
-}
+impl_from_error_for_route!(sqlx::Error);
+impl_from_error_for_route!(mas_storage::DatabaseError);
+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);
+impl_from_error_for_route!(mas_jose::jwt::JwtSignatureError);
 
 #[tracing::instrument(skip_all, err)]
 pub(crate) async fn post(
@@ -203,7 +158,11 @@ pub(crate) async fn post(
 ) -> Result<impl IntoResponse, RouteError> {
     let mut txn = pool.begin().await?;
 
-    let client = client_authorization.credentials.fetch(&mut txn).await?;
+    let client = client_authorization
+        .credentials
+        .fetch(&mut txn)
+        .await?
+        .ok_or(RouteError::ClientNotFound)?;
 
     let method = client
         .token_endpoint_auth_method
@@ -396,8 +355,9 @@ async fn refresh_token_grant(
 ) -> Result<AccessTokenResponse, RouteError> {
     let (clock, mut rng) = crate::rng_and_clock()?;
 
-    let (refresh_token, session) =
-        lookup_active_refresh_token(&mut txn, &grant.refresh_token).await?;
+    let (refresh_token, session) = lookup_active_refresh_token(&mut txn, &grant.refresh_token)
+        .await?
+        .ok_or(RouteError::InvalidGrant)?;
 
     if client.client_id != session.client.client_id {
         // As per https://datatracker.ietf.org/doc/html/rfc6749#section-5.2

crates/handlers/src/upstream_oauth2/callback.rs

@@ -25,9 +25,8 @@ use mas_oidc_client::requests::{
     authorization_code::AuthorizationValidationData, jose::JwtVerificationData,
 };
 use mas_router::{Route, UrlBuilder};
-use mas_storage::{
-    upstream_oauth2::{add_link, complete_session, lookup_link_by_subject, lookup_session},
-    LookupResultExt,
+use mas_storage::upstream_oauth2::{
+    add_link, complete_session, lookup_link_by_subject, lookup_session,
 };
 use oauth2_types::errors::ClientErrorCode;
 use serde::Deserialize;
@@ -97,8 +96,6 @@ pub(crate) enum RouteError {
 }
 
 impl_from_error_for_route!(mas_storage::DatabaseError);
-impl_from_error_for_route!(mas_storage::GenericLookupError);
-impl_from_error_for_route!(mas_storage::upstream_oauth2::SessionLookupError);
 impl_from_error_for_route!(mas_http::ClientInitError);
 impl_from_error_for_route!(sqlx::Error);
 impl_from_error_for_route!(mas_oidc_client::error::DiscoveryError);
@@ -141,8 +138,7 @@ pub(crate) async fn get(
         .map_err(|_| RouteError::MissingCookie)?;
 
     let (provider, session) = lookup_session(&mut txn, session_id)
-        .await
-        .to_option()?
+        .await?
         .ok_or(RouteError::SessionNotFound)?;
 
     if provider.id != provider_id {