From b3756e4ae4446a92b8c41e1e105c63a4a3825065 Mon Sep 17 00:00:00 2001 From: Quentin Gliech Date: Tue, 10 Dec 2024 14:06:06 +0100 Subject: [PATCH] Record the next refresh token ID when refreshing This will help us determine whether we had a double-refresh happening --- crates/data-model/src/tokens.rs | 32 ++++++++++-- crates/handlers/src/oauth2/token.rs | 2 +- ...af36c103e030358262e46a4bd5e4006dc630.json} | 12 +++-- ...019c36bf5983182b0871de6e85036d8df466.json} | 12 +++-- ...29341848fce336d339b6bbf425956f5ed5032.json | 15 ------ ...2980b5fe575ae8432a000e9c4e4307caa2d9b.json | 16 ++++++ ...0115428_oauth_refresh_token_track_next.sql | 9 ++++ crates/storage-pg/src/oauth2/mod.rs | 15 +++++- crates/storage-pg/src/oauth2/refresh_token.rs | 49 ++++++++++++++----- crates/storage/src/oauth2/refresh_token.rs | 3 ++ 10 files changed, 125 insertions(+), 40 deletions(-) rename crates/storage-pg/.sqlx/{query-e709869c062ac50248b1f9f8f808cc2f5e7bef58a6c2e42a7bb0c1cb8f508671.json => query-265a981142194216593cd09cd8f6af36c103e030358262e46a4bd5e4006dc630.json} (74%) rename crates/storage-pg/.sqlx/{query-a6fa7811d0a7c62c7cccff96dc82db5b25462fa7669fde1941ccab4712585b20.json => query-5b94c22d44692a16fa5a6edd5dac019c36bf5983182b0871de6e85036d8df466.json} (74%) delete mode 100644 crates/storage-pg/.sqlx/query-b6a6f5386dc89e4bc2ce56d578a29341848fce336d339b6bbf425956f5ed5032.json create mode 100644 crates/storage-pg/.sqlx/query-ffbfef8b7e72ec4bae02b6bbe862980b5fe575ae8432a000e9c4e4307caa2d9b.json create mode 100644 crates/storage-pg/migrations/20241210115428_oauth_refresh_token_track_next.sql diff --git a/crates/data-model/src/tokens.rs b/crates/data-model/src/tokens.rs index b4c8ede28..210eae0cc 100644 --- a/crates/data-model/src/tokens.rs +++ b/crates/data-model/src/tokens.rs @@ -109,6 +109,7 @@ pub enum RefreshTokenState { Valid, Consumed { consumed_at: DateTime, + next_refresh_token_id: Option, }, } @@ -118,9 +119,16 @@ impl RefreshTokenState { /// # Errors /// /// Returns an error if the refresh token is already consumed. - fn consume(self, consumed_at: DateTime) -> Result { + fn consume( + self, + consumed_at: DateTime, + replaced_by: &RefreshToken, + ) -> Result { match self { - Self::Valid => Ok(Self::Consumed { consumed_at }), + Self::Valid => Ok(Self::Consumed { + consumed_at, + next_refresh_token_id: Some(replaced_by.id), + }), Self::Consumed { .. } => Err(InvalidTransitionError), } } @@ -140,6 +148,18 @@ impl RefreshTokenState { pub fn is_consumed(&self) -> bool { matches!(self, Self::Consumed { .. }) } + + /// Returns the next refresh token ID, if any. + #[must_use] + pub fn next_refresh_token_id(&self) -> Option { + match self { + Self::Valid => None, + Self::Consumed { + next_refresh_token_id, + .. + } => *next_refresh_token_id, + } + } } #[derive(Debug, Clone, PartialEq, Eq)] @@ -171,8 +191,12 @@ 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)?; + pub fn consume( + mut self, + consumed_at: DateTime, + replaced_by: &Self, + ) -> Result { + self.state = self.state.consume(consumed_at, replaced_by)?; Ok(self) } } diff --git a/crates/handlers/src/oauth2/token.rs b/crates/handlers/src/oauth2/token.rs index a228f746b..1222b32e5 100644 --- a/crates/handlers/src/oauth2/token.rs +++ b/crates/handlers/src/oauth2/token.rs @@ -544,7 +544,7 @@ 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 { diff --git a/crates/storage-pg/.sqlx/query-e709869c062ac50248b1f9f8f808cc2f5e7bef58a6c2e42a7bb0c1cb8f508671.json b/crates/storage-pg/.sqlx/query-265a981142194216593cd09cd8f6af36c103e030358262e46a4bd5e4006dc630.json similarity index 74% rename from crates/storage-pg/.sqlx/query-e709869c062ac50248b1f9f8f808cc2f5e7bef58a6c2e42a7bb0c1cb8f508671.json rename to crates/storage-pg/.sqlx/query-265a981142194216593cd09cd8f6af36c103e030358262e46a4bd5e4006dc630.json index 411b02cbe..df5d58132 100644 --- a/crates/storage-pg/.sqlx/query-e709869c062ac50248b1f9f8f808cc2f5e7bef58a6c2e42a7bb0c1cb8f508671.json +++ b/crates/storage-pg/.sqlx/query-265a981142194216593cd09cd8f6af36c103e030358262e46a4bd5e4006dc630.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 , 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": [ { @@ -32,6 +32,11 @@ "ordinal": 5, "name": "oauth2_session_id", "type_info": "Uuid" + }, + { + "ordinal": 6, + "name": "next_oauth2_refresh_token_id", + "type_info": "Uuid" } ], "parameters": { @@ -45,8 +50,9 @@ false, true, true, - false + false, + true ] }, - "hash": "e709869c062ac50248b1f9f8f808cc2f5e7bef58a6c2e42a7bb0c1cb8f508671" + "hash": "265a981142194216593cd09cd8f6af36c103e030358262e46a4bd5e4006dc630" } diff --git a/crates/storage-pg/.sqlx/query-a6fa7811d0a7c62c7cccff96dc82db5b25462fa7669fde1941ccab4712585b20.json b/crates/storage-pg/.sqlx/query-5b94c22d44692a16fa5a6edd5dac019c36bf5983182b0871de6e85036d8df466.json similarity index 74% rename from crates/storage-pg/.sqlx/query-a6fa7811d0a7c62c7cccff96dc82db5b25462fa7669fde1941ccab4712585b20.json rename to crates/storage-pg/.sqlx/query-5b94c22d44692a16fa5a6edd5dac019c36bf5983182b0871de6e85036d8df466.json index bfa6485e8..be9fa3e32 100644 --- a/crates/storage-pg/.sqlx/query-a6fa7811d0a7c62c7cccff96dc82db5b25462fa7669fde1941ccab4712585b20.json +++ b/crates/storage-pg/.sqlx/query-5b94c22d44692a16fa5a6edd5dac019c36bf5983182b0871de6e85036d8df466.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 , 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": [ { @@ -32,6 +32,11 @@ "ordinal": 5, "name": "oauth2_session_id", "type_info": "Uuid" + }, + { + "ordinal": 6, + "name": "next_oauth2_refresh_token_id", + "type_info": "Uuid" } ], "parameters": { @@ -45,8 +50,9 @@ false, true, true, - false + false, + true ] }, - "hash": "a6fa7811d0a7c62c7cccff96dc82db5b25462fa7669fde1941ccab4712585b20" + "hash": "5b94c22d44692a16fa5a6edd5dac019c36bf5983182b0871de6e85036d8df466" } diff --git a/crates/storage-pg/.sqlx/query-b6a6f5386dc89e4bc2ce56d578a29341848fce336d339b6bbf425956f5ed5032.json b/crates/storage-pg/.sqlx/query-b6a6f5386dc89e4bc2ce56d578a29341848fce336d339b6bbf425956f5ed5032.json deleted file mode 100644 index 30aa4caf3..000000000 --- a/crates/storage-pg/.sqlx/query-b6a6f5386dc89e4bc2ce56d578a29341848fce336d339b6bbf425956f5ed5032.json +++ /dev/null @@ -1,15 +0,0 @@ -{ - "db_name": "PostgreSQL", - "query": "\n UPDATE oauth2_refresh_tokens\n SET consumed_at = $2\n WHERE oauth2_refresh_token_id = $1\n ", - "describe": { - "columns": [], - "parameters": { - "Left": [ - "Uuid", - "Timestamptz" - ] - }, - "nullable": [] - }, - "hash": "b6a6f5386dc89e4bc2ce56d578a29341848fce336d339b6bbf425956f5ed5032" -} 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/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..1f0ddf0f4 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 @@ -36,23 +36,39 @@ struct OAuth2RefreshTokenLookup { consumed_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.consumed_at, value.next_oauth2_refresh_token_id) { + (None, None) => RefreshTokenState::Valid, + (Some(consumed_at), None) => RefreshTokenState::Consumed { + consumed_at, + next_refresh_token_id: None, + }, + (Some(consumed_at), Some(id)) => RefreshTokenState::Consumed { + consumed_at, + next_refresh_token_id: Some(Ulid::from(id)), + }, + (None, Some(_)) => { + 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), - } + }) } } @@ -79,18 +95,20 @@ impl OAuth2RefreshTokenRepository for PgOAuth2RefreshTokenRepository<'_> { , consumed_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( @@ -114,6 +132,7 @@ impl OAuth2RefreshTokenRepository for PgOAuth2RefreshTokenRepository<'_> { , consumed_at , oauth2_access_token_id , oauth2_session_id + , next_oauth2_refresh_token_id FROM oauth2_refresh_tokens WHERE refresh_token = $1 @@ -126,7 +145,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 +213,28 @@ 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) + .consume(consumed_at, replaced_by) .map_err(DatabaseError::to_invalid_operation) } } diff --git a/crates/storage/src/oauth2/refresh_token.rs b/crates/storage/src/oauth2/refresh_token.rs index c003e9270..5612b3647 100644 --- a/crates/storage/src/oauth2/refresh_token.rs +++ b/crates/storage/src/oauth2/refresh_token.rs @@ -80,6 +80,7 @@ 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 /// @@ -89,6 +90,7 @@ pub trait OAuth2RefreshTokenRepository: Send + Sync { &mut self, clock: &dyn Clock, refresh_token: RefreshToken, + replaced_by: &RefreshToken, ) -> Result; } @@ -113,5 +115,6 @@ repository_impl!(OAuth2RefreshTokenRepository: &mut self, clock: &dyn Clock, refresh_token: RefreshToken, + replaced_by: &RefreshToken, ) -> Result; );