diff --git a/.sqlx/query-1ee6eda3e44660e7f14fcfe56adc2d41c72901b9c701fc7b992314e5370b32dc.json b/.sqlx/query-0bb2cb6af37bff735b6b380697fc8e1fa2034ca0600e8c16e1e362b722192327.json similarity index 57% rename from .sqlx/query-1ee6eda3e44660e7f14fcfe56adc2d41c72901b9c701fc7b992314e5370b32dc.json rename to .sqlx/query-0bb2cb6af37bff735b6b380697fc8e1fa2034ca0600e8c16e1e362b722192327.json index a302cfc..ac0e971 100644 --- a/.sqlx/query-1ee6eda3e44660e7f14fcfe56adc2d41c72901b9c701fc7b992314e5370b32dc.json +++ b/.sqlx/query-0bb2cb6af37bff735b6b380697fc8e1fa2034ca0600e8c16e1e362b722192327.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "UPDATE invite_codes SET available_uses = available_uses - 1 WHERE code = $1", + "query": "UPDATE invite_codes SET available_uses = available_uses - 1 WHERE code = $1 AND available_uses > 0 AND COALESCE(disabled, false) = false", "describe": { "columns": [], "parameters": { @@ -10,5 +10,5 @@ }, "nullable": [] }, - "hash": "1ee6eda3e44660e7f14fcfe56adc2d41c72901b9c701fc7b992314e5370b32dc" + "hash": "0bb2cb6af37bff735b6b380697fc8e1fa2034ca0600e8c16e1e362b722192327" } diff --git a/crates/tranquil-api/src/delegation.rs b/crates/tranquil-api/src/delegation.rs index b2a30dc..f6bfcb0 100644 --- a/crates/tranquil-api/src/delegation.rs +++ b/crates/tranquil-api/src/delegation.rs @@ -330,7 +330,6 @@ pub struct CreateDelegatedAccountInput { pub handle: String, pub email: Option, pub controller_scopes: ValidatedDelegationScope, - pub invite_code: Option, } #[derive(Debug, Serialize)] @@ -362,19 +361,6 @@ pub async fn create_delegated_account( return Err(ApiError::InvalidEmail); } - let validated_invite_code = if let Some(ref code) = input.invite_code { - match state.repos.infra.validate_invite_code(code).await { - Ok(validated) => Some(validated), - Err(_) => return Err(ApiError::InvalidInviteCode), - } - } else { - let invite_required = tranquil_config::get().server.invite_code_required; - if invite_required { - return Err(ApiError::InviteCodeRequired); - } - None - }; - let plc = create_plc_did(&state, &handle).await.map_err(|e| { tracing::error!("PLC DID creation failed: {:?}", e); e @@ -397,16 +383,15 @@ pub async fn create_delegated_account( commit_cid: repo.commit_cid.to_string(), repo_rev: repo.repo_rev.clone(), genesis_block_cids: repo.genesis_block_cids, - invite_code: input.invite_code.clone(), }; - let user_id = match state + match state .repos .user .create_delegated_account(&create_input) .await { - Ok(id) => id, + Ok(_) => {} Err(tranquil_db_traits::CreateAccountError::HandleTaken) => { return Err(ApiError::HandleNotAvailable(None)); } @@ -417,16 +402,6 @@ pub async fn create_delegated_account( error!("Error creating delegated account: {:?}", e); return Err(ApiError::InternalError(None)); } - }; - - if let Some(validated) = validated_invite_code - && let Err(e) = state - .repos - .infra - .record_invite_code_use(&validated, user_id) - .await - { - warn!("Failed to record invite code use for {}: {:?}", did, e); } crate::identity::provision::sequence_new_account( diff --git a/crates/tranquil-api/src/identity/account.rs b/crates/tranquil-api/src/identity/account.rs index 463157b..9c520c4 100644 --- a/crates/tranquil-api/src/identity/account.rs +++ b/crates/tranquil-api/src/identity/account.rs @@ -10,6 +10,7 @@ use serde::{Deserialize, Serialize}; use serde_json::json; use tracing::{debug, error, info}; use tranquil_pds::api::error::ApiError; +use tranquil_pds::api::invite::check_registration_invite; use tranquil_pds::auth::{ServiceTokenVerifier, extract_auth_token_from_header, is_service_token}; use tranquil_pds::rate_limit::{AccountCreationLimit, RateLimited}; use tranquil_pds::state::AppState; @@ -415,40 +416,11 @@ pub async fn create_account( return ApiError::HandleTaken.into_response(); } - let is_bootstrap = state.bootstrap_invite_code.is_some() - && state.repos.user.count_users().await.unwrap_or(1) == 0; - - if is_bootstrap { - match input.invite_code.as_deref() { - Some(code) if Some(code) == state.bootstrap_invite_code.as_deref() => {} - _ => return ApiError::InvalidInviteCode.into_response(), - } - } else { - let invite_code_required = tranquil_config::get().server.invite_code_required; - if invite_code_required - && input - .invite_code - .as_ref() - .map(|c| c.trim().is_empty()) - .unwrap_or(true) - { - return ApiError::InviteCodeRequired.into_response(); - } - if let Some(code) = &input.invite_code - && !code.trim().is_empty() - { - let valid = match state.repos.user.check_and_consume_invite_code(code).await { - Ok(v) => v, - Err(e) => { - error!("Error checking invite code: {:?}", e); - return ApiError::InternalError(None).into_response(); - } - }; - if !valid { - return ApiError::InvalidInviteCode.into_response(); - } - } - } + let invite_registration = + match check_registration_invite(&state, input.invite_code.as_deref()).await { + Ok(outcome) => outcome, + Err(e) => return e.into_response(), + }; if let Err(e) = validate_password(&input.password) { return ApiError::InvalidRequest(e.to_string()).into_response(); @@ -517,11 +489,7 @@ pub async fn create_account( commit_cid: commit_cid_str.clone(), repo_rev: rev_str.clone(), genesis_block_cids: repo.genesis_block_cids, - invite_code: if is_bootstrap { - None - } else { - input.invite_code.clone() - }, + invite_code: invite_registration.into_invite_code(), birthdate_pref, }; @@ -541,6 +509,9 @@ pub async fn create_account( Err(tranquil_db_traits::CreateAccountError::DidExists) => { return ApiError::AccountAlreadyExists.into_response(); } + Err(tranquil_db_traits::CreateAccountError::InviteCodeUnavailable) => { + return ApiError::InvalidInviteCode.into_response(); + } Err(e) => { error!("Error creating password account: {:?}", e); return ApiError::InternalError(None).into_response(); diff --git a/crates/tranquil-api/src/server/passkey_account.rs b/crates/tranquil-api/src/server/passkey_account.rs index b678e43..d54a7a0 100644 --- a/crates/tranquil-api/src/server/passkey_account.rs +++ b/crates/tranquil-api/src/server/passkey_account.rs @@ -7,6 +7,7 @@ use serde_json::json; use tracing::{debug, error, info, warn}; use tranquil_db_traits::WebauthnChallengeType; use tranquil_pds::api::error::ApiError; +use tranquil_pds::api::invite::check_registration_invite; use tranquil_pds::api::{OptionsResponse, SuccessResponse}; use tranquil_pds::auth::NormalizedLoginIdentifier; @@ -119,26 +120,8 @@ pub async fn create_passkey_account( return Err(ApiError::InvalidEmail); } - let is_bootstrap = state.bootstrap_invite_code.is_some() - && state.repos.user.count_users().await.unwrap_or(1) == 0; - - let _validated_invite_code = if is_bootstrap { - match input.invite_code.as_deref() { - Some(code) if Some(code) == state.bootstrap_invite_code.as_deref() => None, - _ => return Err(ApiError::InvalidInviteCode), - } - } else if let Some(ref code) = input.invite_code { - match state.repos.infra.validate_invite_code(code).await { - Ok(validated) => Some(validated), - Err(_) => return Err(ApiError::InvalidInviteCode), - } - } else { - let invite_required = tranquil_config::get().server.invite_code_required; - if invite_required { - return Err(ApiError::InviteCodeRequired); - } - None - }; + let invite_registration = + check_registration_invite(&state, input.invite_code.as_deref()).await?; let verification_channel = input .verification_channel @@ -343,11 +326,7 @@ pub async fn create_passkey_account( commit_cid: repo.commit_cid.to_string(), repo_rev: repo.repo_rev.clone(), genesis_block_cids: repo.genesis_block_cids, - invite_code: if is_bootstrap { - None - } else { - input.invite_code.clone() - }, + invite_code: invite_registration.into_invite_code(), birthdate_pref, }; @@ -359,6 +338,9 @@ pub async fn create_passkey_account( Err(tranquil_db_traits::CreateAccountError::EmailTaken) => { return Err(ApiError::EmailTaken); } + Err(tranquil_db_traits::CreateAccountError::InviteCodeUnavailable) => { + return Err(ApiError::InvalidInviteCode); + } Err(e) => { error!("Error creating passkey account: {:?}", e); return Err(ApiError::InternalError(None)); diff --git a/crates/tranquil-db-traits/src/infra.rs b/crates/tranquil-db-traits/src/infra.rs index eb337df..566e341 100644 --- a/crates/tranquil-db-traits/src/infra.rs +++ b/crates/tranquil-db-traits/src/infra.rs @@ -268,17 +268,6 @@ pub trait InfraRepository: Send + Sync { code: &'a str, ) -> Result, InviteCodeError>; - async fn decrement_invite_code_uses( - &self, - code: &ValidatedInviteCode<'_>, - ) -> Result<(), DbError>; - - async fn record_invite_code_use( - &self, - code: &ValidatedInviteCode<'_>, - used_by_user: Uuid, - ) -> Result<(), DbError>; - async fn get_invite_codes_for_account( &self, for_account: &Did, diff --git a/crates/tranquil-db-traits/src/user.rs b/crates/tranquil-db-traits/src/user.rs index 17ff302..b6bfa86 100644 --- a/crates/tranquil-db-traits/src/user.rs +++ b/crates/tranquil-db-traits/src/user.rs @@ -589,8 +589,6 @@ pub trait UserRepository: Send + Sync { async fn cleanup_expired_handle_reservations(&self) -> Result; - async fn check_and_consume_invite_code(&self, code: &str) -> Result; - async fn complete_passkey_setup( &self, input: &CompletePasskeySetupInput, @@ -1019,6 +1017,7 @@ pub enum CreateAccountError { EmailTaken, DidExists, InvalidToken, + InviteCodeUnavailable, Database(String), } @@ -1034,7 +1033,6 @@ pub struct CreateDelegatedAccountInput { pub commit_cid: String, pub repo_rev: String, pub genesis_block_cids: Vec>, - pub invite_code: Option, } #[derive(Debug, Clone)] diff --git a/crates/tranquil-db/src/postgres/infra.rs b/crates/tranquil-db/src/postgres/infra.rs index adac68d..a875de5 100644 --- a/crates/tranquil-db/src/postgres/infra.rs +++ b/crates/tranquil-db/src/postgres/infra.rs @@ -225,38 +225,6 @@ impl InfraRepository for PostgresInfraRepository { } } - async fn decrement_invite_code_uses( - &self, - code: &ValidatedInviteCode<'_>, - ) -> Result<(), DbError> { - sqlx::query!( - "UPDATE invite_codes SET available_uses = available_uses - 1 WHERE code = $1", - code.code() - ) - .execute(&self.pool) - .await - .map_err(map_sqlx_error)?; - - Ok(()) - } - - async fn record_invite_code_use( - &self, - code: &ValidatedInviteCode<'_>, - used_by_user: Uuid, - ) -> Result<(), DbError> { - sqlx::query!( - "INSERT INTO invite_code_uses (code, used_by_user) VALUES ($1, $2)", - code.code(), - used_by_user - ) - .execute(&self.pool) - .await - .map_err(map_sqlx_error)?; - - Ok(()) - } - async fn get_invite_codes_for_account( &self, for_account: &Did, diff --git a/crates/tranquil-db/src/postgres/user.rs b/crates/tranquil-db/src/postgres/user.rs index 63131ee..4a574bb 100644 --- a/crates/tranquil-db/src/postgres/user.rs +++ b/crates/tranquil-db/src/postgres/user.rs @@ -43,6 +43,38 @@ pub(crate) fn map_sqlx_error(e: sqlx::Error) -> DbError { } } +async fn consume_invite_code( + conn: &mut sqlx::PgConnection, + code: &str, + user_id: Uuid, +) -> Result<(), tranquil_db_traits::CreateAccountError> { + let map_err = |e: sqlx::Error| tranquil_db_traits::CreateAccountError::Database(e.to_string()); + + let decremented = sqlx::query!( + "UPDATE invite_codes SET available_uses = available_uses - 1 WHERE code = $1 AND available_uses > 0 AND COALESCE(disabled, false) = false", + code + ) + .execute(&mut *conn) + .await + .map_err(map_err)? + .rows_affected(); + + if decremented == 0 { + return Err(tranquil_db_traits::CreateAccountError::InviteCodeUnavailable); + } + + sqlx::query!( + "INSERT INTO invite_code_uses (code, used_by_user) VALUES ($1, $2)", + code, + user_id + ) + .execute(&mut *conn) + .await + .map_err(map_err)?; + + Ok(()) +} + #[async_trait] impl UserRepository for PostgresUserRepository { async fn get_by_did(&self, did: &Did) -> Result, DbError> { @@ -2521,20 +2553,7 @@ impl UserRepository for PostgresUserRepository { })?; if let Some(code) = &input.invite_code { - let _ = sqlx::query!( - "UPDATE invite_codes SET available_uses = available_uses - 1 WHERE code = $1", - code - ) - .execute(&mut *tx) - .await; - - let _ = sqlx::query!( - "INSERT INTO invite_code_uses (code, used_by_user) VALUES ($1, $2)", - code, - user_id - ) - .execute(&mut *tx) - .await; + consume_invite_code(&mut tx, code, user_id).await?; } if let Some(birthdate_pref) = &input.birthdate_pref { @@ -2650,23 +2669,6 @@ impl UserRepository for PostgresUserRepository { tranquil_db_traits::CreateAccountError::Database(e.to_string()) })?; - if let Some(code) = &input.invite_code { - let _ = sqlx::query!( - "UPDATE invite_codes SET available_uses = available_uses - 1 WHERE code = $1", - code - ) - .execute(&mut *tx) - .await; - - let _ = sqlx::query!( - "INSERT INTO invite_code_uses (code, used_by_user) VALUES ($1, $2)", - code, - user_id - ) - .execute(&mut *tx) - .await; - } - tx.commit().await.map_err(|e: sqlx::Error| { tranquil_db_traits::CreateAccountError::Database(e.to_string()) })?; @@ -2784,20 +2786,7 @@ impl UserRepository for PostgresUserRepository { })?; if let Some(code) = &input.invite_code { - let _ = sqlx::query!( - "UPDATE invite_codes SET available_uses = available_uses - 1 WHERE code = $1", - code - ) - .execute(&mut *tx) - .await; - - let _ = sqlx::query!( - "INSERT INTO invite_code_uses (code, used_by_user) VALUES ($1, $2)", - code, - user_id - ) - .execute(&mut *tx) - .await; + consume_invite_code(&mut tx, code, user_id).await?; } if let Some(birthdate_pref) = &input.birthdate_pref { @@ -2932,20 +2921,7 @@ impl UserRepository for PostgresUserRepository { })?; if let Some(code) = &input.invite_code { - let _ = sqlx::query!( - "UPDATE invite_codes SET available_uses = available_uses - 1 WHERE code = $1", - code - ) - .execute(&mut *tx) - .await; - - let _ = sqlx::query!( - "INSERT INTO invite_code_uses (code, used_by_user) VALUES ($1, $2)", - code, - user_id - ) - .execute(&mut *tx) - .await; + consume_invite_code(&mut tx, code, user_id).await?; } if let Some(birthdate_pref) = &input.birthdate_pref { @@ -3124,38 +3100,6 @@ impl UserRepository for PostgresUserRepository { Ok(result.rows_affected()) } - async fn check_and_consume_invite_code(&self, code: &str) -> Result { - let mut tx = self.pool.begin().await.map_err(map_sqlx_error)?; - - let invite = sqlx::query!( - "SELECT available_uses FROM invite_codes WHERE code = $1 FOR UPDATE", - code - ) - .fetch_optional(&mut *tx) - .await - .map_err(map_sqlx_error)?; - - let Some(row) = invite else { - return Ok(false); - }; - - if row.available_uses <= 0 { - return Ok(false); - } - - sqlx::query!( - "UPDATE invite_codes SET available_uses = available_uses - 1 WHERE code = $1", - code - ) - .execute(&mut *tx) - .await - .map_err(map_sqlx_error)?; - - tx.commit().await.map_err(map_sqlx_error)?; - - Ok(true) - } - async fn complete_passkey_setup( &self, input: &tranquil_db_traits::CompletePasskeySetupInput, diff --git a/crates/tranquil-oauth-server/src/sso_endpoints.rs b/crates/tranquil-oauth-server/src/sso_endpoints.rs index 4b1b354..7c1e652 100644 --- a/crates/tranquil-oauth-server/src/sso_endpoints.rs +++ b/crates/tranquil-oauth-server/src/sso_endpoints.rs @@ -10,6 +10,7 @@ use tranquil_db_traits::{SsoAction, SsoProviderType}; use tranquil_types::RequestId; use tranquil_pds::api::error::ApiError; +use tranquil_pds::api::invite::check_registration_invite; use tranquil_pds::auth::extractor::extract_auth_token_from_header; use tranquil_pds::auth::{generate_app_password, validate_bearer_token_cached}; use tranquil_pds::rate_limit::{ @@ -985,18 +986,8 @@ pub async fn complete_registration( None => None, }; - let _validated_invite_code = if let Some(ref code) = input.invite_code { - match state.repos.infra.validate_invite_code(code).await { - Ok(validated) => Some(validated), - Err(_) => return Err(ApiError::InvalidInviteCode), - } - } else { - let invite_required = tranquil_config::get().server.invite_code_required; - if invite_required { - return Err(ApiError::InviteCodeRequired); - } - None - }; + let invite_registration = + check_registration_invite(&state, input.invite_code.as_deref()).await?; let handle_typed: tranquil_pds::types::Handle = handle.parse().map_err(|_| ApiError::InvalidHandle(None))?; @@ -1169,7 +1160,7 @@ pub async fn complete_registration( commit_cid: commit_cid.to_string(), repo_rev: rev.as_ref().to_string(), genesis_block_cids, - invite_code: input.invite_code.clone(), + invite_code: invite_registration.into_invite_code(), birthdate_pref, sso_provider: pending_preview.provider, sso_provider_user_id: pending_preview.provider_user_id.clone().into_inner(), @@ -1196,6 +1187,9 @@ pub async fn complete_registration( Err(tranquil_db_traits::CreateAccountError::InvalidToken) => { return Err(ApiError::SsoSessionExpired); } + Err(tranquil_db_traits::CreateAccountError::InviteCodeUnavailable) => { + return Err(ApiError::InvalidInviteCode); + } Err(e) => { tracing::error!("Error creating SSO account: {:?}", e); return Err(ApiError::InternalError(None)); diff --git a/crates/tranquil-pds/src/api/invite.rs b/crates/tranquil-pds/src/api/invite.rs new file mode 100644 index 0000000..1716a36 --- /dev/null +++ b/crates/tranquil-pds/src/api/invite.rs @@ -0,0 +1,51 @@ +use tranquil_db_traits::InviteCodeError; + +use crate::api::error::ApiError; +use crate::state::AppState; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum InviteRegistration { + Bootstrap, + Standard(Option), +} + +impl InviteRegistration { + pub fn into_invite_code(self) -> Option { + match self { + InviteRegistration::Bootstrap => None, + InviteRegistration::Standard(code) => code, + } + } +} + +pub async fn check_registration_invite( + state: &AppState, + invite_code: Option<&str>, +) -> Result { + let is_bootstrap = state.bootstrap_invite_code.is_some() + && state.repos.user.count_users().await.unwrap_or(1) == 0; + + if is_bootstrap { + return match invite_code { + Some(code) if Some(code) == state.bootstrap_invite_code.as_deref() => { + Ok(InviteRegistration::Bootstrap) + } + _ => Err(ApiError::InvalidInviteCode), + }; + } + + match invite_code.map(str::trim).filter(|code| !code.is_empty()) { + Some(code) => match state.repos.infra.validate_invite_code(code).await { + Ok(_) => Ok(InviteRegistration::Standard(Some(code.to_owned()))), + Err(InviteCodeError::DatabaseError(e)) => { + tracing::error!("failed to validate invite code: {e:?}"); + Err(ApiError::InternalError(None)) + } + Err(_) => Err(ApiError::InvalidInviteCode), + }, + None => match tranquil_config::get().server.invite_code_required { + true => Err(ApiError::InviteCodeRequired), + false => Ok(InviteRegistration::Standard(None)), + }, + } +} diff --git a/crates/tranquil-pds/src/api/mod.rs b/crates/tranquil-pds/src/api/mod.rs index 53de750..4e17789 100644 --- a/crates/tranquil-pds/src/api/mod.rs +++ b/crates/tranquil-pds/src/api/mod.rs @@ -1,4 +1,5 @@ pub mod error; +pub mod invite; pub mod proxy; pub mod proxy_client; pub mod responses; diff --git a/crates/tranquil-pds/tests/invite_registration.rs b/crates/tranquil-pds/tests/invite_registration.rs new file mode 100644 index 0000000..c825486 --- /dev/null +++ b/crates/tranquil-pds/tests/invite_registration.rs @@ -0,0 +1,121 @@ +mod common; +use common::*; +use reqwest::{Client, StatusCode}; +use serde_json::{Value, json}; +use tranquil_pds::api::error::ApiError; +use tranquil_pds::api::invite::{InviteRegistration, check_registration_invite}; + +async fn create_invite_code(client: &Client, admin_jwt: &str, use_count: u32) -> String { + let res = client + .post(format!( + "{}/xrpc/com.atproto.server.createInviteCode", + base_url().await + )) + .bearer_auth(admin_jwt) + .json(&json!({ "useCount": use_count })) + .send() + .await + .expect("failed to create invite code"); + assert_eq!(res.status(), StatusCode::OK); + let body: Value = res.json().await.expect("invite code response not json"); + body["code"].as_str().expect("missing code").to_string() +} + +#[tokio::test] +async fn check_registration_invite_validates_without_consuming() { + let state = get_test_app_state().await; + + assert_eq!( + state.repos.user.count_users().await.unwrap(), + 0, + "bootstrap branch needs a zero-user instance" + ); + + let mut bootstrap = state.clone(); + bootstrap.bootstrap_invite_code = Some("squid-bootstrap".to_string()); + + assert_eq!( + check_registration_invite(&bootstrap, Some("squid-bootstrap")) + .await + .unwrap(), + InviteRegistration::Bootstrap + ); + assert!(matches!( + check_registration_invite(&bootstrap, Some("whelk")).await, + Err(ApiError::InvalidInviteCode) + )); + assert!(matches!( + check_registration_invite(&bootstrap, None).await, + Err(ApiError::InvalidInviteCode) + )); + + let client = client(); + let (admin_jwt, _did) = create_admin_account_and_login(&client).await; + let code = create_invite_code(&client, &admin_jwt, 1).await; + + assert_eq!( + check_registration_invite(state, Some(&code)).await.unwrap(), + InviteRegistration::Standard(Some(code.clone())) + ); + assert_eq!( + state + .repos + .infra + .get_invite_code_available_uses(&code) + .await + .unwrap(), + Some(1), + "validation must not consume the invite" + ); + + assert_eq!( + check_registration_invite(state, Some(&format!(" {code} "))) + .await + .unwrap(), + InviteRegistration::Standard(Some(code.clone())), + "surrounding whitespace must be trimmed into the validated code" + ); + + assert!(matches!( + check_registration_invite(state, Some("whelk")).await, + Err(ApiError::InvalidInviteCode) + )); +} + +#[tokio::test] +async fn create_account_consumes_invite_code_exactly_once() { + let client = client(); + let (admin_jwt, _did) = create_admin_account_and_login(&client).await; + let code = create_invite_code(&client, &admin_jwt, 2).await; + + let handle = format!("u{}", &uuid::Uuid::new_v4().simple().to_string()[..12]); + let res = client + .post(format!( + "{}/xrpc/com.atproto.server.createAccount", + base_url().await + )) + .json(&json!({ + "handle": handle, + "email": format!("{handle}@nel.pet"), + "password": "Testpass123!", + "inviteCode": code, + })) + .send() + .await + .expect("createAccount request failed"); + let status = res.status(); + let text = res.text().await.unwrap_or_default(); + assert_eq!(status, StatusCode::OK, "createAccount failed: {text}"); + + let state = get_test_app_state().await; + assert_eq!( + state + .repos + .infra + .get_invite_code_available_uses(&code) + .await + .unwrap(), + Some(1), + "a single registration must consume exactly one invite use" + ); +} diff --git a/crates/tranquil-server/src/http3.rs b/crates/tranquil-server/src/http3.rs index 15efe6f..d018255 100644 --- a/crates/tranquil-server/src/http3.rs +++ b/crates/tranquil-server/src/http3.rs @@ -78,7 +78,10 @@ pub fn with_host_from_authority(app: Router) -> Router { .uri() .authority() .map(|a| HeaderValue::from_str(a.as_str())); - match (request.headers().contains_key(http::header::HOST), authority) { + match ( + request.headers().contains_key(http::header::HOST), + authority, + ) { (false, Some(Ok(value))) => { request.headers_mut().insert(http::header::HOST, value); request diff --git a/crates/tranquil-store/src/metastore/client.rs b/crates/tranquil-store/src/metastore/client.rs index ddee573..c38e586 100644 --- a/crates/tranquil-store/src/metastore/client.rs +++ b/crates/tranquil-store/src/metastore/client.rs @@ -1916,35 +1916,6 @@ impl tranquil_db_traits::InfraRepository for MetastoreCl Ok(ValidatedInviteCode::new_validated(code)) } - async fn decrement_invite_code_uses( - &self, - code: &ValidatedInviteCode<'_>, - ) -> Result<(), DbError> { - let (tx, rx) = oneshot::channel(); - self.pool.send(MetastoreRequest::Infra( - InfraRequest::DecrementInviteCodeUses { - code: code.code().to_owned(), - tx, - }, - ))?; - recv(rx).await - } - - async fn record_invite_code_use( - &self, - code: &ValidatedInviteCode<'_>, - used_by_user: Uuid, - ) -> Result<(), DbError> { - let (tx, rx) = oneshot::channel(); - self.pool - .send(MetastoreRequest::Infra(InfraRequest::RecordInviteCodeUse { - code: code.code().to_owned(), - used_by_user, - tx, - }))?; - recv(rx).await - } - async fn get_invite_codes_for_account( &self, for_account: &Did, @@ -5009,17 +4980,6 @@ impl tranquil_db_traits::UserRepository for MetastoreCli recv(rx).await } - async fn check_and_consume_invite_code(&self, code: &str) -> Result { - let (tx, rx) = oneshot::channel(); - self.pool.send(MetastoreRequest::User( - UserRequest::CheckAndConsumeInviteCode { - code: code.to_owned(), - tx, - }, - ))?; - recv(rx).await - } - async fn complete_passkey_setup( &self, input: &CompletePasskeySetupInput, diff --git a/crates/tranquil-store/src/metastore/handler.rs b/crates/tranquil-store/src/metastore/handler.rs index 6b373bd..fea5b39 100644 --- a/crates/tranquil-store/src/metastore/handler.rs +++ b/crates/tranquil-store/src/metastore/handler.rs @@ -38,6 +38,7 @@ use uuid::Uuid; use super::MetastoreError; use super::commit_ops::CommitOps; use super::event_ops::EventOps; +use super::infra_ops::InfraOps; use super::keys::UserHash; use super::record_ops::ListRecordsQuery; use super::user_hash::UserHashMap; @@ -49,6 +50,59 @@ use crate::metastore::Metastore; type Tx = oneshot::Sender>; +fn reserve_invite(infra: &InfraOps, code: Option<&str>) -> Result<(), CreateAccountError> { + match code { + Some(code) => infra.reserve_invite_code(code).map_err(|e| match e { + InviteCodeError::DatabaseError(e) => CreateAccountError::Database(e.to_string()), + _ => CreateAccountError::InviteCodeUnavailable, + }), + None => Ok(()), + } +} + +fn record_invite_use( + infra: &InfraOps, + code: Option<&str>, + user_id: Uuid, +) -> Result<(), CreateAccountError> { + match code { + Some(code) => { + let validated = ValidatedInviteCode::new_validated(code); + infra + .record_invite_code_use(&validated, user_id) + .map_err(|e| CreateAccountError::Database(e.to_string())) + } + None => Ok(()), + } +} + +fn refund_invite(infra: &InfraOps, code: Option<&str>) { + if let Some(code) = code + && let Err(e) = infra.refund_invite_code(code) + { + tracing::error!("failed to refund invite code after account creation failed: {e:?}"); + } +} + +fn finalize_account( + infra: &InfraOps, + code: Option<&str>, + created: Result, + after: impl FnOnce(&T) -> Result, +) -> Result { + match created { + Ok(result) => { + let user_id = after(&result)?; + record_invite_use(infra, code, user_id)?; + Ok(result) + } + Err(e) => { + refund_invite(infra, code); + Err(e) + } + } +} + fn metastore_to_db(e: MetastoreError) -> DbError { match e { MetastoreError::Fjall(e) => DbError::Query(e.to_string()), @@ -1546,10 +1600,6 @@ pub enum UserRequest { CleanupExpiredHandleReservations { tx: Tx, }, - CheckAndConsumeInviteCode { - code: String, - tx: Tx, - }, CompletePasskeySetup { input: CompletePasskeySetupInput, tx: Tx<()>, @@ -1749,7 +1799,6 @@ impl UserRequest { | Self::GetUserForPasskeyRecovery { .. } | Self::GetAccountsScheduledForDeletion { .. } | Self::CleanupExpiredHandleReservations { .. } - | Self::CheckAndConsumeInviteCode { .. } | Self::GetPasswordResetInfo { .. } | Self::ExpirePasswordResetCode { .. } | Self::SaveDiscoverableChallenge { .. } @@ -1810,15 +1859,6 @@ pub enum InfraRequest { code: String, tx: oneshot::Sender>, }, - DecrementInviteCodeUses { - code: String, - tx: Tx<()>, - }, - RecordInviteCodeUse { - code: String, - used_by_user: Uuid, - tx: Tx<()>, - }, GetInviteCodesForAccount { for_account: Did, tx: Tx>, @@ -2061,9 +2101,6 @@ impl InfraRequest { Self::CreateInviteCodesBatch { created_by_user, .. } => uuid_to_routing(user_hashes, created_by_user), - Self::RecordInviteCodeUse { used_by_user, .. } => { - uuid_to_routing(user_hashes, used_by_user) - } Self::EnqueueComms { user_id: Some(uid), .. } => uuid_to_routing(user_hashes, uid), @@ -3936,28 +3973,6 @@ fn dispatch_infra(state: &HandlerState, req: InfraRequest) { .map(|_| ()); let _ = tx.send(result); } - InfraRequest::DecrementInviteCodeUses { code, tx } => { - let validated = ValidatedInviteCode::new_validated(&code); - let result = state - .metastore - .infra_ops() - .decrement_invite_code_uses(&validated) - .map_err(metastore_to_db); - let _ = tx.send(result); - } - InfraRequest::RecordInviteCodeUse { - code, - used_by_user, - tx, - } => { - let validated = ValidatedInviteCode::new_validated(&code); - let result = state - .metastore - .infra_ops() - .record_invite_code_use(&validated, used_by_user) - .map_err(metastore_to_db); - let _ = tx.send(result); - } InfraRequest::GetInviteCodesForAccount { for_account, tx } => { let result = state .metastore @@ -5799,15 +5814,17 @@ fn dispatch_user(state: &HandlerState, req: UserReque let _ = tx.send(result.map(|_| ())); } UserRequest::CreatePasswordAccount { input, tx } => { - let result = user.create_password_account(&input).and_then(|result| { - if let Some(key_id) = input.reserved_key_id { - state - .metastore - .infra_ops() - .mark_signing_key_used(key_id) - .map_err(|e| CreateAccountError::Database(e.to_string()))?; - } - Ok(result) + let infra = state.metastore.infra_ops(); + let code = input.invite_code.as_deref(); + let result = reserve_invite(&infra, code).and_then(|()| { + finalize_account(&infra, code, user.create_password_account(&input), |result| { + if let Some(key_id) = input.reserved_key_id { + infra + .mark_signing_key_used(key_id) + .map_err(|e| CreateAccountError::Database(e.to_string()))?; + } + Ok(result.user_id) + }) }); let _ = tx.send(result); } @@ -5838,35 +5855,41 @@ fn dispatch_user(state: &HandlerState, req: UserReque let _ = tx.send(result); } UserRequest::CreatePasskeyAccount { input, tx } => { - let result = user.create_passkey_account(&input).and_then(|result| { - if let Some(key_id) = input.reserved_key_id { - state - .metastore - .infra_ops() - .mark_signing_key_used(key_id) - .map_err(|e| CreateAccountError::Database(e.to_string()))?; - } - Ok(result) + let infra = state.metastore.infra_ops(); + let code = input.invite_code.as_deref(); + let result = reserve_invite(&infra, code).and_then(|()| { + finalize_account(&infra, code, user.create_passkey_account(&input), |result| { + if let Some(key_id) = input.reserved_key_id { + infra + .mark_signing_key_used(key_id) + .map_err(|e| CreateAccountError::Database(e.to_string()))?; + } + Ok(result.user_id) + }) }); let _ = tx.send(result); } UserRequest::CreateSsoAccount { input, tx } => { let sso_ops = state.metastore.sso_ops(); + let infra = state.metastore.infra_ops(); + let code = input.invite_code.as_deref(); let result = sso_ops .consume_pending_registration(&input.pending_registration_token) .map_err(|e| CreateAccountError::Database(e.to_string())) .and_then(|consumed| match consumed { - Some(_) => user.create_sso_account(&input).and_then(|result| { - sso_ops - .create_external_identity( - &input.did, - input.sso_provider, - &input.sso_provider_user_id, - input.sso_provider_username.as_deref(), - input.sso_provider_email.as_deref(), - ) - .map_err(|e| CreateAccountError::Database(e.to_string()))?; - Ok(result) + Some(_) => reserve_invite(&infra, code).and_then(|()| { + finalize_account(&infra, code, user.create_sso_account(&input), |result| { + sso_ops + .create_external_identity( + &input.did, + input.sso_provider, + &input.sso_provider_user_id, + input.sso_provider_username.as_deref(), + input.sso_provider_email.as_deref(), + ) + .map_err(|e| CreateAccountError::Database(e.to_string()))?; + Ok(result.user_id) + }) }), None => Err(CreateAccountError::InvalidToken), }); @@ -5903,17 +5926,6 @@ fn dispatch_user(state: &HandlerState, req: UserReque .map_err(metastore_to_db), ); } - UserRequest::CheckAndConsumeInviteCode { code, tx } => { - let infra = state.metastore.infra_ops(); - let result = match infra.validate_invite_code(&code) { - Ok(validated) => infra - .decrement_invite_code_uses(&validated) - .map(|()| true) - .map_err(metastore_to_db), - Err(_) => Ok(false), - }; - let _ = tx.send(result); - } UserRequest::CompletePasskeySetup { input, tx } => { let _ = tx.send(user.complete_passkey_setup(&input).map_err(metastore_to_db)); } @@ -6187,6 +6199,84 @@ mod tests { assert_eq!(repo.repo_rev.as_deref(), Some("rev1")); } + #[test] + fn reserve_invite_code_consumes_once_and_guards_exhaustion() { + let dir = tempfile::TempDir::new().unwrap(); + let metastore = Metastore::open( + dir.path(), + MetastoreConfig { + cache_size_bytes: 64 * 1024 * 1024, + }, + ) + .unwrap(); + let infra = metastore.infra_ops(); + + assert!(infra.create_invite_code("squid-invite", 1, None).unwrap()); + + infra.reserve_invite_code("squid-invite").unwrap(); + assert_eq!( + infra + .get_invite_code_available_uses("squid-invite") + .unwrap(), + Some(0) + ); + + assert!(matches!( + infra.reserve_invite_code("squid-invite"), + Err(InviteCodeError::ExhaustedUses) + )); + assert_eq!( + infra + .get_invite_code_available_uses("squid-invite") + .unwrap(), + Some(0), + "a rejected reservation must not decrement below zero" + ); + + assert!(matches!( + infra.reserve_invite_code("whelk"), + Err(InviteCodeError::NotFound) + )); + } + + #[test] + fn refund_invite_code_restores_a_reserved_use() { + let dir = tempfile::TempDir::new().unwrap(); + let metastore = Metastore::open( + dir.path(), + MetastoreConfig { + cache_size_bytes: 64 * 1024 * 1024, + }, + ) + .unwrap(); + let infra = metastore.infra_ops(); + + assert!(infra.create_invite_code("squid-invite", 1, None).unwrap()); + + infra.reserve_invite_code("squid-invite").unwrap(); + infra.refund_invite_code("squid-invite").unwrap(); + assert_eq!( + infra + .get_invite_code_available_uses("squid-invite") + .unwrap(), + Some(1), + "refund must return the reserved use so a failed signup does not burn it" + ); + + infra.reserve_invite_code("squid-invite").unwrap(); + assert_eq!( + infra + .get_invite_code_available_uses("squid-invite") + .unwrap(), + Some(0) + ); + + assert!(matches!( + infra.refund_invite_code("whelk"), + Err(InviteCodeError::NotFound) + )); + } + #[test] fn routing_determinism() { let user_id = Uuid::from_u128(0x12345678); diff --git a/crates/tranquil-store/src/metastore/infra_ops.rs b/crates/tranquil-store/src/metastore/infra_ops.rs index cef98dd..89f8c9f 100644 --- a/crates/tranquil-store/src/metastore/infra_ops.rs +++ b/crates/tranquil-store/src/metastore/infra_ops.rs @@ -37,6 +37,7 @@ pub struct InfraOps { users: Keyspace, user_hashes: Arc, comms_seq: Arc, + counter_lock: Arc>, } impl InfraOps { @@ -47,6 +48,7 @@ impl InfraOps { users: Keyspace, user_hashes: Arc, comms_seq: Arc, + counter_lock: Arc>, ) -> Self { Self { db, @@ -55,6 +57,7 @@ impl InfraOps { users, user_hashes, comms_seq, + counter_lock, } } @@ -450,6 +453,36 @@ impl InfraOps { } } + pub fn reserve_invite_code(&self, code: &str) -> Result<(), InviteCodeError> { + let _guard = self.counter_lock.lock(); + let validated = self.validate_invite_code(code)?; + self.decrement_invite_code_uses(&validated).map_err(|e| { + InviteCodeError::DatabaseError(tranquil_db_traits::DbError::Query(e.to_string())) + }) + } + + pub fn refund_invite_code(&self, code: &str) -> Result<(), InviteCodeError> { + let _guard = self.counter_lock.lock(); + let key = invite_code_key(code); + let mut val: InviteCodeValue = point_lookup( + &self.infra, + key.as_slice(), + InviteCodeValue::deserialize, + "corrupt invite code", + ) + .map_err(|e| { + InviteCodeError::DatabaseError(tranquil_db_traits::DbError::Query(e.to_string())) + })? + .ok_or(InviteCodeError::NotFound)?; + + val.available_uses += 1; + self.infra + .insert(key.as_slice(), val.serialize()) + .map_err(|e| { + InviteCodeError::DatabaseError(tranquil_db_traits::DbError::Query(e.to_string())) + }) + } + pub fn decrement_invite_code_uses( &self, code: &ValidatedInviteCode<'_>, diff --git a/crates/tranquil-store/src/metastore/mod.rs b/crates/tranquil-store/src/metastore/mod.rs index 39af6ed..9132f34 100644 --- a/crates/tranquil-store/src/metastore/mod.rs +++ b/crates/tranquil-store/src/metastore/mod.rs @@ -339,6 +339,7 @@ impl Metastore { self.partitions[Partition::Users.index()].clone(), Arc::clone(&self.user_hashes), Arc::clone(&self.comms_seq), + Arc::clone(&self.counter_lock), ) }