storage: cleanup access/refresh token lookups

crates/handlers/src/oauth2/introspection.rs

@@ -24,7 +24,8 @@ use mas_keystore::Encrypter;
 use mas_storage::{
     compat::{find_compat_access_token, find_compat_refresh_token, lookup_compat_session},
     oauth2::{
-        access_token::lookup_active_access_token, refresh_token::lookup_active_refresh_token,
+        access_token::find_access_token, refresh_token::lookup_refresh_token,
+        OAuth2SessionRepository,
     },
     user::{BrowserSessionRepository, UserRepository},
     Clock, Repository,
@@ -168,8 +169,17 @@ pub(crate) async fn post(
 
     let reply = match token_type {
         TokenType::AccessToken => {
-            let (token, session) = lookup_active_access_token(&mut conn, token)
+            let token = find_access_token(&mut conn, token)
                 .await?
+                .filter(|t| t.is_valid(clock.now()))
+                .ok_or(RouteError::UnknownToken)?;
+
+            let session = conn
+                .oauth2_session()
+                .lookup(token.session_id)
+                .await?
+                .filter(|s| s.is_valid())
+                // XXX: is that the right error to bubble up?
                 .ok_or(RouteError::UnknownToken)?;
 
             let browser_session = conn
@@ -191,13 +201,22 @@ pub(crate) async fn post(
                 sub: Some(browser_session.user.sub),
                 aud: None,
                 iss: None,
-                jti: None,
+                jti: Some(token.jti()),
             }
         }
 
         TokenType::RefreshToken => {
-            let (token, session) = lookup_active_refresh_token(&mut conn, token)
+            let token = lookup_refresh_token(&mut conn, token)
                 .await?
+                .filter(|t| t.is_valid())
+                .ok_or(RouteError::UnknownToken)?;
+
+            let session = conn
+                .oauth2_session()
+                .lookup(token.session_id)
+                .await?
+                .filter(|s| s.is_valid())
+                // XXX: is that the right error to bubble up?
                 .ok_or(RouteError::UnknownToken)?;
 
             let browser_session = conn
@@ -219,7 +238,7 @@ pub(crate) async fn post(
                 sub: Some(browser_session.user.sub),
                 aud: None,
                 iss: None,
-                jti: None,
+                jti: Some(token.jti()),
             }
         }
 

crates/handlers/src/oauth2/token.rs

@@ -33,9 +33,9 @@ use mas_keystore::{Encrypter, Keystore};
 use mas_router::UrlBuilder;
 use mas_storage::{
     oauth2::{
-        access_token::{add_access_token, revoke_access_token},
+        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_active_refresh_token},
+        refresh_token::{add_refresh_token, consume_refresh_token, lookup_refresh_token},
         OAuth2SessionRepository,
     },
     user::BrowserSessionRepository,
@@ -374,10 +374,20 @@ async fn refresh_token_grant(
 ) -> Result<AccessTokenResponse, RouteError> {
     let (clock, mut rng) = crate::clock_and_rng();
 
-    let (refresh_token, session) = lookup_active_refresh_token(&mut txn, &grant.refresh_token)
+    let refresh_token = lookup_refresh_token(&mut txn, &grant.refresh_token)
         .await?
         .ok_or(RouteError::InvalidGrant)?;
 
+    let session = txn
+        .oauth2_session()
+        .lookup(refresh_token.session_id)
+        .await?
+        .ok_or(RouteError::NoSuchOAuthSession)?;
+
+    if !refresh_token.is_valid() || !session.is_valid() {
+        return Err(RouteError::InvalidGrant);
+    }
+
     if client.id != session.client_id {
         // As per https://datatracker.ietf.org/doc/html/rfc6749#section-5.2
         return Err(RouteError::InvalidGrant);
@@ -407,10 +417,12 @@ async fn refresh_token_grant(
     )
     .await?;
 
-    consume_refresh_token(&mut txn, &clock, &refresh_token).await?;
+    let refresh_token = consume_refresh_token(&mut txn, &clock, refresh_token).await?;
 
     if let Some(access_token_id) = refresh_token.access_token_id {
-        revoke_access_token(&mut txn, &clock, access_token_id).await?;
+        if let Some(access_token) = lookup_access_token(&mut txn, access_token_id).await? {
+            revoke_access_token(&mut txn, &clock, access_token).await?;
+        }
     }
 
     let params = AccessTokenResponse::new(access_token_str)

crates/handlers/src/oauth2/userinfo.rs

@@ -101,10 +101,10 @@ pub async fn get(
     State(key_store): State<Keystore>,
     user_authorization: UserAuthorization,
 ) -> Result<Response, RouteError> {
-    let (_clock, mut rng) = crate::clock_and_rng();
+    let (clock, mut rng) = crate::clock_and_rng();
     let mut conn = pool.acquire().await?;
 
-    let session = user_authorization.protected(&mut conn).await?;
+    let session = user_authorization.protected(&mut conn, clock.now()).await?;
 
     let browser_session = conn
         .browser_session()