diff --git a/crates/axum-utils/src/user_authorization.rs b/crates/axum-utils/src/user_authorization.rs index 66ec7cd3d..eb71ad5a4 100644 --- a/crates/axum-utils/src/user_authorization.rs +++ b/crates/axum-utils/src/user_authorization.rs @@ -117,6 +117,11 @@ impl UserAuthorization { return Err(AuthorizationVerificationError::InvalidToken); } + if !token.is_used() { + // Mark the token as used + repo.oauth2_access_token().mark_used(clock, token).await?; + } + Ok(session) } } diff --git a/crates/data-model/src/tokens.rs b/crates/data-model/src/tokens.rs index b4c8ede28..9e94f7171 100644 --- a/crates/data-model/src/tokens.rs +++ b/crates/data-model/src/tokens.rs @@ -55,6 +55,7 @@ pub struct AccessToken { pub access_token: String, pub created_at: DateTime, pub expires_at: Option>, + pub first_used_at: Option>, } impl AccessToken { @@ -88,6 +89,12 @@ impl AccessToken { } } + /// Whether the access token was used at least once + #[must_use] + pub fn is_used(&self) -> bool { + self.first_used_at.is_some() + } + /// Mark the access token as revoked /// /// # Parameters @@ -109,6 +116,10 @@ pub enum RefreshTokenState { Valid, Consumed { consumed_at: DateTime, + next_refresh_token_id: Option, + }, + Revoked { + revoked_at: DateTime, }, } @@ -117,11 +128,30 @@ impl RefreshTokenState { /// /// # Errors /// - /// Returns an error if the refresh token is already consumed. - fn consume(self, consumed_at: DateTime) -> Result { + /// Returns an error if the refresh token is revoked. + fn consume( + self, + consumed_at: DateTime, + replaced_by: &RefreshToken, + ) -> Result { + match self { + Self::Valid | Self::Consumed { .. } => Ok(Self::Consumed { + consumed_at, + next_refresh_token_id: Some(replaced_by.id), + }), + Self::Revoked { .. } => Err(InvalidTransitionError), + } + } + + /// Revoke the refresh token, returning a new state. + /// + /// # Errors + /// + /// Returns an error if the refresh token is already consumed or revoked. + pub fn revoke(self, revoked_at: DateTime) -> Result { match self { - Self::Valid => Ok(Self::Consumed { consumed_at }), - Self::Consumed { .. } => Err(InvalidTransitionError), + Self::Valid => Ok(Self::Revoked { revoked_at }), + Self::Consumed { .. } | Self::Revoked { .. } => Err(InvalidTransitionError), } } @@ -133,12 +163,16 @@ impl RefreshTokenState { matches!(self, Self::Valid) } - /// Returns `true` if the refresh token state is [`Consumed`]. - /// - /// [`Consumed`]: RefreshTokenState::Consumed + /// Returns the next refresh token ID, if any. #[must_use] - pub fn is_consumed(&self) -> bool { - matches!(self, Self::Consumed { .. }) + pub fn next_refresh_token_id(&self) -> Option { + match self { + Self::Valid | Self::Revoked { .. } => None, + Self::Consumed { + next_refresh_token_id, + .. + } => *next_refresh_token_id, + } } } @@ -170,9 +204,23 @@ impl RefreshToken { /// /// # Errors /// - /// Returns an error if the refresh token is already consumed. - pub fn consume(mut self, consumed_at: DateTime) -> Result { - self.state = self.state.consume(consumed_at)?; + /// Returns an error if the refresh token is revoked. + pub fn consume( + mut self, + consumed_at: DateTime, + replaced_by: &Self, + ) -> Result { + self.state = self.state.consume(consumed_at, replaced_by)?; + Ok(self) + } + + /// Revokes the refresh token and returns a new revoked token + /// + /// # Errors + /// + /// Returns an error if the refresh token is already revoked. + pub fn revoke(mut self, revoked_at: DateTime) -> Result { + self.state = self.state.revoke(revoked_at)?; Ok(self) } } diff --git a/crates/handlers/src/oauth2/introspection.rs b/crates/handlers/src/oauth2/introspection.rs index 4c534f8d3..fcb0a205f 100644 --- a/crates/handlers/src/oauth2/introspection.rs +++ b/crates/handlers/src/oauth2/introspection.rs @@ -207,7 +207,7 @@ pub(crate) async fn post( let reply = match token_type { TokenType::AccessToken => { - let access_token = repo + let mut access_token = repo .oauth2_access_token() .find_by_token(token) .await? @@ -227,6 +227,14 @@ pub(crate) async fn post( return Err(RouteError::InvalidOAuthSession); } + // If this is the first time we're using this token, mark it as used + if !access_token.is_used() { + access_token = repo + .oauth2_access_token() + .mark_used(&clock, access_token) + .await?; + } + // The session might not have a user on it (for Client Credentials grants for // example), so we're optionally fetching the user let (sub, username) = if let Some(user_id) = session.user_id { @@ -443,6 +451,8 @@ pub(crate) async fn post( } }; + repo.save().await?; + Ok(Json(reply)) } @@ -625,6 +635,16 @@ mod tests { .unwrap() .unwrap(); assert_eq!(session.last_active_at, Some(state.clock.now())); + + // And recorded the access token as used + let access_token_lookup = repo + .oauth2_access_token() + .find_by_token(&access_token) + .await + .unwrap() + .unwrap(); + assert!(access_token_lookup.is_used()); + assert_eq!(access_token_lookup.first_used_at, Some(state.clock.now())); repo.cancel().await.unwrap(); // Advance the clock to invalidate the access token diff --git a/crates/handlers/src/oauth2/token.rs b/crates/handlers/src/oauth2/token.rs index a228f746b..ac64da50d 100644 --- a/crates/handlers/src/oauth2/token.rs +++ b/crates/handlers/src/oauth2/token.rs @@ -39,7 +39,7 @@ use oauth2_types::{ scope, }; use thiserror::Error; -use tracing::debug; +use tracing::{debug, info}; use ulid::Ulid; use super::{generate_id_token, generate_token_pair}; @@ -98,6 +98,20 @@ pub(crate) enum RouteError { #[error("failed to load oauth session")] NoSuchOAuthSession, + #[error( + "failed to load the next refresh token ({next:?}) from the previous one ({previous:?})" + )] + NoSuchNextRefreshToken { next: Ulid, previous: Ulid }, + + #[error("failed to load the access token ({access_token:?}) associated with the next refresh token ({refresh_token:?})")] + NoSuchNextAccessToken { + access_token: Ulid, + refresh_token: Ulid, + }, + + #[error("no access token associated with the refresh token {refresh_token:?}")] + NoAccessTokenOnRefreshToken { refresh_token: Ulid }, + #[error("device code grant expired")] DeviceCodeExpired, @@ -122,7 +136,10 @@ impl IntoResponse for RouteError { Self::Internal(_) | Self::NoSuchBrowserSession | Self::NoSuchOAuthSession - | Self::ProvisionDeviceFailed(_) => ( + | Self::ProvisionDeviceFailed(_) + | Self::NoSuchNextRefreshToken { .. } + | Self::NoSuchNextAccessToken { .. } + | Self::NoAccessTokenOnRefreshToken { .. } => ( StatusCode::INTERNAL_SERVER_ERROR, Json(ClientError::from(ClientErrorCode::ServerError)), ), @@ -482,6 +499,7 @@ async fn authorization_code_grant( Ok((params, repo)) } +#[allow(clippy::too_many_lines)] async fn refresh_token_grant( rng: &mut BoxRng, clock: &impl Clock, @@ -518,10 +536,6 @@ async fn refresh_token_grant( .await?; } - if !refresh_token.is_valid() { - return Err(RouteError::RefreshTokenInvalid(refresh_token.id)); - } - if !session.is_valid() { return Err(RouteError::SessionInvalid(session.id)); } @@ -534,6 +548,77 @@ async fn refresh_token_grant( }); } + if !refresh_token.is_valid() { + // We're seing a refresh token that already has been consumed, this might be a + // double-refresh or a replay attack + + // First, get the next refresh token + let Some(next_refresh_token_id) = refresh_token.next_refresh_token_id() else { + // If we don't have a 'next' refresh token, it may just be because this was + // before we were recording those. Let's just treat it as a replay. + return Err(RouteError::RefreshTokenInvalid(refresh_token.id)); + }; + + let Some(next_refresh_token) = repo + .oauth2_refresh_token() + .lookup(next_refresh_token_id) + .await? + else { + return Err(RouteError::NoSuchNextRefreshToken { + next: next_refresh_token_id, + previous: refresh_token.id, + }); + }; + + // Check if the next refresh token was already consumed or not + if !next_refresh_token.is_valid() { + // XXX: This is a replay, we *may* want to invalidate the session + return Err(RouteError::RefreshTokenInvalid(next_refresh_token.id)); + } + + // Check if the associated access token was already used + let Some(access_token_id) = next_refresh_token.access_token_id else { + // This should in theory not happen: this means an access token got cleaned up, + // but the refresh token was still valid. + return Err(RouteError::NoAccessTokenOnRefreshToken { + refresh_token: next_refresh_token.id, + }); + }; + + // Load it + let next_access_token = repo + .oauth2_access_token() + .lookup(access_token_id) + .await? + .ok_or(RouteError::NoSuchNextAccessToken { + access_token: access_token_id, + refresh_token: next_refresh_token_id, + })?; + + if next_access_token.is_used() { + // XXX: This is a replay, we *may* want to invalidate the session + return Err(RouteError::RefreshTokenInvalid(next_refresh_token.id)); + } + + // Looks like it's a double-refresh, client lost their refresh token on + // the way back. Let's revoke the unused access and refresh tokens, and + // issue new ones + info!( + oauth2_session.id = %session.id, + oauth2_client.id = %client.id, + %refresh_token.id, + "Refresh token already used, but issued refresh and access tokens are unused. Assuming those were lost; revoking those and reissuing new ones." + ); + + repo.oauth2_access_token() + .revoke(clock, next_access_token) + .await?; + + repo.oauth2_refresh_token() + .revoke(clock, next_refresh_token) + .await?; + } + activity_tracker .record_oauth2_session(clock, &session) .await; @@ -544,15 +629,18 @@ async fn refresh_token_grant( let refresh_token = repo .oauth2_refresh_token() - .consume(clock, refresh_token) + .consume(clock, refresh_token, &new_refresh_token) .await?; if let Some(access_token_id) = refresh_token.access_token_id { let access_token = repo.oauth2_access_token().lookup(access_token_id).await?; if let Some(access_token) = access_token { - repo.oauth2_access_token() - .revoke(clock, access_token) - .await?; + // If it is a double-refresh, it might already be revoked + if !access_token.state.is_revoked() { + repo.oauth2_access_token() + .revoke(clock, access_token) + .await?; + } } } @@ -1123,6 +1211,175 @@ mod tests { let _: AccessTokenResponse = response.json(); } + #[sqlx::test(migrator = "mas_storage_pg::MIGRATOR")] + async fn test_double_refresh(pool: PgPool) { + setup(); + let state = TestState::from_pool(pool).await.unwrap(); + + // Provision a client + let request = + Request::post(mas_router::OAuth2RegistrationEndpoint::PATH).json(serde_json::json!({ + "client_uri": "https://example.com/", + "redirect_uris": ["https://example.com/callback"], + "token_endpoint_auth_method": "none", + "response_types": ["code"], + "grant_types": ["authorization_code", "refresh_token"], + })); + + let response = state.request(request).await; + response.assert_status(StatusCode::CREATED); + + let ClientRegistrationResponse { client_id, .. } = response.json(); + + // Let's provision a user and create a session for them. This part is hard to + // test with just HTTP requests, so we'll use the repository directly. + let mut repo = state.repository().await.unwrap(); + + let user = repo + .user() + .add(&mut state.rng(), &state.clock, "alice".to_owned()) + .await + .unwrap(); + + let browser_session = repo + .browser_session() + .add(&mut state.rng(), &state.clock, &user, None) + .await + .unwrap(); + + // Lookup the client in the database. + let client = repo + .oauth2_client() + .find_by_client_id(&client_id) + .await + .unwrap() + .unwrap(); + + // Get a token pair + let session = repo + .oauth2_session() + .add_from_browser_session( + &mut state.rng(), + &state.clock, + &client, + &browser_session, + Scope::from_iter([OPENID]), + ) + .await + .unwrap(); + + let (AccessToken { access_token, .. }, RefreshToken { refresh_token, .. }) = + generate_token_pair( + &mut state.rng(), + &state.clock, + &mut repo, + &session, + Duration::microseconds(5 * 60 * 1000 * 1000), + ) + .await + .unwrap(); + + repo.save().await.unwrap(); + + // First check that the token is valid + assert!(state.is_access_token_valid(&access_token).await); + + // Now call the token endpoint to get an access token. + let request = + Request::post(mas_router::OAuth2TokenEndpoint::PATH).form(serde_json::json!({ + "grant_type": "refresh_token", + "refresh_token": refresh_token, + "client_id": client.client_id, + })); + + let first_response = state.request(request).await; + first_response.assert_status(StatusCode::OK); + let first_response: AccessTokenResponse = first_response.json(); + + // Call a second time, it should work, as we haven't done anything yet with the + // token + let request = + Request::post(mas_router::OAuth2TokenEndpoint::PATH).form(serde_json::json!({ + "grant_type": "refresh_token", + "refresh_token": refresh_token, + "client_id": client.client_id, + })); + + let second_response = state.request(request).await; + second_response.assert_status(StatusCode::OK); + let second_response: AccessTokenResponse = second_response.json(); + + // Check that we got new tokens + assert_ne!(first_response.access_token, second_response.access_token); + assert_ne!(first_response.refresh_token, second_response.refresh_token); + + // Check that the old-new token is invalid + assert!( + !state + .is_access_token_valid(&first_response.access_token) + .await + ); + + // Check that the new-new token is valid + assert!( + state + .is_access_token_valid(&second_response.access_token) + .await + ); + + // Do a third refresh, this one should not work, as we've used the new + // access token + let request = + Request::post(mas_router::OAuth2TokenEndpoint::PATH).form(serde_json::json!({ + "grant_type": "refresh_token", + "refresh_token": refresh_token, + "client_id": client.client_id, + })); + + let third_response = state.request(request).await; + third_response.assert_status(StatusCode::BAD_REQUEST); + + // The other reason we consider a new refresh token to be 'used' is if + // it was already used in a refresh + // So, if we do a refresh with the second_response.refresh_token, then + // another refresh with the result, redoing one with + // second_response.refresh_token again should fail + let request = + Request::post(mas_router::OAuth2TokenEndpoint::PATH).form(serde_json::json!({ + "grant_type": "refresh_token", + "refresh_token": second_response.refresh_token, + "client_id": client.client_id, + })); + + // This one is fine + let fourth_response = state.request(request).await; + fourth_response.assert_status(StatusCode::OK); + let fourth_response: AccessTokenResponse = fourth_response.json(); + + // Do another one, it should be fine as well + let request = + Request::post(mas_router::OAuth2TokenEndpoint::PATH).form(serde_json::json!({ + "grant_type": "refresh_token", + "refresh_token": fourth_response.refresh_token, + "client_id": client.client_id, + })); + + let fifth_response = state.request(request).await; + fifth_response.assert_status(StatusCode::OK); + + // But now, if we re-do with the second_response.refresh_token, it should + // fail + let request = + Request::post(mas_router::OAuth2TokenEndpoint::PATH).form(serde_json::json!({ + "grant_type": "refresh_token", + "refresh_token": second_response.refresh_token, + "client_id": client.client_id, + })); + + let sixth_response = state.request(request).await; + sixth_response.assert_status(StatusCode::BAD_REQUEST); + } + #[sqlx::test(migrator = "mas_storage_pg::MIGRATOR")] async fn test_client_credentials(pool: PgPool) { setup(); diff --git a/crates/handlers/src/oauth2/userinfo.rs b/crates/handlers/src/oauth2/userinfo.rs index 6618f063b..c3223482b 100644 --- a/crates/handlers/src/oauth2/userinfo.rs +++ b/crates/handlers/src/oauth2/userinfo.rs @@ -142,6 +142,8 @@ pub async fn get( .await? .ok_or(RouteError::NoSuchClient)?; + repo.save().await?; + if let Some(alg) = client.userinfo_signed_response_alg { let key = key_store .signing_key_for_algorithm(&alg) diff --git a/crates/storage-pg/.sqlx/query-b6a6f5386dc89e4bc2ce56d578a29341848fce336d339b6bbf425956f5ed5032.json b/crates/storage-pg/.sqlx/query-66693f31eff5673e88ca516ee727a709b06455e08b9fd75cc08f142070f330b3.json similarity index 57% rename from crates/storage-pg/.sqlx/query-b6a6f5386dc89e4bc2ce56d578a29341848fce336d339b6bbf425956f5ed5032.json rename to crates/storage-pg/.sqlx/query-66693f31eff5673e88ca516ee727a709b06455e08b9fd75cc08f142070f330b3.json index 30aa4caf3..716bbf8c7 100644 --- a/crates/storage-pg/.sqlx/query-b6a6f5386dc89e4bc2ce56d578a29341848fce336d339b6bbf425956f5ed5032.json +++ b/crates/storage-pg/.sqlx/query-66693f31eff5673e88ca516ee727a709b06455e08b9fd75cc08f142070f330b3.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "\n UPDATE oauth2_refresh_tokens\n SET consumed_at = $2\n WHERE oauth2_refresh_token_id = $1\n ", + "query": "\n UPDATE oauth2_refresh_tokens\n SET revoked_at = $2\n WHERE oauth2_refresh_token_id = $1\n ", "describe": { "columns": [], "parameters": { @@ -11,5 +11,5 @@ }, "nullable": [] }, - "hash": "b6a6f5386dc89e4bc2ce56d578a29341848fce336d339b6bbf425956f5ed5032" + "hash": "66693f31eff5673e88ca516ee727a709b06455e08b9fd75cc08f142070f330b3" } diff --git a/crates/storage-pg/.sqlx/query-a6fa7811d0a7c62c7cccff96dc82db5b25462fa7669fde1941ccab4712585b20.json b/crates/storage-pg/.sqlx/query-6d71188dffc492ddc8f7f21476516d3b08fd5d736ecf36845e6fd4bfc515b2cf.json similarity index 63% rename from crates/storage-pg/.sqlx/query-a6fa7811d0a7c62c7cccff96dc82db5b25462fa7669fde1941ccab4712585b20.json rename to crates/storage-pg/.sqlx/query-6d71188dffc492ddc8f7f21476516d3b08fd5d736ecf36845e6fd4bfc515b2cf.json index bfa6485e8..1acfa2784 100644 --- a/crates/storage-pg/.sqlx/query-a6fa7811d0a7c62c7cccff96dc82db5b25462fa7669fde1941ccab4712585b20.json +++ b/crates/storage-pg/.sqlx/query-6d71188dffc492ddc8f7f21476516d3b08fd5d736ecf36845e6fd4bfc515b2cf.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "\n SELECT oauth2_refresh_token_id\n , refresh_token\n , created_at\n , consumed_at\n , oauth2_access_token_id\n , oauth2_session_id\n FROM oauth2_refresh_tokens\n\n WHERE oauth2_refresh_token_id = $1\n ", + "query": "\n SELECT oauth2_refresh_token_id\n , refresh_token\n , created_at\n , consumed_at\n , revoked_at\n , oauth2_access_token_id\n , oauth2_session_id\n , next_oauth2_refresh_token_id\n FROM oauth2_refresh_tokens\n\n WHERE oauth2_refresh_token_id = $1\n ", "describe": { "columns": [ { @@ -25,13 +25,23 @@ }, { "ordinal": 4, + "name": "revoked_at", + "type_info": "Timestamptz" + }, + { + "ordinal": 5, "name": "oauth2_access_token_id", "type_info": "Uuid" }, { - "ordinal": 5, + "ordinal": 6, "name": "oauth2_session_id", "type_info": "Uuid" + }, + { + "ordinal": 7, + "name": "next_oauth2_refresh_token_id", + "type_info": "Uuid" } ], "parameters": { @@ -45,8 +55,10 @@ false, true, true, - false + true, + false, + true ] }, - "hash": "a6fa7811d0a7c62c7cccff96dc82db5b25462fa7669fde1941ccab4712585b20" + "hash": "6d71188dffc492ddc8f7f21476516d3b08fd5d736ecf36845e6fd4bfc515b2cf" } diff --git a/crates/storage-pg/.sqlx/query-7189b6136fd08ac9ae7c51bff06fb2254d1bf9e8a97cd7d32ba789c740e0fbdb.json b/crates/storage-pg/.sqlx/query-7189b6136fd08ac9ae7c51bff06fb2254d1bf9e8a97cd7d32ba789c740e0fbdb.json new file mode 100644 index 000000000..a0695db46 --- /dev/null +++ b/crates/storage-pg/.sqlx/query-7189b6136fd08ac9ae7c51bff06fb2254d1bf9e8a97cd7d32ba789c740e0fbdb.json @@ -0,0 +1,15 @@ +{ + "db_name": "PostgreSQL", + "query": "\n UPDATE oauth2_access_tokens\n SET first_used_at = $2\n WHERE oauth2_access_token_id = $1\n ", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Uuid", + "Timestamptz" + ] + }, + "nullable": [] + }, + "hash": "7189b6136fd08ac9ae7c51bff06fb2254d1bf9e8a97cd7d32ba789c740e0fbdb" +} diff --git a/crates/storage-pg/.sqlx/query-dd16942318bf38d9a245b2c86fedd3cbd6b65e7a13465552d79cd3c022122fd4.json b/crates/storage-pg/.sqlx/query-875294dc5cf87bcf302fb9e87933745cc1c57bbe3c3c69110592a07400116c7f.json similarity index 71% rename from crates/storage-pg/.sqlx/query-dd16942318bf38d9a245b2c86fedd3cbd6b65e7a13465552d79cd3c022122fd4.json rename to crates/storage-pg/.sqlx/query-875294dc5cf87bcf302fb9e87933745cc1c57bbe3c3c69110592a07400116c7f.json index c9bd70555..b82d7c463 100644 --- a/crates/storage-pg/.sqlx/query-dd16942318bf38d9a245b2c86fedd3cbd6b65e7a13465552d79cd3c022122fd4.json +++ b/crates/storage-pg/.sqlx/query-875294dc5cf87bcf302fb9e87933745cc1c57bbe3c3c69110592a07400116c7f.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "\n SELECT oauth2_access_token_id\n , access_token\n , created_at\n , expires_at\n , revoked_at\n , oauth2_session_id\n\n FROM oauth2_access_tokens\n\n WHERE oauth2_access_token_id = $1\n ", + "query": "\n SELECT oauth2_access_token_id\n , access_token\n , created_at\n , expires_at\n , revoked_at\n , oauth2_session_id\n , first_used_at\n\n FROM oauth2_access_tokens\n\n WHERE oauth2_access_token_id = $1\n ", "describe": { "columns": [ { @@ -32,6 +32,11 @@ "ordinal": 5, "name": "oauth2_session_id", "type_info": "Uuid" + }, + { + "ordinal": 6, + "name": "first_used_at", + "type_info": "Timestamptz" } ], "parameters": { @@ -45,8 +50,9 @@ false, true, true, - false + false, + true ] }, - "hash": "dd16942318bf38d9a245b2c86fedd3cbd6b65e7a13465552d79cd3c022122fd4" + "hash": "875294dc5cf87bcf302fb9e87933745cc1c57bbe3c3c69110592a07400116c7f" } diff --git a/crates/storage-pg/.sqlx/query-e709869c062ac50248b1f9f8f808cc2f5e7bef58a6c2e42a7bb0c1cb8f508671.json b/crates/storage-pg/.sqlx/query-a75a6a08c9639053cfc3cffa9d4a009785f358b334f5c586c2e358f0d0b4d856.json similarity index 64% rename from crates/storage-pg/.sqlx/query-e709869c062ac50248b1f9f8f808cc2f5e7bef58a6c2e42a7bb0c1cb8f508671.json rename to crates/storage-pg/.sqlx/query-a75a6a08c9639053cfc3cffa9d4a009785f358b334f5c586c2e358f0d0b4d856.json index 411b02cbe..c7146b7d1 100644 --- a/crates/storage-pg/.sqlx/query-e709869c062ac50248b1f9f8f808cc2f5e7bef58a6c2e42a7bb0c1cb8f508671.json +++ b/crates/storage-pg/.sqlx/query-a75a6a08c9639053cfc3cffa9d4a009785f358b334f5c586c2e358f0d0b4d856.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "\n SELECT oauth2_refresh_token_id\n , refresh_token\n , created_at\n , consumed_at\n , oauth2_access_token_id\n , oauth2_session_id\n FROM oauth2_refresh_tokens\n\n WHERE refresh_token = $1\n ", + "query": "\n SELECT oauth2_refresh_token_id\n , refresh_token\n , created_at\n , consumed_at\n , revoked_at\n , oauth2_access_token_id\n , oauth2_session_id\n , next_oauth2_refresh_token_id\n FROM oauth2_refresh_tokens\n\n WHERE refresh_token = $1\n ", "describe": { "columns": [ { @@ -25,13 +25,23 @@ }, { "ordinal": 4, + "name": "revoked_at", + "type_info": "Timestamptz" + }, + { + "ordinal": 5, "name": "oauth2_access_token_id", "type_info": "Uuid" }, { - "ordinal": 5, + "ordinal": 6, "name": "oauth2_session_id", "type_info": "Uuid" + }, + { + "ordinal": 7, + "name": "next_oauth2_refresh_token_id", + "type_info": "Uuid" } ], "parameters": { @@ -45,8 +55,10 @@ false, true, true, - false + true, + false, + true ] }, - "hash": "e709869c062ac50248b1f9f8f808cc2f5e7bef58a6c2e42a7bb0c1cb8f508671" + "hash": "a75a6a08c9639053cfc3cffa9d4a009785f358b334f5c586c2e358f0d0b4d856" } diff --git a/crates/storage-pg/.sqlx/query-477f79556e5777b38feb85013b4f04dbb8230e4b0b0bcc45f669d7b8d0b91db4.json b/crates/storage-pg/.sqlx/query-c09e0bb0378d9dfb15de7f2f1209fab6ea87589819128e6fc9ed5da11dfc2770.json similarity index 72% rename from crates/storage-pg/.sqlx/query-477f79556e5777b38feb85013b4f04dbb8230e4b0b0bcc45f669d7b8d0b91db4.json rename to crates/storage-pg/.sqlx/query-c09e0bb0378d9dfb15de7f2f1209fab6ea87589819128e6fc9ed5da11dfc2770.json index bc6c2fb72..2ef7b86cf 100644 --- a/crates/storage-pg/.sqlx/query-477f79556e5777b38feb85013b4f04dbb8230e4b0b0bcc45f669d7b8d0b91db4.json +++ b/crates/storage-pg/.sqlx/query-c09e0bb0378d9dfb15de7f2f1209fab6ea87589819128e6fc9ed5da11dfc2770.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "\n SELECT oauth2_access_token_id\n , access_token\n , created_at\n , expires_at\n , revoked_at\n , oauth2_session_id\n\n FROM oauth2_access_tokens\n\n WHERE access_token = $1\n ", + "query": "\n SELECT oauth2_access_token_id\n , access_token\n , created_at\n , expires_at\n , revoked_at\n , oauth2_session_id\n , first_used_at\n\n FROM oauth2_access_tokens\n\n WHERE access_token = $1\n ", "describe": { "columns": [ { @@ -32,6 +32,11 @@ "ordinal": 5, "name": "oauth2_session_id", "type_info": "Uuid" + }, + { + "ordinal": 6, + "name": "first_used_at", + "type_info": "Timestamptz" } ], "parameters": { @@ -45,8 +50,9 @@ false, true, true, - false + false, + true ] }, - "hash": "477f79556e5777b38feb85013b4f04dbb8230e4b0b0bcc45f669d7b8d0b91db4" + "hash": "c09e0bb0378d9dfb15de7f2f1209fab6ea87589819128e6fc9ed5da11dfc2770" } diff --git a/crates/storage-pg/.sqlx/query-1a8701f5672de052bb766933f60b93249acc7237b996e8b93cd61b9f69c902ff.json b/crates/storage-pg/.sqlx/query-c68b28232b42a62907709403caafe1cc5267780232cd615468d50f5c0af2bedb.json similarity index 63% rename from crates/storage-pg/.sqlx/query-1a8701f5672de052bb766933f60b93249acc7237b996e8b93cd61b9f69c902ff.json rename to crates/storage-pg/.sqlx/query-c68b28232b42a62907709403caafe1cc5267780232cd615468d50f5c0af2bedb.json index 07ab1afa8..5da8bf5bc 100644 --- a/crates/storage-pg/.sqlx/query-1a8701f5672de052bb766933f60b93249acc7237b996e8b93cd61b9f69c902ff.json +++ b/crates/storage-pg/.sqlx/query-c68b28232b42a62907709403caafe1cc5267780232cd615468d50f5c0af2bedb.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "\n DELETE FROM oauth2_access_tokens\n WHERE expires_at < $1\n ", + "query": "\n DELETE FROM oauth2_access_tokens\n WHERE revoked_at < $1\n ", "describe": { "columns": [], "parameters": { @@ -10,5 +10,5 @@ }, "nullable": [] }, - "hash": "1a8701f5672de052bb766933f60b93249acc7237b996e8b93cd61b9f69c902ff" + "hash": "c68b28232b42a62907709403caafe1cc5267780232cd615468d50f5c0af2bedb" } diff --git a/crates/storage-pg/.sqlx/query-ffbfef8b7e72ec4bae02b6bbe862980b5fe575ae8432a000e9c4e4307caa2d9b.json b/crates/storage-pg/.sqlx/query-ffbfef8b7e72ec4bae02b6bbe862980b5fe575ae8432a000e9c4e4307caa2d9b.json new file mode 100644 index 000000000..77f807125 --- /dev/null +++ b/crates/storage-pg/.sqlx/query-ffbfef8b7e72ec4bae02b6bbe862980b5fe575ae8432a000e9c4e4307caa2d9b.json @@ -0,0 +1,16 @@ +{ + "db_name": "PostgreSQL", + "query": "\n UPDATE oauth2_refresh_tokens\n SET consumed_at = $2,\n next_oauth2_refresh_token_id = $3\n WHERE oauth2_refresh_token_id = $1\n ", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Uuid", + "Timestamptz", + "Uuid" + ] + }, + "nullable": [] + }, + "hash": "ffbfef8b7e72ec4bae02b6bbe862980b5fe575ae8432a000e9c4e4307caa2d9b" +} diff --git a/crates/storage-pg/migrations/20241210115428_oauth_refresh_token_track_next.sql b/crates/storage-pg/migrations/20241210115428_oauth_refresh_token_track_next.sql new file mode 100644 index 000000000..b1bd2c3f6 --- /dev/null +++ b/crates/storage-pg/migrations/20241210115428_oauth_refresh_token_track_next.sql @@ -0,0 +1,9 @@ +-- Copyright 2024 New Vector Ltd. +-- +-- SPDX-License-Identifier: AGPL-3.0-only +-- Please see LICENSE in the repository root for full details. + +-- Add a reference to the 'next' refresh token when it was consumed and replaced +ALTER TABLE oauth2_refresh_tokens + ADD COLUMN "next_oauth2_refresh_token_id" UUID + REFERENCES oauth2_refresh_tokens (oauth2_refresh_token_id); diff --git a/crates/storage-pg/migrations/20241210133651_oauth2_access_token_first_used.sql b/crates/storage-pg/migrations/20241210133651_oauth2_access_token_first_used.sql new file mode 100644 index 000000000..8d9507a69 --- /dev/null +++ b/crates/storage-pg/migrations/20241210133651_oauth2_access_token_first_used.sql @@ -0,0 +1,8 @@ +-- Copyright 2024 New Vector Ltd. +-- +-- SPDX-License-Identifier: AGPL-3.0-only +-- Please see LICENSE in the repository root for full details. + +-- Track when the access token was first used. A NULL value means it was never used. +ALTER TABLE oauth2_access_tokens + ADD COLUMN "first_used_at" TIMESTAMP WITH TIME ZONE; diff --git a/crates/storage-pg/src/oauth2/access_token.rs b/crates/storage-pg/src/oauth2/access_token.rs index ea8be74d3..48801ead6 100644 --- a/crates/storage-pg/src/oauth2/access_token.rs +++ b/crates/storage-pg/src/oauth2/access_token.rs @@ -36,6 +36,7 @@ struct OAuth2AccessTokenLookup { created_at: DateTime, expires_at: Option>, revoked_at: Option>, + first_used_at: Option>, } impl From for AccessToken { @@ -52,6 +53,7 @@ impl From for AccessToken { access_token: value.access_token, created_at: value.created_at, expires_at: value.expires_at, + first_used_at: value.first_used_at, } } } @@ -70,6 +72,7 @@ impl OAuth2AccessTokenRepository for PgOAuth2AccessTokenRepository<'_> { , expires_at , revoked_at , oauth2_session_id + , first_used_at FROM oauth2_access_tokens @@ -106,6 +109,7 @@ impl OAuth2AccessTokenRepository for PgOAuth2AccessTokenRepository<'_> { , expires_at , revoked_at , oauth2_session_id + , first_used_at FROM oauth2_access_tokens @@ -170,9 +174,20 @@ impl OAuth2AccessTokenRepository for PgOAuth2AccessTokenRepository<'_> { session_id: session.id, created_at, expires_at, + first_used_at: None, }) } + #[tracing::instrument( + name = "db.oauth2_access_token.revoke", + skip_all, + fields( + db.query.text, + session.id = %access_token.session_id, + %access_token.id, + ), + err, + )] async fn revoke( &mut self, clock: &dyn Clock, @@ -188,6 +203,7 @@ impl OAuth2AccessTokenRepository for PgOAuth2AccessTokenRepository<'_> { Uuid::from(access_token.id), revoked_at, ) + .traced() .execute(&mut *self.conn) .await?; @@ -198,16 +214,60 @@ impl OAuth2AccessTokenRepository for PgOAuth2AccessTokenRepository<'_> { .map_err(DatabaseError::to_invalid_operation) } - async fn cleanup_expired(&mut self, clock: &dyn Clock) -> Result { - // Cleanup token which expired more than 15 minutes ago - let threshold = clock.now() - Duration::microseconds(15 * 60 * 1000 * 1000); + #[tracing::instrument( + name = "db.oauth2_access_token.mark_used", + skip_all, + fields( + db.query.text, + session.id = %access_token.session_id, + %access_token.id, + ), + err, + )] + async fn mark_used( + &mut self, + clock: &dyn Clock, + mut access_token: AccessToken, + ) -> Result { + let now = clock.now(); + let res = sqlx::query!( + r#" + UPDATE oauth2_access_tokens + SET first_used_at = $2 + WHERE oauth2_access_token_id = $1 + "#, + Uuid::from(access_token.id), + now, + ) + .execute(&mut *self.conn) + .await?; + + DatabaseError::ensure_affected_rows(&res, 1)?; + + access_token.first_used_at = Some(now); + + Ok(access_token) + } + + #[tracing::instrument( + name = "db.oauth2_access_token.cleanup_revoked", + skip_all, + fields( + db.query.text, + ), + err, + )] + async fn cleanup_revoked(&mut self, clock: &dyn Clock) -> Result { + // Cleanup token that were revoked more than an hour ago + let threshold = clock.now() - Duration::microseconds(60 * 60 * 1000 * 1000); let res = sqlx::query!( r#" DELETE FROM oauth2_access_tokens - WHERE expires_at < $1 + WHERE revoked_at < $1 "#, threshold, ) + .traced() .execute(&mut *self.conn) .await?; diff --git a/crates/storage-pg/src/oauth2/mod.rs b/crates/storage-pg/src/oauth2/mod.rs index c55b4d70e..54ce32fb7 100644 --- a/crates/storage-pg/src/oauth2/mod.rs +++ b/crates/storage-pg/src/oauth2/mod.rs @@ -341,6 +341,19 @@ mod tests { clock.advance(Duration::try_minutes(-6).unwrap()); // Go back in time assert!(access_token.is_valid(clock.now())); + // Create a new refresh token to be able to consume the old one + let new_refresh_token = repo + .oauth2_refresh_token() + .add( + &mut rng, + &clock, + &session, + &access_token, + "ddeeff".to_owned(), + ) + .await + .unwrap(); + // Mark the access token as revoked let access_token = repo .oauth2_access_token() @@ -353,7 +366,7 @@ mod tests { assert!(refresh_token.is_valid()); let refresh_token = repo .oauth2_refresh_token() - .consume(&clock, refresh_token) + .consume(&clock, refresh_token, &new_refresh_token) .await .unwrap(); assert!(!refresh_token.is_valid()); diff --git a/crates/storage-pg/src/oauth2/refresh_token.rs b/crates/storage-pg/src/oauth2/refresh_token.rs index bc0480a67..ee4412ffd 100644 --- a/crates/storage-pg/src/oauth2/refresh_token.rs +++ b/crates/storage-pg/src/oauth2/refresh_token.rs @@ -13,7 +13,7 @@ use sqlx::PgConnection; use ulid::Ulid; use uuid::Uuid; -use crate::{tracing::ExecuteExt, DatabaseError}; +use crate::{tracing::ExecuteExt, DatabaseError, DatabaseInconsistencyError}; /// An implementation of [`OAuth2RefreshTokenRepository`] for a PostgreSQL /// connection @@ -34,25 +34,47 @@ struct OAuth2RefreshTokenLookup { refresh_token: String, created_at: DateTime, consumed_at: Option>, + revoked_at: Option>, oauth2_access_token_id: Option, oauth2_session_id: Uuid, + next_oauth2_refresh_token_id: Option, } -impl From for RefreshToken { - fn from(value: OAuth2RefreshTokenLookup) -> Self { - let state = match value.consumed_at { - None => RefreshTokenState::Valid, - Some(consumed_at) => RefreshTokenState::Consumed { consumed_at }, +impl TryFrom for RefreshToken { + type Error = DatabaseInconsistencyError; + + fn try_from(value: OAuth2RefreshTokenLookup) -> Result { + let id = value.oauth2_refresh_token_id.into(); + let state = match ( + value.revoked_at, + value.consumed_at, + value.next_oauth2_refresh_token_id, + ) { + (None, None, None) => RefreshTokenState::Valid, + (Some(revoked_at), None, None) => RefreshTokenState::Revoked { revoked_at }, + (None, Some(consumed_at), None) => RefreshTokenState::Consumed { + consumed_at, + next_refresh_token_id: None, + }, + (None, Some(consumed_at), Some(id)) => RefreshTokenState::Consumed { + consumed_at, + next_refresh_token_id: Some(Ulid::from(id)), + }, + _ => { + return Err(DatabaseInconsistencyError::on("oauth2_refresh_tokens") + .column("next_oauth2_refresh_token_id") + .row(id)) + } }; - RefreshToken { - id: value.oauth2_refresh_token_id.into(), + Ok(RefreshToken { + id, 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), - } + }) } } @@ -77,20 +99,23 @@ impl OAuth2RefreshTokenRepository for PgOAuth2RefreshTokenRepository<'_> { , refresh_token , created_at , consumed_at + , revoked_at , oauth2_access_token_id , oauth2_session_id + , next_oauth2_refresh_token_id FROM oauth2_refresh_tokens WHERE oauth2_refresh_token_id = $1 "#, Uuid::from(id), ) + .traced() .fetch_optional(&mut *self.conn) .await?; let Some(res) = res else { return Ok(None) }; - Ok(Some(res.into())) + Ok(Some(res.try_into()?)) } #[tracing::instrument( @@ -112,8 +137,10 @@ impl OAuth2RefreshTokenRepository for PgOAuth2RefreshTokenRepository<'_> { , refresh_token , created_at , consumed_at + , revoked_at , oauth2_access_token_id , oauth2_session_id + , next_oauth2_refresh_token_id FROM oauth2_refresh_tokens WHERE refresh_token = $1 @@ -126,7 +153,7 @@ impl OAuth2RefreshTokenRepository for PgOAuth2RefreshTokenRepository<'_> { let Some(res) = res else { return Ok(None) }; - Ok(Some(res.into())) + Ok(Some(res.try_into()?)) } #[tracing::instrument( @@ -194,24 +221,64 @@ impl OAuth2RefreshTokenRepository for PgOAuth2RefreshTokenRepository<'_> { &mut self, clock: &dyn Clock, refresh_token: RefreshToken, + replaced_by: &RefreshToken, ) -> Result { let consumed_at = clock.now(); let res = sqlx::query!( r#" UPDATE oauth2_refresh_tokens - SET consumed_at = $2 + SET consumed_at = $2, + next_oauth2_refresh_token_id = $3 WHERE oauth2_refresh_token_id = $1 "#, Uuid::from(refresh_token.id), consumed_at, + Uuid::from(replaced_by.id), + ) + .traced() + .execute(&mut *self.conn) + .await?; + + DatabaseError::ensure_affected_rows(&res, 1)?; + + refresh_token + .consume(consumed_at, replaced_by) + .map_err(DatabaseError::to_invalid_operation) + } + + #[tracing::instrument( + name = "db.oauth2_refresh_token.revoke", + skip_all, + fields( + db.query.text, + %refresh_token.id, + session.id = %refresh_token.session_id, + ), + err, + )] + async fn revoke( + &mut self, + clock: &dyn Clock, + refresh_token: RefreshToken, + ) -> Result { + let revoked_at = clock.now(); + let res = sqlx::query!( + r#" + UPDATE oauth2_refresh_tokens + SET revoked_at = $2 + WHERE oauth2_refresh_token_id = $1 + "#, + Uuid::from(refresh_token.id), + revoked_at, ) + .traced() .execute(&mut *self.conn) .await?; DatabaseError::ensure_affected_rows(&res, 1)?; refresh_token - .consume(consumed_at) + .revoke(revoked_at) .map_err(DatabaseError::to_invalid_operation) } } diff --git a/crates/storage/src/oauth2/access_token.rs b/crates/storage/src/oauth2/access_token.rs index 79354ef6b..989dd9861 100644 --- a/crates/storage/src/oauth2/access_token.rs +++ b/crates/storage/src/oauth2/access_token.rs @@ -91,7 +91,23 @@ pub trait OAuth2AccessTokenRepository: Send + Sync { access_token: AccessToken, ) -> Result; - /// Cleanup expired access tokens + /// Mark the access token as used, to track when it was first used + /// + /// # Parameters + /// + /// * `clock`: The clock used to generate timestamps + /// * `access_token`: The access token to mark as used + /// + /// # Errors + /// + /// Returns [`Self::Error`] if the underlying repository fails + async fn mark_used( + &mut self, + clock: &dyn Clock, + access_token: AccessToken, + ) -> Result; + + /// Cleanup revoked access tokens /// /// Returns the number of access tokens that were cleaned up /// @@ -102,7 +118,7 @@ pub trait OAuth2AccessTokenRepository: Send + Sync { /// # Errors /// /// Returns [`Self::Error`] if the underlying repository fails - async fn cleanup_expired(&mut self, clock: &dyn Clock) -> Result; + async fn cleanup_revoked(&mut self, clock: &dyn Clock) -> Result; } repository_impl!(OAuth2AccessTokenRepository: @@ -128,5 +144,11 @@ repository_impl!(OAuth2AccessTokenRepository: access_token: AccessToken, ) -> Result; - async fn cleanup_expired(&mut self, clock: &dyn Clock) -> Result; + async fn mark_used( + &mut self, + clock: &dyn Clock, + access_token: AccessToken, + ) -> Result; + + async fn cleanup_revoked(&mut self, clock: &dyn Clock) -> Result; ); diff --git a/crates/storage/src/oauth2/refresh_token.rs b/crates/storage/src/oauth2/refresh_token.rs index c003e9270..6e240f7b9 100644 --- a/crates/storage/src/oauth2/refresh_token.rs +++ b/crates/storage/src/oauth2/refresh_token.rs @@ -80,15 +80,36 @@ pub trait OAuth2RefreshTokenRepository: Send + Sync { /// /// * `clock`: The clock used to generate timestamps /// * `refresh_token`: The [`RefreshToken`] to consume + /// * `replaced_by`: The [`RefreshToken`] which replaced this one /// /// # Errors /// /// Returns [`Self::Error`] if the underlying repository fails, or if the - /// token was already consumed + /// token was already consumed or revoked async fn consume( &mut self, clock: &dyn Clock, refresh_token: RefreshToken, + replaced_by: &RefreshToken, + ) -> Result; + + /// Revoke a refresh token + /// + /// Returns the updated [`RefreshToken`] + /// + /// # Parameters + /// + /// * `clock`: The clock used to generate timestamps + /// * `refresh_token`: The [`RefreshToken`] to revoke + /// + /// # Errors + /// + /// Returns [`Self::Error`] if the underlying repository fails, or if the + /// token was already revoked or consumed + async fn revoke( + &mut self, + clock: &dyn Clock, + refresh_token: RefreshToken, ) -> Result; } @@ -113,5 +134,12 @@ repository_impl!(OAuth2RefreshTokenRepository: &mut self, clock: &dyn Clock, refresh_token: RefreshToken, + replaced_by: &RefreshToken, + ) -> Result; + + async fn revoke( + &mut self, + clock: &dyn Clock, + refresh_token: RefreshToken, ) -> Result; ); diff --git a/crates/tasks/src/database.rs b/crates/tasks/src/database.rs index b64e1defb..74903fcd1 100644 --- a/crates/tasks/src/database.rs +++ b/crates/tasks/src/database.rs @@ -24,7 +24,7 @@ impl RunnableJob for CleanupExpiredTokensJob { let count = repo .oauth2_access_token() - .cleanup_expired(&clock) + .cleanup_revoked(&clock) .await .map_err(JobError::retry)?; repo.save().await.map_err(JobError::retry)?; @@ -32,7 +32,7 @@ impl RunnableJob for CleanupExpiredTokensJob { if count == 0 { debug!("no token to clean up"); } else { - info!(count, "cleaned up expired tokens"); + info!(count, "cleaned up revoked tokens"); } Ok(()) diff --git a/crates/tasks/src/lib.rs b/crates/tasks/src/lib.rs index ecfd1ca18..eb9f19306 100644 --- a/crates/tasks/src/lib.rs +++ b/crates/tasks/src/lib.rs @@ -121,7 +121,7 @@ pub async fn init( .register_handler::() .add_schedule( "cleanup-expired-tokens", - "*/15 * * * * *".parse()?, + "0 0 * * * *".parse()?, mas_storage::queue::CleanupExpiredTokensJob, );