diff --git a/crates/handlers/src/admin/v1/oauth2_clients/get.rs b/crates/handlers/src/admin/v1/oauth2_clients/get.rs index aa941915b..f70de8747 100644 --- a/crates/handlers/src/admin/v1/oauth2_clients/get.rs +++ b/crates/handlers/src/admin/v1/oauth2_clients/get.rs @@ -95,9 +95,9 @@ mod tests { // Add a client we know the ID of let mut repo = state.repository().await.unwrap(); - let client = repo + let (client, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut state.rng(), &state.clock, vec!["https://example.com/redirect".parse().unwrap()], diff --git a/crates/handlers/src/admin/v1/oauth2_clients/list.rs b/crates/handlers/src/admin/v1/oauth2_clients/list.rs index 5f59222ce..7003abffe 100644 --- a/crates/handlers/src/admin/v1/oauth2_clients/list.rs +++ b/crates/handlers/src/admin/v1/oauth2_clients/list.rs @@ -252,7 +252,7 @@ mod tests { // Add a dynamically-registered client repo.oauth2_client() - .add( + .add_or_reuse( &mut state.rng(), &state.clock, vec!["https://first.example.com/redirect".parse().unwrap()], @@ -278,7 +278,7 @@ mod tests { // Add another dynamically-registered client repo.oauth2_client() - .add( + .add_or_reuse( &mut state.rng(), &state.clock, vec!["https://second.example.com/redirect".parse().unwrap()], @@ -424,7 +424,7 @@ mod tests { // One authorization_code client and one client_credentials client let mut repo = state.repository().await.unwrap(); repo.oauth2_client() - .add( + .add_or_reuse( &mut state.rng(), &state.clock, vec!["https://code.example.com/redirect".parse().unwrap()], @@ -448,7 +448,7 @@ mod tests { .await .unwrap(); repo.oauth2_client() - .add( + .add_or_reuse( &mut state.rng(), &state.clock, vec![], @@ -510,9 +510,9 @@ mod tests { // A client with an active client-credentials session, and one without let mut repo = state.repository().await.unwrap(); - let with_session = repo + let (with_session, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut state.rng(), &state.clock, vec![], @@ -536,7 +536,7 @@ mod tests { .await .unwrap(); repo.oauth2_client() - .add( + .add_or_reuse( &mut state.rng(), &state.clock, vec![], diff --git a/crates/handlers/src/admin/v1/oauth2_sessions/list.rs b/crates/handlers/src/admin/v1/oauth2_sessions/list.rs index b0c2f9f3f..56300c18b 100644 --- a/crates/handlers/src/admin/v1/oauth2_sessions/list.rs +++ b/crates/handlers/src/admin/v1/oauth2_sessions/list.rs @@ -548,9 +548,9 @@ mod tests { let mut rng = state.rng(); let mut repo = state.repository().await.unwrap(); - let client = repo + let (client, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut rng, &state.clock, vec!["https://example.com/redirect".parse().unwrap()], @@ -622,9 +622,9 @@ mod tests { let mut rng = state.rng(); let mut repo = state.repository().await.unwrap(); - let client = repo + let (client, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut rng, &state.clock, vec!["https://example.com/redirect".parse().unwrap()], @@ -708,9 +708,9 @@ mod tests { // Provision two clients and a session for each, both using the // client_credentials flow so that they don't depend on a user. let mut repo = state.repository().await.unwrap(); - let client_a = repo + let (client_a, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut rng, &state.clock, vec!["https://a.example.com/redirect".parse().unwrap()], @@ -733,9 +733,9 @@ mod tests { ) .await .unwrap(); - let client_b = repo + let (client_b, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut rng, &state.clock, vec!["https://b.example.com/redirect".parse().unwrap()], diff --git a/crates/handlers/src/admin/v1/users/list.rs b/crates/handlers/src/admin/v1/users/list.rs index fde06716e..d4c44bb35 100644 --- a/crates/handlers/src/admin/v1/users/list.rs +++ b/crates/handlers/src/admin/v1/users/list.rs @@ -553,7 +553,7 @@ mod tests { rng: &mut rand_chacha::ChaChaRng, host: &str| { repo.oauth2_client() - .add( + .add_or_reuse( rng, &state.clock, vec![format!("https://{host}/redirect").parse().unwrap()], @@ -576,6 +576,7 @@ mod tests { ) .await .unwrap() + .0 }; let client_x = make_client(&mut repo, &mut rng, "x.example.com").await; @@ -724,9 +725,9 @@ mod tests { let mut repo = state.repository().await.unwrap(); - let client = repo + let (client, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut rng, &state.clock, vec!["https://example.com/redirect".parse().unwrap()], diff --git a/crates/handlers/src/cleanup_tests.rs b/crates/handlers/src/cleanup_tests.rs index a0bcdba30..035588318 100644 --- a/crates/handlers/src/cleanup_tests.rs +++ b/crates/handlers/src/cleanup_tests.rs @@ -131,9 +131,9 @@ async fn create_session_hierarchy( // Create OAuth2 session if requested let oauth2_session = if with_oauth2 { // First create an OAuth2 client - let client = repo + let (client, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut rng, &state.clock, vec!["https://example.com/callback".parse().unwrap()], diff --git a/crates/handlers/src/graphql/tests.rs b/crates/handlers/src/graphql/tests.rs index 898df9a56..fde54bd9b 100644 --- a/crates/handlers/src/graphql/tests.rs +++ b/crates/handlers/src/graphql/tests.rs @@ -28,9 +28,9 @@ async fn create_test_client(state: &TestState) -> Client { let mut repo = state.repository().await.unwrap(); let mut rng = state.rng(); - let client = repo + let (client, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut rng, &state.clock, vec![], diff --git a/crates/handlers/src/oauth2/registration.rs b/crates/handlers/src/oauth2/registration.rs index d3176adad..8439b2002 100644 --- a/crates/handlers/src/oauth2/registration.rs +++ b/crates/handlers/src/oauth2/registration.rs @@ -320,60 +320,54 @@ pub(crate) async fn post( }; // If the client doesn't have a secret, we may be able to deduplicate it. To - // do so, we hash the client metadata, and look for it in the database - let (digest_hash, existing_client) = if client_secret.is_none() { + // do so, we hash the client metadata, and reuse any client with the same hash + let digest_hash = client_secret.is_none().then(|| { // XXX: One interesting caveat is that we hash *before* saving to the database. // It means it takes into account fields that we don't care about *yet*. // // This means that if later we start supporting a particular field, we // will still serve the 'old' client_id, without updating the client in the // database - let hash = sha2::Sha256::digest(body_json); - let hash = hex::encode(hash); - let client = repo.oauth2_client().find_by_metadata_digest(&hash).await?; - (Some(hash), client) - } else { - (None, None) - }; + hex::encode(sha2::Sha256::digest(body_json)) + }); - let client = if let Some(client) = existing_client { - tracing::info!(%client.id, "Reusing existing client"); - REGISTRATION_COUNTER.add(1, &[KeyValue::new(RESULT, "reused")]); - client - } else { - let client = repo - .oauth2_client() - .add( - &mut rng, - &clock, - metadata.redirect_uris().to_vec(), - digest_hash, - encrypted_client_secret, - metadata.application_type.clone(), - //&metadata.response_types(), - metadata.grant_types().to_vec(), - metadata - .client_name - .clone() - .map(Localized::to_non_localized), - metadata.logo_uri.clone().map(Localized::to_non_localized), - metadata.client_uri.clone().map(Localized::to_non_localized), - metadata.policy_uri.clone().map(Localized::to_non_localized), - metadata.tos_uri.clone().map(Localized::to_non_localized), - metadata.jwks_uri.clone(), - metadata.jwks.clone(), - // XXX: those might not be right, should be function calls - metadata.id_token_signed_response_alg.clone(), - metadata.userinfo_signed_response_alg.clone(), - metadata.token_endpoint_auth_method.clone(), - metadata.token_endpoint_auth_signing_alg.clone(), - metadata.initiate_login_uri.clone(), - ) - .await?; + let (client, created) = repo + .oauth2_client() + .add_or_reuse( + &mut rng, + &clock, + metadata.redirect_uris().to_vec(), + digest_hash, + encrypted_client_secret, + metadata.application_type.clone(), + //&metadata.response_types(), + metadata.grant_types().to_vec(), + metadata + .client_name + .clone() + .map(Localized::to_non_localized), + metadata.logo_uri.clone().map(Localized::to_non_localized), + metadata.client_uri.clone().map(Localized::to_non_localized), + metadata.policy_uri.clone().map(Localized::to_non_localized), + metadata.tos_uri.clone().map(Localized::to_non_localized), + metadata.jwks_uri.clone(), + metadata.jwks.clone(), + // XXX: those might not be right, should be function calls + metadata.id_token_signed_response_alg.clone(), + metadata.userinfo_signed_response_alg.clone(), + metadata.token_endpoint_auth_method.clone(), + metadata.token_endpoint_auth_signing_alg.clone(), + metadata.initiate_login_uri.clone(), + ) + .await?; + + if created { tracing::info!(%client.id, "Registered new client"); REGISTRATION_COUNTER.add(1, &[KeyValue::new(RESULT, "created")]); - client - }; + } else { + tracing::info!(%client.id, "Reusing existing client"); + REGISTRATION_COUNTER.add(1, &[KeyValue::new(RESULT, "reused")]); + } let response = ClientRegistrationResponse { client_id: client.client_id.clone(), diff --git a/crates/storage-pg/.sqlx/query-17e390ee0b52d0c1c60b0951811fa6c8385fad48d470685353d074e34f5d16c9.json b/crates/storage-pg/.sqlx/query-17e390ee0b52d0c1c60b0951811fa6c8385fad48d470685353d074e34f5d16c9.json deleted file mode 100644 index 9ef809756..000000000 --- a/crates/storage-pg/.sqlx/query-17e390ee0b52d0c1c60b0951811fa6c8385fad48d470685353d074e34f5d16c9.json +++ /dev/null @@ -1,148 +0,0 @@ -{ - "db_name": "PostgreSQL", - "query": "\n SELECT oauth2_client_id\n , metadata_digest\n , encrypted_client_secret\n , application_type\n , redirect_uris\n , grant_type_authorization_code\n , grant_type_refresh_token\n , grant_type_client_credentials\n , grant_type_device_code\n , client_name\n , logo_uri\n , client_uri\n , policy_uri\n , tos_uri\n , jwks_uri\n , jwks\n , id_token_signed_response_alg\n , userinfo_signed_response_alg\n , token_endpoint_auth_method\n , token_endpoint_auth_signing_alg\n , initiate_login_uri\n , is_static\n FROM oauth2_clients\n WHERE metadata_digest = $1\n ", - "describe": { - "columns": [ - { - "ordinal": 0, - "name": "oauth2_client_id", - "type_info": "Uuid" - }, - { - "ordinal": 1, - "name": "metadata_digest", - "type_info": "Text" - }, - { - "ordinal": 2, - "name": "encrypted_client_secret", - "type_info": "Text" - }, - { - "ordinal": 3, - "name": "application_type", - "type_info": "Text" - }, - { - "ordinal": 4, - "name": "redirect_uris", - "type_info": "TextArray" - }, - { - "ordinal": 5, - "name": "grant_type_authorization_code", - "type_info": "Bool" - }, - { - "ordinal": 6, - "name": "grant_type_refresh_token", - "type_info": "Bool" - }, - { - "ordinal": 7, - "name": "grant_type_client_credentials", - "type_info": "Bool" - }, - { - "ordinal": 8, - "name": "grant_type_device_code", - "type_info": "Bool" - }, - { - "ordinal": 9, - "name": "client_name", - "type_info": "Text" - }, - { - "ordinal": 10, - "name": "logo_uri", - "type_info": "Text" - }, - { - "ordinal": 11, - "name": "client_uri", - "type_info": "Text" - }, - { - "ordinal": 12, - "name": "policy_uri", - "type_info": "Text" - }, - { - "ordinal": 13, - "name": "tos_uri", - "type_info": "Text" - }, - { - "ordinal": 14, - "name": "jwks_uri", - "type_info": "Text" - }, - { - "ordinal": 15, - "name": "jwks", - "type_info": "Jsonb" - }, - { - "ordinal": 16, - "name": "id_token_signed_response_alg", - "type_info": "Text" - }, - { - "ordinal": 17, - "name": "userinfo_signed_response_alg", - "type_info": "Text" - }, - { - "ordinal": 18, - "name": "token_endpoint_auth_method", - "type_info": "Text" - }, - { - "ordinal": 19, - "name": "token_endpoint_auth_signing_alg", - "type_info": "Text" - }, - { - "ordinal": 20, - "name": "initiate_login_uri", - "type_info": "Text" - }, - { - "ordinal": 21, - "name": "is_static", - "type_info": "Bool" - } - ], - "parameters": { - "Left": [ - "Text" - ] - }, - "nullable": [ - false, - true, - true, - true, - false, - false, - false, - false, - false, - true, - true, - true, - true, - true, - true, - true, - true, - true, - true, - true, - true, - false - ] - }, - "hash": "17e390ee0b52d0c1c60b0951811fa6c8385fad48d470685353d074e34f5d16c9" -} diff --git a/crates/storage-pg/.sqlx/query-4d0386ad2fe47f1aded46917abe6141752ba90d36467693a68318573171d57b0.json b/crates/storage-pg/.sqlx/query-4d0386ad2fe47f1aded46917abe6141752ba90d36467693a68318573171d57b0.json deleted file mode 100644 index 0b5fb6205..000000000 --- a/crates/storage-pg/.sqlx/query-4d0386ad2fe47f1aded46917abe6141752ba90d36467693a68318573171d57b0.json +++ /dev/null @@ -1,34 +0,0 @@ -{ - "db_name": "PostgreSQL", - "query": "\n INSERT INTO oauth2_clients\n ( oauth2_client_id\n , metadata_digest\n , encrypted_client_secret\n , application_type\n , redirect_uris\n , grant_type_authorization_code\n , grant_type_refresh_token\n , grant_type_client_credentials\n , grant_type_device_code\n , client_name\n , logo_uri\n , client_uri\n , policy_uri\n , tos_uri\n , jwks_uri\n , jwks\n , id_token_signed_response_alg\n , userinfo_signed_response_alg\n , token_endpoint_auth_method\n , token_endpoint_auth_signing_alg\n , initiate_login_uri\n , is_static\n )\n VALUES\n ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13,\n $14, $15, $16, $17, $18, $19, $20, $21, FALSE)\n ", - "describe": { - "columns": [], - "parameters": { - "Left": [ - "Uuid", - "Text", - "Text", - "Text", - "TextArray", - "Bool", - "Bool", - "Bool", - "Bool", - "Text", - "Text", - "Text", - "Text", - "Text", - "Text", - "Jsonb", - "Text", - "Text", - "Text", - "Text", - "Text" - ] - }, - "nullable": [] - }, - "hash": "4d0386ad2fe47f1aded46917abe6141752ba90d36467693a68318573171d57b0" -} diff --git a/crates/storage-pg/.sqlx/query-9093029989557aee00c580baa9b7795318b8c393146b039d90d296c21b417ee6.json b/crates/storage-pg/.sqlx/query-9093029989557aee00c580baa9b7795318b8c393146b039d90d296c21b417ee6.json new file mode 100644 index 000000000..ad72f64ba --- /dev/null +++ b/crates/storage-pg/.sqlx/query-9093029989557aee00c580baa9b7795318b8c393146b039d90d296c21b417ee6.json @@ -0,0 +1,168 @@ +{ + "db_name": "PostgreSQL", + "query": "\n INSERT INTO oauth2_clients\n ( oauth2_client_id\n , metadata_digest\n , encrypted_client_secret\n , application_type\n , redirect_uris\n , grant_type_authorization_code\n , grant_type_refresh_token\n , grant_type_client_credentials\n , grant_type_device_code\n , client_name\n , logo_uri\n , client_uri\n , policy_uri\n , tos_uri\n , jwks_uri\n , jwks\n , id_token_signed_response_alg\n , userinfo_signed_response_alg\n , token_endpoint_auth_method\n , token_endpoint_auth_signing_alg\n , initiate_login_uri\n , is_static\n )\n VALUES\n ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13,\n $14, $15, $16, $17, $18, $19, $20, $21, FALSE)\n -- RETURNING yields nothing for rows skipped by DO NOTHING, so this\n -- no-op update is what makes it return the existing client, including\n -- one committed while this insert was waiting on it\n ON CONFLICT (metadata_digest) DO UPDATE\n SET metadata_digest = EXCLUDED.metadata_digest\n RETURNING oauth2_client_id\n , metadata_digest\n , encrypted_client_secret\n , application_type\n , redirect_uris\n , grant_type_authorization_code\n , grant_type_refresh_token\n , grant_type_client_credentials\n , grant_type_device_code\n , client_name\n , logo_uri\n , client_uri\n , policy_uri\n , tos_uri\n , jwks_uri\n , jwks\n , id_token_signed_response_alg\n , userinfo_signed_response_alg\n , token_endpoint_auth_method\n , token_endpoint_auth_signing_alg\n , initiate_login_uri\n , is_static\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "oauth2_client_id", + "type_info": "Uuid" + }, + { + "ordinal": 1, + "name": "metadata_digest", + "type_info": "Text" + }, + { + "ordinal": 2, + "name": "encrypted_client_secret", + "type_info": "Text" + }, + { + "ordinal": 3, + "name": "application_type", + "type_info": "Text" + }, + { + "ordinal": 4, + "name": "redirect_uris", + "type_info": "TextArray" + }, + { + "ordinal": 5, + "name": "grant_type_authorization_code", + "type_info": "Bool" + }, + { + "ordinal": 6, + "name": "grant_type_refresh_token", + "type_info": "Bool" + }, + { + "ordinal": 7, + "name": "grant_type_client_credentials", + "type_info": "Bool" + }, + { + "ordinal": 8, + "name": "grant_type_device_code", + "type_info": "Bool" + }, + { + "ordinal": 9, + "name": "client_name", + "type_info": "Text" + }, + { + "ordinal": 10, + "name": "logo_uri", + "type_info": "Text" + }, + { + "ordinal": 11, + "name": "client_uri", + "type_info": "Text" + }, + { + "ordinal": 12, + "name": "policy_uri", + "type_info": "Text" + }, + { + "ordinal": 13, + "name": "tos_uri", + "type_info": "Text" + }, + { + "ordinal": 14, + "name": "jwks_uri", + "type_info": "Text" + }, + { + "ordinal": 15, + "name": "jwks", + "type_info": "Jsonb" + }, + { + "ordinal": 16, + "name": "id_token_signed_response_alg", + "type_info": "Text" + }, + { + "ordinal": 17, + "name": "userinfo_signed_response_alg", + "type_info": "Text" + }, + { + "ordinal": 18, + "name": "token_endpoint_auth_method", + "type_info": "Text" + }, + { + "ordinal": 19, + "name": "token_endpoint_auth_signing_alg", + "type_info": "Text" + }, + { + "ordinal": 20, + "name": "initiate_login_uri", + "type_info": "Text" + }, + { + "ordinal": 21, + "name": "is_static", + "type_info": "Bool" + } + ], + "parameters": { + "Left": [ + "Uuid", + "Text", + "Text", + "Text", + "TextArray", + "Bool", + "Bool", + "Bool", + "Bool", + "Text", + "Text", + "Text", + "Text", + "Text", + "Text", + "Jsonb", + "Text", + "Text", + "Text", + "Text", + "Text" + ] + }, + "nullable": [ + false, + true, + true, + true, + false, + false, + false, + false, + false, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + false + ] + }, + "hash": "9093029989557aee00c580baa9b7795318b8c393146b039d90d296c21b417ee6" +} diff --git a/crates/storage-pg/src/app_session.rs b/crates/storage-pg/src/app_session.rs index 1f1762c35..002dc2f44 100644 --- a/crates/storage-pg/src/app_session.rs +++ b/crates/storage-pg/src/app_session.rs @@ -655,9 +655,9 @@ mod tests { ); // Start an OAuth2 session - let client = repo + let (client, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut rng, &clock, vec!["https://example.com/redirect".parse().unwrap()], @@ -824,9 +824,9 @@ mod tests { clock.advance(Duration::try_minutes(1).unwrap()); // Create an OAuth2 session for the same user - let client = repo + let (client, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut rng, &clock, vec!["https://example.com/redirect".parse().unwrap()], diff --git a/crates/storage-pg/src/oauth2/client.rs b/crates/storage-pg/src/oauth2/client.rs index 7a723323a..37c8aed01 100644 --- a/crates/storage-pg/src/oauth2/client.rs +++ b/crates/storage-pg/src/oauth2/client.rs @@ -372,57 +372,6 @@ impl OAuth2ClientRepository for PgOAuth2ClientRepository<'_> { Ok(Some(res.try_into()?)) } - #[tracing::instrument( - name = "db.oauth2_client.find_by_metadata_digest", - skip_all, - fields( - db.query.text, - ), - err, - )] - async fn find_by_metadata_digest( - &mut self, - digest: &str, - ) -> Result, Self::Error> { - let res = sqlx::query_as!( - OAuth2ClientLookup, - r#" - SELECT oauth2_client_id - , metadata_digest - , encrypted_client_secret - , application_type - , redirect_uris - , grant_type_authorization_code - , grant_type_refresh_token - , grant_type_client_credentials - , grant_type_device_code - , client_name - , logo_uri - , client_uri - , policy_uri - , tos_uri - , jwks_uri - , jwks - , id_token_signed_response_alg - , userinfo_signed_response_alg - , token_endpoint_auth_method - , token_endpoint_auth_signing_alg - , initiate_login_uri - , is_static - FROM oauth2_clients - WHERE metadata_digest = $1 - "#, - digest, - ) - .traced() - .fetch_optional(&mut *self.conn) - .await?; - - let Some(res) = res else { return Ok(None) }; - - Ok(Some(res.try_into()?)) - } - #[tracing::instrument( name = "db.oauth2_client.load_batch", skip_all, @@ -481,7 +430,7 @@ impl OAuth2ClientRepository for PgOAuth2ClientRepository<'_> { } #[tracing::instrument( - name = "db.oauth2_client.add", + name = "db.oauth2_client.add_or_reuse", skip_all, fields( db.query.text, @@ -490,7 +439,7 @@ impl OAuth2ClientRepository for PgOAuth2ClientRepository<'_> { ), err, )] - async fn add( + async fn add_or_reuse( &mut self, rng: &mut (dyn Rng + Send), clock: &dyn Clock, @@ -511,10 +460,13 @@ impl OAuth2ClientRepository for PgOAuth2ClientRepository<'_> { token_endpoint_auth_method: Option, token_endpoint_auth_signing_alg: Option, initiate_login_uri: Option, - ) -> Result { + ) -> Result<(Client, bool), Self::Error> { let now = clock.now(); let id = Ulid::from_datetime_with_rng(now, rng); - tracing::Span::current().record("client.id", tracing::field::display(id)); + + if jwks.is_some() && jwks_uri.is_some() { + return Err(DatabaseError::invalid_operation()); + } let jwks_json = jwks .as_ref() @@ -524,7 +476,8 @@ impl OAuth2ClientRepository for PgOAuth2ClientRepository<'_> { let redirect_uris_array = redirect_uris.iter().map(Url::to_string).collect::>(); - sqlx::query!( + let res = sqlx::query_as!( + OAuth2ClientLookup, r#" INSERT INTO oauth2_clients ( oauth2_client_id @@ -553,6 +506,33 @@ impl OAuth2ClientRepository for PgOAuth2ClientRepository<'_> { VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, FALSE) + -- RETURNING yields nothing for rows skipped by DO NOTHING, so this + -- no-op update is what makes it return the existing client, including + -- one committed while this insert was waiting on it + ON CONFLICT (metadata_digest) DO UPDATE + SET metadata_digest = EXCLUDED.metadata_digest + RETURNING oauth2_client_id + , metadata_digest + , encrypted_client_secret + , application_type + , redirect_uris + , grant_type_authorization_code + , grant_type_refresh_token + , grant_type_client_credentials + , grant_type_device_code + , client_name + , logo_uri + , client_uri + , policy_uri + , tos_uri + , jwks_uri + , jwks + , id_token_signed_response_alg + , userinfo_signed_response_alg + , token_endpoint_auth_method + , token_endpoint_auth_signing_alg + , initiate_login_uri + , is_static "#, Uuid::from(id), metadata_digest, @@ -583,37 +563,14 @@ impl OAuth2ClientRepository for PgOAuth2ClientRepository<'_> { initiate_login_uri.as_ref().map(Url::as_str), ) .traced() - .execute(&mut *self.conn) + .fetch_one(&mut *self.conn) .await?; - let jwks = match (jwks, jwks_uri) { - (None, None) => None, - (Some(jwks), None) => Some(JwksOrJwksUri::Jwks(jwks)), - (None, Some(jwks_uri)) => Some(JwksOrJwksUri::JwksUri(jwks_uri)), - _ => return Err(DatabaseError::invalid_operation()), - }; + let client: Client = res.try_into()?; + tracing::Span::current().record("client.id", tracing::field::display(client.id)); + let inserted = client.id == id; - Ok(Client { - id, - client_id: id.to_string(), - metadata_digest: None, - encrypted_client_secret, - application_type, - redirect_uris, - grant_types, - client_name, - logo_uri, - client_uri, - policy_uri, - tos_uri, - jwks, - id_token_signed_response_alg, - userinfo_signed_response_alg, - token_endpoint_auth_method, - token_endpoint_auth_signing_alg, - initiate_login_uri, - is_static: false, - }) + Ok((client, inserted)) } #[tracing::instrument( diff --git a/crates/storage-pg/src/oauth2/mod.rs b/crates/storage-pg/src/oauth2/mod.rs index 0bdaf3f4a..b9d14d205 100644 --- a/crates/storage-pg/src/oauth2/mod.rs +++ b/crates/storage-pg/src/oauth2/mod.rs @@ -25,10 +25,10 @@ pub use self::{ #[cfg(test)] mod tests { use chrono::Duration; - use mas_data_model::{AuthorizationCode, Clock, UlidExt as _, clock::MockClock}; + use mas_data_model::{AuthorizationCode, Client, Clock, UlidExt as _, clock::MockClock}; use mas_iana::oauth::OAuthClientAuthenticationMethod; use mas_storage::{ - Pagination, + BoxRepository, Pagination, oauth2::{ OAuth2ClientFilter, OAuth2DeviceCodeGrantParams, OAuth2SessionFilter, OAuth2SessionRepository, @@ -64,9 +64,9 @@ mod tests { assert_eq!(client, None); // Create a client - let client = repo + let (client, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut rng, &clock, vec!["https://example.com/redirect".parse().unwrap()], @@ -414,9 +414,9 @@ mod tests { .unwrap(); // Create two clients - let client1 = repo + let (client1, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut rng, &clock, vec!["https://first.example.com/redirect".parse().unwrap()], @@ -439,9 +439,9 @@ mod tests { ) .await .unwrap(); - let client2 = repo + let (client2, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut rng, &clock, vec!["https://second.example.com/redirect".parse().unwrap()], @@ -748,9 +748,9 @@ mod tests { .add(&mut rng, &clock, &user, None) .await .unwrap(); - let client = repo + let (client, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut rng, &clock, vec!["https://example.com/redirect".parse().unwrap()], @@ -849,9 +849,9 @@ mod tests { // Provision three clients let mut clients = Vec::new(); for label in ["first", "second", "third"] { - let client = repo + let (client, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut rng, &clock, vec![ @@ -969,9 +969,9 @@ mod tests { let mut repo = PgRepository::from_pool(&pool).await.unwrap().boxed(); // Provision a client - let client = repo + let (client, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut rng, &clock, vec!["https://example.com/redirect".parse().unwrap()], @@ -1211,9 +1211,9 @@ mod tests { assert!(!page.has_next_page); // Add a couple of clients - let client1 = repo + let (client1, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut rng, &clock, vec!["https://first.example.com/redirect".parse().unwrap()], @@ -1238,9 +1238,9 @@ mod tests { .unwrap(); clock.advance(Duration::try_minutes(1).unwrap()); - let client2 = repo + let (client2, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut rng, &clock, vec!["https://second.example.com/redirect".parse().unwrap()], @@ -1363,9 +1363,9 @@ mod tests { let mut repo = PgRepository::from_pool(&pool).await.unwrap().boxed(); // A client supporting authorization_code (+ refresh_token) - let auth_code_client = repo + let (auth_code_client, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut rng, &clock, vec!["https://code.example.com/redirect".parse().unwrap()], @@ -1390,9 +1390,9 @@ mod tests { .unwrap(); // A client supporting only client_credentials - let client_credentials_client = repo + let (client_credentials_client, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut rng, &clock, vec![], @@ -1459,9 +1459,9 @@ mod tests { let mut repo = PgRepository::from_pool(&pool).await.unwrap().boxed(); // A client that will have an active session - let with_session = repo + let (with_session, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut rng, &clock, vec![], @@ -1486,9 +1486,9 @@ mod tests { .unwrap(); // A client without any session - let without_session = repo + let (without_session, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut rng, &clock, vec![], @@ -1553,4 +1553,72 @@ mod tests { let filter = OAuth2ClientFilter::new().with_active_sessions(false); assert_eq!(repo.oauth2_client().count(filter).await.unwrap(), 2); } + + #[sqlx::test(migrator = "crate::MIGRATOR")] + async fn test_add_client_with_concurrent_metadata_digest(pool: PgPool) { + async fn add( + repo: &mut BoxRepository, + rng: &mut ChaChaRng, + clock: &MockClock, + ) -> (Client, bool) { + repo.oauth2_client() + .add_or_reuse( + rng, + clock, + vec!["https://example.com/redirect".parse().unwrap()], + Some("digest".to_owned()), + None, + None, + vec![GrantType::AuthorizationCode], + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + ) + .await + .unwrap() + } + + let clock = MockClock::default(); + let mut repo1 = PgRepository::from_pool(&pool).await.unwrap().boxed(); + let mut repo2 = PgRepository::from_pool(&pool).await.unwrap().boxed(); + + let (client1, created) = add(&mut repo1, &mut ChaChaRng::seed_from_u64(1), &clock).await; + assert!(created); + assert_eq!(client1.metadata_digest.as_deref(), Some("digest")); + + // The second insert either waits for the first transaction or sees its + // committed row + let second = tokio::spawn(async move { + let client = add( + &mut repo2, + &mut ChaChaRng::seed_from_u64(2), + &MockClock::default(), + ) + .await; + repo2.save().await.unwrap(); + client + }); + repo1.save().await.unwrap(); + let (client2, created) = second.await.unwrap(); + + assert!(!created); + assert_eq!(client1.id, client2.id); + let mut repo = PgRepository::from_pool(&pool).await.unwrap().boxed(); + assert_eq!( + repo.oauth2_client() + .count(OAuth2ClientFilter::new()) + .await + .unwrap(), + 1 + ); + } } diff --git a/crates/storage-pg/src/user/tests.rs b/crates/storage-pg/src/user/tests.rs index e46941245..cc5da4fd4 100644 --- a/crates/storage-pg/src/user/tests.rs +++ b/crates/storage-pg/src/user/tests.rs @@ -891,7 +891,7 @@ async fn test_user_filter_active_oauth2_session_for_any_of_clients(pool: PgPool) clock: &MockClock, host: &str| { repo.oauth2_client() - .add( + .add_or_reuse( rng, clock, vec![format!("https://{host}/redirect").parse().unwrap()], @@ -914,6 +914,7 @@ async fn test_user_filter_active_oauth2_session_for_any_of_clients(pool: PgPool) ) .await .unwrap() + .0 }; let client_x = make_client(&mut repo, &mut rng, &clock, "x.example.com").await; @@ -1101,9 +1102,9 @@ async fn test_user_filter_active_oauth2_session(pool: PgPool) { .await .unwrap(); - let client = repo + let (client, _) = repo .oauth2_client() - .add( + .add_or_reuse( &mut rng, &clock, vec!["https://example.com/redirect".parse().unwrap()], diff --git a/crates/storage/src/oauth2/client.rs b/crates/storage/src/oauth2/client.rs index 8f37877e3..0c2854b4e 100644 --- a/crates/storage/src/oauth2/client.rs +++ b/crates/storage/src/oauth2/client.rs @@ -161,23 +161,6 @@ pub trait OAuth2ClientRepository: Send + Sync { self.lookup(id).await } - /// Find an OAuth client by its metadata digest - /// - /// Returns `None` if the client does not exist - /// - /// # Parameters - /// - /// * `digest`: The metadata digest (SHA-256 hash encoded in hex) of the - /// client to find - /// - /// # Errors - /// - /// Returns [`Self::Error`] if the underlying repository fails - async fn find_by_metadata_digest( - &mut self, - digest: &str, - ) -> Result, Self::Error>; - /// Load a batch of OAuth clients by their IDs /// /// Returns a map of client IDs to clients. If a client does not exist, it @@ -195,9 +178,12 @@ pub trait OAuth2ClientRepository: Send + Sync { ids: BTreeSet, ) -> Result, Self::Error>; - /// Add a new OAuth client + /// Add a new OAuth client, or reuse the existing one with the same metadata + /// digest /// - /// Returns the client that was added + /// Returns the client and whether it was inserted. If `metadata_digest` is + /// set and a client with the same digest already exists, that client is + /// returned unchanged instead. /// /// # Parameters /// @@ -230,7 +216,7 @@ pub trait OAuth2ClientRepository: Send + Sync { /// /// Returns [`Self::Error`] if the underlying repository fails #[expect(clippy::too_many_arguments)] - async fn add( + async fn add_or_reuse( &mut self, rng: &mut (dyn Rng + Send), clock: &dyn Clock, @@ -251,7 +237,7 @@ pub trait OAuth2ClientRepository: Send + Sync { token_endpoint_auth_method: Option, token_endpoint_auth_signing_alg: Option, initiate_login_uri: Option, - ) -> Result; + ) -> Result<(Client, bool), Self::Error>; /// Add or replace a static client /// @@ -348,17 +334,12 @@ pub trait OAuth2ClientRepository: Send + Sync { repository_impl!(OAuth2ClientRepository: async fn lookup(&mut self, id: Ulid) -> Result, Self::Error>; - async fn find_by_metadata_digest( - &mut self, - digest: &str, - ) -> Result, Self::Error>; - async fn load_batch( &mut self, ids: BTreeSet, ) -> Result, Self::Error>; - async fn add( + async fn add_or_reuse( &mut self, rng: &mut (dyn Rng + Send), clock: &dyn Clock, @@ -379,7 +360,7 @@ repository_impl!(OAuth2ClientRepository: token_endpoint_auth_method: Option, token_endpoint_auth_signing_alg: Option, initiate_login_uri: Option, - ) -> Result; + ) -> Result<(Client, bool), Self::Error>; async fn upsert_static( &mut self,