mirrors/tranquil-pds
1
0
Fork 0
mirror of https://tangled.org/tranquil.farm/tranquil-pds synced 2026-08-17 06:46:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

invite codes: dedup consumption, iron out kinks

Browse Source 
Lewis: May this revision serve well! <lu5a@proton.me>
This commit is contained in:
Lewis
2026-06-26 13:28:49 +03:00
committed by Tangled
parent 05ab0b7423
commit 39a2e40b35
17 changed files with 444 additions and 363 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.sqlx/query-1ee6eda3e44660e7f14fcfe56adc2d41c72901b9c701fc7b992314e5370b32dc.json → .sqlx/query-0bb2cb6af37bff735b6b380697fc8e1fa2034ca0600e8c16e1e362b722192327.json
Generated
+2 -2
View File
@@ -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"
}
crates/tranquil-api/src/delegation.rs
+2 -27
View File
@@ -330,7 +330,6 @@ pub struct CreateDelegatedAccountInput {
pub handle: String,
pub email: Option<String>,
pub controller_scopes: ValidatedDelegationScope,
pub invite_code: Option<String>,
}
#[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(
crates/tranquil-api/src/identity/account.rs
+10 -39
View File
@@ -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();
crates/tranquil-api/src/server/passkey_account.rs
+7 -25
View File
@@ -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));
crates/tranquil-db-traits/src/infra.rs
-11
View File
@@ -268,17 +268,6 @@ pub trait InfraRepository: Send + Sync {
code: &'a str,
) -> Result<ValidatedInviteCode<'a>, 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,
crates/tranquil-db-traits/src/user.rs
+1 -3
View File
@@ -589,8 +589,6 @@ pub trait UserRepository: Send + Sync {
async fn cleanup_expired_handle_reservations(&self) -> Result<u64, DbError>;
async fn check_and_consume_invite_code(&self, code: &str) -> Result<bool, DbError>;
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<Vec<u8>>,
pub invite_code: Option<String>,
}
#[derive(Debug, Clone)]
crates/tranquil-db/src/postgres/infra.rs
-32
View File
@@ -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,
crates/tranquil-db/src/postgres/user.rs
+35 -91
View File
@@ -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<Option<UserRow>, 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<bool, DbError> {
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,
crates/tranquil-oauth-server/src/sso_endpoints.rs
+7 -13
View File
@@ -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));
crates/tranquil-pds/src/api/invite.rs
+51
View File
@@ -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<String>),
}
impl InviteRegistration {
pub fn into_invite_code(self) -> Option<String> {
match self {
InviteRegistration::Bootstrap => None,
InviteRegistration::Standard(code) => code,
}
}
}
pub async fn check_registration_invite(
state: &AppState,
invite_code: Option<&str>,
) -> Result<InviteRegistration, ApiError> {
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)),
},
}
}
crates/tranquil-pds/src/api/mod.rs
+1
View File
@@ -1,4 +1,5 @@
pub mod error;
pub mod invite;
pub mod proxy;
pub mod proxy_client;
pub mod responses;
crates/tranquil-pds/tests/invite_registration.rs
+121
View File
@@ -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"
);
}
crates/tranquil-server/src/http3.rs
+4 -1
View File
@@ -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
crates/tranquil-store/src/metastore/client.rs
-40
View File
@@ -1916,35 +1916,6 @@ impl<S: StorageIO + 'static> 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<S: StorageIO + 'static> tranquil_db_traits::UserRepository for MetastoreCli
recv(rx).await
}
async fn check_and_consume_invite_code(&self, code: &str) -> Result<bool, DbError> {
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,
crates/tranquil-store/src/metastore/handler.rs
+169 -79
View File
@@ -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<T> = oneshot::Sender<Result<T, DbError>>;
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<T>(
infra: &InfraOps,
code: Option<&str>,
created: Result<T, CreateAccountError>,
after: impl FnOnce(&T) -> Result<Uuid, CreateAccountError>,
) -> Result<T, CreateAccountError> {
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<u64>,
},
CheckAndConsumeInviteCode {
code: String,
tx: Tx<bool>,
},
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<Result<(), InviteCodeError>>,
},
DecrementInviteCodeUses {
code: String,
tx: Tx<()>,
},
RecordInviteCodeUse {
code: String,
used_by_user: Uuid,
tx: Tx<()>,
},
GetInviteCodesForAccount {
for_account: Did,
tx: Tx<Vec<InviteCodeInfo>>,
@@ -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<S: StorageIO>(state: &HandlerState<S>, 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<S: StorageIO + 'static>(state: &HandlerState<S>, 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<S: StorageIO + 'static>(state: &HandlerState<S>, 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<S: StorageIO + 'static>(state: &HandlerState<S>, 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);
crates/tranquil-store/src/metastore/infra_ops.rs
+33
View File
@@ -37,6 +37,7 @@ pub struct InfraOps {
users: Keyspace,
user_hashes: Arc<UserHashMap>,
comms_seq: Arc<std::sync::atomic::AtomicU32>,
counter_lock: Arc<parking_lot::Mutex<()>>,
}
impl InfraOps {
@@ -47,6 +48,7 @@ impl InfraOps {
users: Keyspace,
user_hashes: Arc<UserHashMap>,
comms_seq: Arc<std::sync::atomic::AtomicU32>,
counter_lock: Arc<parking_lot::Mutex<()>>,
) -> 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<'_>,
crates/tranquil-store/src/metastore/mod.rs
+1
View File
@@ -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),
)
}