From cad6d54ddb4c849b5673a59130e2a8b28f59931b Mon Sep 17 00:00:00 2001
From: Quentin Gliech <quenting@element.io>
Date: Fri, 25 Feb 2022 11:28:23 +0100
Subject: [PATCH] Reply with proper errors on the OAuth token endpoint

---
 crates/handlers/src/oauth2/token.rs        | 37 +++++++++++++++-------
 crates/storage/src/oauth2/refresh_token.rs | 28 +++++++++++++---
 2 files changed, 50 insertions(+), 15 deletions(-)

diff --git a/crates/handlers/src/oauth2/token.rs b/crates/handlers/src/oauth2/token.rs
index b97f5aff..391dccb6 100644
--- a/crates/handlers/src/oauth2/token.rs
+++ b/crates/handlers/src/oauth2/token.rs
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::{collections::HashMap, sync::Arc};
+use std::{collections::HashMap, convert::Infallible, sync::Arc};
 
 use anyhow::Context;
 use chrono::{DateTime, Duration, Utc};
@@ -31,7 +31,10 @@ use mas_storage::{
         access_token::{add_access_token, revoke_access_token},
         authorization_grant::{exchange_grant, lookup_grant_by_code},
         end_oauth_session,
-        refresh_token::{add_refresh_token, lookup_active_refresh_token, replace_refresh_token},
+        refresh_token::{
+            add_refresh_token, lookup_active_refresh_token, replace_refresh_token,
+            RefreshTokenLookupError,
+        },
     },
     DatabaseInconsistencyError,
 };
@@ -41,7 +44,9 @@ use mas_warp_utils::{
     reply::with_typed_header,
 };
 use oauth2_types::{
-    errors::{InvalidGrant, InvalidRequest, OAuth2Error, OAuth2ErrorCode, UnauthorizedClient},
+    errors::{
+        InvalidGrant, InvalidRequest, OAuth2Error, OAuth2ErrorCode, ServerError, UnauthorizedClient,
+    },
     requests::{
         AccessTokenRequest, AccessTokenResponse, AuthorizationCodeGrant, RefreshTokenGrant,
     },
@@ -122,12 +127,23 @@ pub fn filter(
         .boxed()
 }
 
-async fn recover(rejection: Rejection) -> Result<Box<dyn Reply>, Rejection> {
-    if let Some(Error { json, status }) = rejection.find::<Error>() {
-        Ok(Box::new(with_status(warp::reply::json(json), *status)))
-    } else {
-        Err(rejection)
+async fn recover(rejection: Rejection) -> Result<Box<dyn Reply>, Infallible> {
+    fn reply<E: OAuth2ErrorCode>(err: E) -> Box<dyn Reply> {
+        let status = err.status();
+        Box::new(with_status(warp::reply::json(&err.into_response()), status))
     }
+
+    if let Some(Error { json, status }) = rejection.find::<Error>() {
+        return Ok(Box::new(with_status(warp::reply::json(json), *status)));
+    }
+
+    if let Some(e) = rejection.find::<RefreshTokenLookupError>() {
+        if e.not_found() {
+            return Ok(reply(InvalidGrant));
+        }
+    };
+
+    Ok(reply(ServerError))
 }
 
 async fn token(
@@ -333,9 +349,8 @@ async fn refresh_token_grant(
     conn: &mut PoolConnection<Postgres>,
 ) -> Result<AccessTokenResponse, Rejection> {
     let mut txn = conn.begin().await.wrap_error()?;
-    let (refresh_token, session) = lookup_active_refresh_token(&mut txn, &grant.refresh_token)
-        .await
-        .wrap_error()?;
+    let (refresh_token, session) =
+        lookup_active_refresh_token(&mut txn, &grant.refresh_token).await?;
 
     if client.client_id != session.client.client_id {
         // As per https://datatracker.ietf.org/doc/html/rfc6749#section-5.2
diff --git a/crates/storage/src/oauth2/refresh_token.rs b/crates/storage/src/oauth2/refresh_token.rs
index d54ed589..18bc1a9c 100644
--- a/crates/storage/src/oauth2/refresh_token.rs
+++ b/crates/storage/src/oauth2/refresh_token.rs
@@ -18,6 +18,8 @@ use mas_data_model::{
     AccessToken, Authentication, BrowserSession, Client, RefreshToken, Session, User, UserEmail,
 };
 use sqlx::PgExecutor;
+use thiserror::Error;
+use warp::reject::Reject;
 
 use crate::{DatabaseInconsistencyError, IdAndCreationTime, PostgresqlBackend};
 
@@ -76,11 +78,28 @@ struct OAuth2RefreshTokenLookup {
     user_email_confirmed_at: Option<DateTime<Utc>>,
 }
 
+#[derive(Error, Debug)]
+#[error("could not lookup refresh token")]
+pub enum RefreshTokenLookupError {
+    Fetch(#[from] sqlx::Error),
+    Conversion(#[from] DatabaseInconsistencyError),
+}
+
+impl Reject for RefreshTokenLookupError {}
+
+impl RefreshTokenLookupError {
+    #[must_use]
+    pub fn not_found(&self) -> bool {
+        matches!(self, Self::Fetch(sqlx::Error::RowNotFound))
+    }
+}
+
 #[allow(clippy::too_many_lines)]
 pub async fn lookup_active_refresh_token(
     executor: impl PgExecutor<'_>,
     token: &str,
-) -> anyhow::Result<(RefreshToken<PostgresqlBackend>, Session<PostgresqlBackend>)> {
+) -> Result<(RefreshToken<PostgresqlBackend>, Session<PostgresqlBackend>), RefreshTokenLookupError>
+{
     let res = sqlx::query_as!(
         OAuth2RefreshTokenLookup,
         r#"
@@ -130,8 +149,7 @@ pub async fn lookup_active_refresh_token(
         token,
     )
     .fetch_one(executor)
-    .await
-    .context("failed to fetch oauth2 refresh token")?;
+    .await?;
 
     let access_token = match (
         res.access_token_id,
@@ -204,11 +222,13 @@ pub async fn lookup_active_refresh_token(
         last_authentication,
     };
 
+    let scope = res.scope.parse().map_err(|_e| DatabaseInconsistencyError)?;
+
     let session = Session {
         data: res.session_id,
         client,
         browser_session,
-        scope: res.scope.parse().context("invalid scope in database")?,
+        scope,
     };
 
     Ok((refresh_token, session))