From 03b4eba347ada3b1f5986d5d4721e935fc26ea9d Mon Sep 17 00:00:00 2001 From: Quentin Gliech Date: Fri, 9 Oct 2026 14:05:45 +0200 Subject: [PATCH] Handle concurrent registrations of identical public clients Two registrations with the same metadata could both miss the `find_by_metadata_digest` lookup and then race on the insert, the loser failing with a unique violation on `oauth2_clients_metadata_digest_key` and getting a 500. `OAuth2ClientRepository::add` becomes `add_or_reuse` and inserts with `ON CONFLICT (metadata_digest) DO UPDATE ... RETURNING`, which yields the existing client when the digest is taken, including one committed while the insert was waiting on it. It reports whether it inserted, so the registration handler drops its lookup and still tells created and reused clients apart in the metric. `find_by_metadata_digest` has no callers left and is removed. The returned client is now built from the database row, so it carries the metadata digest instead of always `None`, and the span records the id of the client actually returned. --- .../src/admin/v1/oauth2_clients/get.rs | 4 +- .../src/admin/v1/oauth2_clients/list.rs | 14 +- .../src/admin/v1/oauth2_sessions/list.rs | 16 +- crates/handlers/src/admin/v1/users/list.rs | 7 +- crates/handlers/src/cleanup_tests.rs | 4 +- crates/handlers/src/graphql/tests.rs | 4 +- crates/handlers/src/oauth2/registration.rs | 84 ++++----- ...fa6c8385fad48d470685353d074e34f5d16c9.json | 148 --------------- ...6141752ba90d36467693a68318573171d57b0.json | 34 ---- ...7795318b8c393146b039d90d296c21b417ee6.json | 168 ++++++++++++++++++ crates/storage-pg/src/app_session.rs | 8 +- crates/storage-pg/src/oauth2/client.rs | 125 +++++-------- crates/storage-pg/src/oauth2/mod.rs | 120 ++++++++++--- crates/storage-pg/src/user/tests.rs | 7 +- crates/storage/src/oauth2/client.rs | 37 +--- 15 files changed, 384 insertions(+), 396 deletions(-) delete mode 100644 crates/storage-pg/.sqlx/query-17e390ee0b52d0c1c60b0951811fa6c8385fad48d470685353d074e34f5d16c9.json delete mode 100644 crates/storage-pg/.sqlx/query-4d0386ad2fe47f1aded46917abe6141752ba90d36467693a68318573171d57b0.json create mode 100644 crates/storage-pg/.sqlx/query-9093029989557aee00c580baa9b7795318b8c393146b039d90d296c21b417ee6.json 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 e47546479..e372f65a0 100644 --- a/crates/handlers/src/oauth2/registration.rs +++ b/crates/handlers/src/oauth2/registration.rs @@ -319,60 +319,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 7cc49367e..060b2c2d2 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 RngCore + 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 b1e6017d8..414f46710 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 RngCore + 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 RngCore + 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,