From 63d84d38fb76ed260586991fb14ec2113432fcf5 Mon Sep 17 00:00:00 2001 From: ave Date: Fri, 5 Jun 2026 23:26:44 +0000 Subject: [PATCH] refactor(auth): unify short-code generation onto util::generate_token_code Collapse the three ad-hoc short-code generators into one canonical generator plus a shared normalizer: - util::generate_token_code now emits the uppercase base32 XXXXX-XXXXX display form; new util::normalize_token_code canonicalizes user input (uppercase, strip hyphen/whitespace). - email_token and legacy_2fa now generate via util, store the normalized form, and compare normalized input. Their private generate_short_token/generate_code (and BASE32_CHARS/CODE_LENGTH) are removed. - PLC (request/sign) and password reset inline util::generate_token_code, persist the normalized form, email the display form, and normalize input before lookup. The generate_plc_token/generate_reset_code wrappers are removed. Behavior changes: legacy login-2FA codes go from 8-digit numeric to XXXXX-XXXXX; PLC and password-reset codes go from lowercase to uppercase. All four code types are now accepted case-/hyphen-insensitively. OAuth web-login 2FA, account deletion, and the long verification_token blobs are intentionally untouched. Tests: add util normalize tests + email_token/legacy_2fa case/hyphen tests; update integration tests to expect the canonical stored form and the new emailed format. --- .../tranquil-api/src/identity/plc/request.rs | 11 ++-- crates/tranquil-api/src/identity/plc/sign.rs | 3 +- crates/tranquil-api/src/server/password.rs | 19 +++--- crates/tranquil-pds/src/auth/email_token.rs | 58 +++++++++-------- crates/tranquil-pds/src/auth/legacy_2fa.rs | 64 ++++++++++++------- crates/tranquil-pds/src/util.rs | 51 ++++++++++++++- crates/tranquil-pds/tests/legacy_2fa.rs | 22 +++---- crates/tranquil-pds/tests/password_reset.rs | 12 +++- crates/tranquil-pds/tests/plc_operations.rs | 16 ++++- 9 files changed, 170 insertions(+), 86 deletions(-) diff --git a/crates/tranquil-api/src/identity/plc/request.rs b/crates/tranquil-api/src/identity/plc/request.rs index 77d0354..e1e9a6a 100644 --- a/crates/tranquil-api/src/identity/plc/request.rs +++ b/crates/tranquil-api/src/identity/plc/request.rs @@ -6,10 +6,6 @@ use tranquil_pds::api::error::{ApiError, DbResultExt}; use tranquil_pds::auth::{Auth, Permissive}; use tranquil_pds::state::AppState; -fn generate_plc_token() -> String { - tranquil_pds::util::generate_token_code() -} - pub async fn request_plc_operation_signature( State(state): State, auth: Auth, @@ -28,12 +24,13 @@ pub async fn request_plc_operation_signature( .ok_or(ApiError::AccountNotFound)?; let _ = state.repos.infra.delete_plc_tokens_for_user(user_id).await; - let plc_token = generate_plc_token(); + let display_token = tranquil_pds::util::generate_token_code(); + let stored_token = tranquil_pds::util::normalize_token_code(&display_token); let expires_at = Utc::now() + Duration::minutes(10); state .repos .infra - .insert_plc_token(user_id, &plc_token, expires_at) + .insert_plc_token(user_id, &stored_token, expires_at) .await .log_db_err("creating PLC token")?; @@ -42,7 +39,7 @@ pub async fn request_plc_operation_signature( state.repos.user.as_ref(), state.repos.infra.as_ref(), user_id, - &plc_token, + &display_token, hostname, ) .await diff --git a/crates/tranquil-api/src/identity/plc/sign.rs b/crates/tranquil-api/src/identity/plc/sign.rs index 4d144ab..66dceca 100644 --- a/crates/tranquil-api/src/identity/plc/sign.rs +++ b/crates/tranquil-api/src/identity/plc/sign.rs @@ -43,9 +43,10 @@ pub async fn sign_plc_operation( "PLC operations are only valid for did:plc identities".into(), )); } - let token = input.token.as_ref().ok_or_else(|| { + let raw_token = input.token.as_ref().ok_or_else(|| { ApiError::InvalidRequest("Email confirmation token required to sign PLC operations".into()) })?; + let token = &tranquil_pds::util::normalize_token_code(raw_token); let user_id = state .repos diff --git a/crates/tranquil-api/src/server/password.rs b/crates/tranquil-api/src/server/password.rs index 7d81064..393afbf 100644 --- a/crates/tranquil-api/src/server/password.rs +++ b/crates/tranquil-api/src/server/password.rs @@ -13,10 +13,6 @@ use tranquil_pds::state::AppState; use tranquil_pds::types::PlainPassword; use tranquil_pds::validation::validate_password; -fn generate_reset_code() -> String { - tranquil_pds::util::generate_token_code() -} - #[derive(Deserialize)] pub struct RequestPasswordResetInput { #[serde(alias = "identifier")] @@ -70,12 +66,13 @@ pub async fn request_password_reset( return Err(ApiError::InternalError(None)); } }; - let code = generate_reset_code(); + let display_code = tranquil_pds::util::generate_token_code(); + let stored_code = tranquil_pds::util::normalize_token_code(&display_code); let expires_at = Utc::now() + Duration::minutes(10); if let Err(e) = state .repos .user - .set_password_reset_code(user_id, &code, expires_at) + .set_password_reset_code(user_id, &stored_code, expires_at) .await { error!("DB error setting reset code: {:?}", e); @@ -86,7 +83,7 @@ pub async fn request_password_reset( state.repos.user.as_ref(), state.repos.infra.as_ref(), user_id, - &code, + &display_code, hostname, ) .await @@ -133,7 +130,13 @@ pub async fn reset_password( if let Err(e) = validate_password(password) { return Err(ApiError::InvalidRequest(e.to_string())); } - let user = match state.repos.user.get_user_by_reset_code(token).await { + let normalized_token = tranquil_pds::util::normalize_token_code(token); + let user = match state + .repos + .user + .get_user_by_reset_code(&normalized_token) + .await + { Ok(Some(u)) => u, Ok(None) => { return Err(ApiError::InvalidToken(None)); diff --git a/crates/tranquil-pds/src/auth/email_token.rs b/crates/tranquil-pds/src/auth/email_token.rs index 20b5156..1a0a650 100644 --- a/crates/tranquil-pds/src/auth/email_token.rs +++ b/crates/tranquil-pds/src/auth/email_token.rs @@ -1,11 +1,10 @@ -use rand::Rng; use serde::{Deserialize, Serialize}; use std::time::Duration; use crate::cache::Cache; +use crate::util::{generate_token_code, normalize_token_code}; const TOKEN_TTL_SECS: u64 = 900; -const BASE32_CHARS: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ234567"; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum EmailTokenPurpose { @@ -46,14 +45,6 @@ fn cache_key(did: &str, purpose: EmailTokenPurpose) -> String { format!("email_token:{}:{}", purpose.as_str(), did) } -fn generate_short_token() -> String { - let mut rng = rand::thread_rng(); - let token: String = (0..10) - .map(|_| BASE32_CHARS[rng.gen_range(0..BASE32_CHARS.len())] as char) - .collect(); - format!("{}-{}", &token[0..5], &token[5..10]) -} - fn current_timestamp() -> u64 { u64::try_from(chrono::Utc::now().timestamp()).unwrap_or(0) } @@ -67,9 +58,9 @@ pub async fn create_email_token( return Err(TokenError::CacheUnavailable); } - let token = generate_short_token(); + let token = generate_token_code(); let data = TokenData { - token: token.clone(), + token: normalize_token_code(&token), created_at: current_timestamp(), }; @@ -108,10 +99,9 @@ pub async fn validate_email_token( return Err(TokenError::ExpiredToken); } - let normalized_input = token.to_uppercase().replace('-', ""); - let normalized_stored = data.token.to_uppercase().replace('-', ""); + let normalized_input = normalize_token_code(token); - if !constant_time_eq(normalized_input.as_bytes(), normalized_stored.as_bytes()) { + if !constant_time_eq(normalized_input.as_bytes(), data.token.as_bytes()) { return Err(TokenError::InvalidToken); } @@ -260,20 +250,19 @@ mod tests { #[tokio::test] async fn test_token_format() { - (0..100).for_each(|_| { - let token = generate_short_token(); + // The emitted token is the display form: uppercase `XXXXX-XXXXX`. + let cache = MockCache::new(); + let did = "did:plc:test123"; + (0..50).for_each(|_| { + let token = futures::executor::block_on(create_email_token( + &cache, + did, + EmailTokenPurpose::UpdateEmail, + )) + .unwrap(); assert_eq!(token.len(), 11); assert_eq!(&token[5..6], "-"); - assert!( - token[0..5] - .chars() - .all(|c| BASE32_CHARS.contains(&(c as u8))) - ); - assert!( - token[6..11] - .chars() - .all(|c| BASE32_CHARS.contains(&(c as u8))) - ); + assert_eq!(token, token.to_uppercase()); }); } @@ -292,6 +281,21 @@ mod tests { assert!(result.is_ok()); } + #[tokio::test] + async fn test_hyphen_insensitive_validation() { + let cache = MockCache::new(); + let did = "did:plc:test123"; + + let token = create_email_token(&cache, did, EmailTokenPurpose::UpdateEmail) + .await + .unwrap(); + + let no_hyphen = token.replace('-', ""); + let result = + validate_email_token(&cache, did, EmailTokenPurpose::UpdateEmail, &no_hyphen).await; + assert!(result.is_ok()); + } + #[tokio::test] async fn test_noop_cache_returns_unavailable() { let cache = crate::cache::NoOpCache; diff --git a/crates/tranquil-pds/src/auth/legacy_2fa.rs b/crates/tranquil-pds/src/auth/legacy_2fa.rs index daf6ae6..d526235 100644 --- a/crates/tranquil-pds/src/auth/legacy_2fa.rs +++ b/crates/tranquil-pds/src/auth/legacy_2fa.rs @@ -1,15 +1,14 @@ use chrono::Utc; -use rand::Rng; use serde::{Deserialize, Serialize}; use std::time::Duration; use crate::cache::Cache; use crate::types::Did; +use crate::util::{generate_token_code, normalize_token_code}; const CHALLENGE_TTL_SECS: u64 = 300; const MIN_REMAINING_TTL_SECS: u64 = 10; const MAX_ATTEMPTS: u8 = 5; -const CODE_LENGTH: usize = 8; const COOLDOWN_SECS: u64 = 60; #[derive(Debug, Clone, Serialize, Deserialize)] @@ -94,7 +93,8 @@ async fn validate_challenge_internal( return Err(ValidationError::ChallengeExpired); } - if !constant_time_eq(code.as_bytes(), data.code.as_bytes()) { + let normalized_input = normalize_token_code(code); + if !constant_time_eq(normalized_input.as_bytes(), data.code.as_bytes()) { let updated = ChallengeData { code: data.code, attempts: data.attempts + 1, @@ -127,13 +127,6 @@ fn cooldown_key(did: &str) -> String { format!("legacy_2fa_cooldown:{}", did) } -fn generate_code() -> String { - let mut rng = rand::thread_rng(); - (0..CODE_LENGTH) - .map(|_| rng.gen_range(0..10).to_string()) - .collect() -} - fn current_timestamp() -> u64 { u64::try_from(Utc::now().timestamp()).unwrap_or(0) } @@ -219,11 +212,11 @@ async fn create_challenge_code( return Err(ChallengeError::RateLimited); } - let code = generate_code(); + let display = generate_token_code(); let now = current_timestamp(); let data = ChallengeData { - code: code.clone(), + code: normalize_token_code(&display), attempts: 0, created_at: now, }; @@ -244,7 +237,7 @@ async fn create_challenge_code( .await .map_err(|_| ChallengeError::CacheError)?; - Ok(ChallengeCode(code)) + Ok(ChallengeCode(display)) } #[derive(Debug)] @@ -332,12 +325,46 @@ mod tests { let did = Did::new("did:plc:test123".to_string()).unwrap(); let code = create_challenge(&cache, &did).await.unwrap(); - assert_eq!(code.as_str().len(), CODE_LENGTH); + assert_eq!(code.as_str().len(), 11); let result = validate_challenge(&cache, &did, code.as_str()).await; assert!(result.is_ok()); } + #[tokio::test] + async fn test_challenge_code_format() { + let cache = MockCache::new(); + let did = Did::new("did:plc:test123".to_string()).unwrap(); + + let code = create_challenge(&cache, &did).await.unwrap(); + let code = code.as_str(); + assert_eq!(code.len(), 11); + assert_eq!(&code[5..6], "-"); + assert_eq!(code, code.to_uppercase()); + } + + #[tokio::test] + async fn test_case_insensitive_validation() { + let cache = MockCache::new(); + let did = Did::new("did:plc:test123".to_string()).unwrap(); + + let code = create_challenge(&cache, &did).await.unwrap(); + let lowercase = code.as_str().to_lowercase(); + let result = validate_challenge(&cache, &did, &lowercase).await; + assert!(result.is_ok()); + } + + #[tokio::test] + async fn test_hyphen_insensitive_validation() { + let cache = MockCache::new(); + let did = Did::new("did:plc:test123".to_string()).unwrap(); + + let code = create_challenge(&cache, &did).await.unwrap(); + let no_hyphen = code.as_str().replace('-', ""); + let result = validate_challenge(&cache, &did, &no_hyphen).await; + assert!(result.is_ok()); + } + #[tokio::test] async fn test_invalid_code_rejected() { let cache = MockCache::new(); @@ -396,15 +423,6 @@ mod tests { assert_eq!(result.unwrap_err(), ChallengeError::CacheUnavailable); } - #[tokio::test] - async fn test_code_generation_is_numeric() { - (0..100).for_each(|_| { - let code = generate_code(); - assert!(code.chars().all(|c| c.is_ascii_digit())); - assert_eq!(code.len(), CODE_LENGTH); - }); - } - #[tokio::test] async fn test_constant_time_eq() { assert!(constant_time_eq(b"12345678", b"12345678")); diff --git a/crates/tranquil-pds/src/util.rs b/crates/tranquil-pds/src/util.rs index 31bf546..57e7b75 100644 --- a/crates/tranquil-pds/src/util.rs +++ b/crates/tranquil-pds/src/util.rs @@ -33,7 +33,20 @@ pub fn generate_token_code() -> String { .map(|_| chars[rng.gen_range(0..chars.len())]) .collect() }; - format!("{}-{}", gen_segment(&mut rng), gen_segment(&mut rng)) + // Human-entered short codes are displayed in uppercase; base32 digits are + // unaffected by the conversion. + format!("{}-{}", gen_segment(&mut rng), gen_segment(&mut rng)).to_uppercase() +} + +/// Normalize a user-entered short code so that codes are accepted +/// case-insensitively and regardless of the separating hyphen or surrounding +/// whitespace. +pub fn normalize_token_code(input: &str) -> String { + input + .chars() + .filter(|c| !c.is_whitespace() && *c != '-') + .collect::() + .to_uppercase() } pub fn parse_repeated_query_param(query: Option<&str>, key: &str) -> Vec { @@ -491,10 +504,44 @@ mod tests { assert!( code.chars() .filter(|&c| c != '-') - .all(|c| BASE32_ALPHABET.contains(c)) + .all(|c| BASE32_ALPHABET.to_uppercase().contains(c)) ); } + #[test] + fn test_generate_token_code_is_uppercase() { + (0..100).for_each(|_| { + let code = generate_token_code(); + assert_eq!(code, code.to_uppercase(), "code must be uppercase: {code}"); + }); + } + + #[test] + fn test_normalize_token_code_strips_hyphen_and_uppercases() { + assert_eq!(normalize_token_code("k7m2p-q9rst"), "K7M2PQ9RST"); + assert_eq!(normalize_token_code("K7M2P-Q9RST"), "K7M2PQ9RST"); + } + + #[test] + fn test_normalize_token_code_strips_whitespace() { + assert_eq!(normalize_token_code(" k7m2p-q9rst \n"), "K7M2PQ9RST"); + } + + #[test] + fn test_normalize_token_code_is_idempotent() { + let once = normalize_token_code("k7m2p-q9rst"); + assert_eq!(normalize_token_code(&once), once); + } + + #[test] + fn test_generated_code_round_trips_through_normalize() { + let code = generate_token_code(); + // A user re-typing the displayed code lowercased and without the hyphen + // must normalize to the same canonical form as the code itself. + let retyped = code.to_lowercase().replace('-', ""); + assert_eq!(normalize_token_code(&code), normalize_token_code(&retyped)); + } + #[test] fn test_json_to_ipld_cid_link() { let json = serde_json::json!({ diff --git a/crates/tranquil-pds/tests/legacy_2fa.rs b/crates/tranquil-pds/tests/legacy_2fa.rs index 1fac06c..d7f2d6f 100644 --- a/crates/tranquil-pds/tests/legacy_2fa.rs +++ b/crates/tranquil-pds/tests/legacy_2fa.rs @@ -40,19 +40,17 @@ async fn get_2fa_code_from_queue(did: &str) -> Option { .await .ok()?; + const ALPHABET: &str = "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567"; comms.first().and_then(|c| { - c.body - .lines() - .find(|line: &&str| line.chars().all(|c: char| c.is_ascii_digit()) && line.len() == 8) - .map(|s: &str| s.to_string()) - .or_else(|| { - c.body - .split_whitespace() - .find(|word: &&str| { - word.chars().all(|c: char| c.is_ascii_digit()) && word.len() == 8 - }) - .map(|s: &str| s.to_string()) - }) + c.body.split_whitespace().find_map(|word: &str| { + let candidate = word.trim_matches(|ch: char| !ch.is_ascii_alphanumeric() && ch != '-'); + let normalized = candidate.replace('-', ""); + if normalized.len() == 10 && normalized.chars().all(|ch| ALPHABET.contains(ch)) { + Some(candidate.to_string()) + } else { + None + } + }) }) } diff --git a/crates/tranquil-pds/tests/password_reset.rs b/crates/tranquil-pds/tests/password_reset.rs index 43006e1..82d9ee4 100644 --- a/crates/tranquil-pds/tests/password_reset.rs +++ b/crates/tranquil-pds/tests/password_reset.rs @@ -45,9 +45,12 @@ async fn test_request_password_reset_creates_code() { .expect("user not found"); assert!(info.code.is_some()); assert!(info.expires_at.is_some()); + // The stored code is normalized: uppercase base32, 10 chars, no hyphen. + // The hyphenated display form only appears in the email. let code = info.code.unwrap(); - assert!(code.contains('-')); - assert_eq!(code.len(), 11); + assert!(!code.contains('-')); + assert_eq!(code.len(), 10); + assert_eq!(code, code.to_uppercase()); } #[tokio::test] @@ -109,7 +112,10 @@ async fn test_reset_password_with_valid_token() { .await .expect("failed to look up user") .expect("user not found"); - let token = info.code.expect("No reset code"); + let stored = info.code.expect("No reset code"); + // Submit a variant a user might actually type: lowercased, with the display + // hyphen re-inserted. Normalization must still accept it. + let token = format!("{}-{}", &stored[0..5], &stored[5..10]).to_lowercase(); let res = client .post(format!( "{}/xrpc/com.atproto.server.resetPassword", diff --git a/crates/tranquil-pds/tests/plc_operations.rs b/crates/tranquil-pds/tests/plc_operations.rs index 01150c0..e62908c 100644 --- a/crates/tranquil-pds/tests/plc_operations.rs +++ b/crates/tranquil-pds/tests/plc_operations.rs @@ -188,12 +188,22 @@ async fn test_plc_token_lifecycle() { "PLC token should be created in database" ); let first = &tokens[0]; + // The token is persisted in canonical (normalized) form: uppercase base32, + // 10 chars, no hyphen. The hyphenated display form only appears in the email. assert_eq!( first.token.len(), - 11, - "Token should be in format xxxxx-xxxxx" + 10, + "Stored token should be the 10-char canonical form" + ); + assert!( + !first.token.contains('-'), + "Stored token should not contain a hyphen" + ); + assert_eq!( + first.token, + first.token.to_uppercase(), + "Stored token should be uppercase" ); - assert!(first.token.contains('-'), "Token should contain hyphen"); assert!( first.expires_at > chrono::Utc::now(), "Token should not be expired"