From ac77c0263e3be851f3890d0854fa7813d5063f36 Mon Sep 17 00:00:00 2001 From: Anthony Ronning <101225832+AnthonyRonning@users.noreply.github.com> Date: Tue, 14 Jul 2026 18:08:12 +0000 Subject: [PATCH 1/2] Bound OAuth state lifecycle --- src/oauth.rs | 314 ++++++++++++++++++++++++++++++++++------ src/web/oauth_routes.rs | 5 +- 2 files changed, 270 insertions(+), 49 deletions(-) diff --git a/src/oauth.rs b/src/oauth.rs index bbee334a..2fcdb566 100644 --- a/src/oauth.rs +++ b/src/oauth.rs @@ -8,22 +8,128 @@ use oauth2::{ use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::sync::Arc; -use tokio::sync::RwLock; +use std::time::{Duration, Instant}; +use tokio::sync::Mutex; use tracing::{debug, error, info}; use uuid::Uuid; +// OAuth redirects normally complete within seconds. Ten minutes leaves room for a user to +// authenticate without retaining abandoned CSRF state for the encrypted session lifetime. +const OAUTH_STATE_TTL: Duration = Duration::from_secs(10 * 60); +// This limit applies independently to each configured provider. +const OAUTH_STATE_CAPACITY: usize = 4_096; + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct OAuthState { pub csrf_token: String, pub client_id: Uuid, } +#[derive(Debug, Clone)] +struct OAuthStateStore { + inner: Arc>, + ttl: Duration, + capacity: usize, +} + +#[derive(Debug, Default)] +struct OAuthStateStoreInner { + entries: HashMap, + next_sequence: u64, +} + +#[derive(Debug)] +struct StoredOAuthState { + state: OAuthState, + expires_at: Instant, + sequence: u64, +} + +impl OAuthStateStore { + fn new() -> Self { + Self::with_limits(OAUTH_STATE_TTL, OAUTH_STATE_CAPACITY) + } + + fn with_limits(ttl: Duration, capacity: usize) -> Self { + assert!(capacity > 0, "OAuth state capacity must be positive"); + Self { + inner: Arc::new(Mutex::new(OAuthStateStoreInner::default())), + ttl, + capacity, + } + } + + async fn store(&self, csrf_token: &str, state: OAuthState) { + self.store_at(csrf_token, state, Instant::now()).await; + } + + async fn store_at(&self, csrf_token: &str, state: OAuthState, now: Instant) { + let mut inner = self.inner.lock().await; + inner + .entries + .retain(|_, stored_state| stored_state.expires_at > now); + + if !inner.entries.contains_key(csrf_token) && inner.entries.len() >= self.capacity { + let oldest_key = inner + .entries + .iter() + .min_by_key(|(_, stored_state)| stored_state.sequence) + .map(|(key, _)| key.clone()); + if let Some(oldest_key) = oldest_key { + inner.entries.remove(&oldest_key); + } + } + + let sequence = inner.next_sequence; + inner.next_sequence = inner.next_sequence.saturating_add(1); + inner.entries.insert( + csrf_token.to_string(), + StoredOAuthState { + state, + expires_at: now + .checked_add(self.ttl) + .expect("OAuth state TTL must fit in Instant"), + sequence, + }, + ); + } + + async fn consume(&self, state: &OAuthState) -> bool { + self.consume_at(state, Instant::now()).await + } + + async fn consume_at(&self, state: &OAuthState, now: Instant) -> bool { + let mut inner = self.inner.lock().await; + inner + .entries + .retain(|_, stored_state| stored_state.expires_at > now); + + let matches = inner + .entries + .get(&state.csrf_token) + .is_some_and(|stored_state| stored_state.state == *state); + if matches { + inner.entries.remove(&state.csrf_token); + } + matches + } + + #[cfg(test)] + async fn len_at(&self, now: Instant) -> usize { + let mut inner = self.inner.lock().await; + inner + .entries + .retain(|_, stored_state| stored_state.expires_at > now); + inner.entries.len() + } +} + #[derive(Debug, Clone)] pub struct GithubProvider { pub auth_url: String, pub token_url: String, pub user_info_url: String, - pub state_store: Arc>>, + state_store: OAuthStateStore, } impl GithubProvider { @@ -36,7 +142,7 @@ impl GithubProvider { auth_url, token_url, user_info_url, - state_store: Arc::new(RwLock::new(HashMap::new())), + state_store: OAuthStateStore::new(), }; // Ensure the provider exists in the database @@ -79,19 +185,11 @@ impl GithubProvider { } pub async fn store_state(&self, csrf_token: &str, state: OAuthState) { - self.state_store - .write() - .await - .insert(csrf_token.to_string(), state); + self.state_store.store(csrf_token, state).await; } - pub async fn validate_state(&self, state: &OAuthState) -> bool { - if let Some(stored_state) = self.state_store.read().await.get(&state.csrf_token) { - // Validate both the CSRF token and the client_id match - stored_state == state - } else { - false - } + pub async fn consume_state(&self, state: &OAuthState) -> bool { + self.state_store.consume(state).await } async fn ensure_provider_exists( @@ -133,7 +231,7 @@ pub struct GoogleProvider { pub auth_url: String, pub token_url: String, pub user_info_url: String, - pub state_store: Arc>>, + state_store: OAuthStateStore, } impl GoogleProvider { @@ -146,7 +244,7 @@ impl GoogleProvider { auth_url, token_url, user_info_url, - state_store: Arc::new(RwLock::new(HashMap::new())), + state_store: OAuthStateStore::new(), }; // Ensure the provider exists in the database @@ -190,19 +288,11 @@ impl GoogleProvider { } pub async fn store_state(&self, csrf_token: &str, state: OAuthState) { - self.state_store - .write() - .await - .insert(csrf_token.to_string(), state); + self.state_store.store(csrf_token, state).await; } - pub async fn validate_state(&self, state: &OAuthState) -> bool { - if let Some(stored_state) = self.state_store.read().await.get(&state.csrf_token) { - // Validate both the CSRF token and the client_id match - stored_state == state - } else { - false - } + pub async fn consume_state(&self, state: &OAuthState) -> bool { + self.state_store.consume(state).await } async fn ensure_provider_exists( @@ -244,7 +334,7 @@ pub struct AppleProvider { pub auth_url: String, pub token_url: String, pub jwks_url: String, - pub state_store: Arc>>, + state_store: OAuthStateStore, // No need for user_info_url as Apple doesn't have a separate endpoint - all info is in the ID token } @@ -258,7 +348,7 @@ impl AppleProvider { auth_url, token_url, jwks_url, - state_store: Arc::new(RwLock::new(HashMap::new())), + state_store: OAuthStateStore::new(), }; // Ensure the provider exists in the database @@ -306,19 +396,11 @@ impl AppleProvider { } pub async fn store_state(&self, csrf_token: &str, state: OAuthState) { - self.state_store - .write() - .await - .insert(csrf_token.to_string(), state); + self.state_store.store(csrf_token, state).await; } - pub async fn validate_state(&self, state: &OAuthState) -> bool { - if let Some(stored_state) = self.state_store.read().await.get(&state.csrf_token) { - // Validate both the CSRF token and the client_id match - stored_state == state - } else { - false - } + pub async fn consume_state(&self, state: &OAuthState) -> bool { + self.state_store.consume(state).await } async fn ensure_provider_exists( @@ -369,7 +451,7 @@ pub trait OAuthProvider: Send + Sync + 'static { async fn generate_authorize_url(&self, client: &BasicClient) -> (String, CsrfToken); async fn store_state(&self, csrf_token: &str, state: OAuthState); - async fn validate_state(&self, state: &OAuthState) -> bool; + async fn consume_state(&self, state: &OAuthState) -> bool; async fn build_client( &self, client_id: String, @@ -412,8 +494,8 @@ impl OAuthProvider for GithubProvider { self.store_state(csrf_token, state).await } - async fn validate_state(&self, state: &OAuthState) -> bool { - self.validate_state(state).await + async fn consume_state(&self, state: &OAuthState) -> bool { + self.consume_state(state).await } async fn build_client( @@ -441,8 +523,8 @@ impl OAuthProvider for GoogleProvider { self.store_state(csrf_token, state).await } - async fn validate_state(&self, state: &OAuthState) -> bool { - self.validate_state(state).await + async fn consume_state(&self, state: &OAuthState) -> bool { + self.consume_state(state).await } async fn build_client( @@ -470,8 +552,8 @@ impl OAuthProvider for AppleProvider { self.store_state(csrf_token, state).await } - async fn validate_state(&self, state: &OAuthState) -> bool { - self.validate_state(state).await + async fn consume_state(&self, state: &OAuthState) -> bool { + self.consume_state(state).await } async fn build_client( @@ -484,3 +566,141 @@ impl OAuthProvider for AppleProvider { .await } } + +#[cfg(test)] +mod tests { + use super::*; + use tokio::sync::Barrier; + + fn state(csrf_token: &str, client_id: u128) -> OAuthState { + OAuthState { + csrf_token: csrf_token.to_string(), + client_id: Uuid::from_u128(client_id), + } + } + + #[tokio::test] + async fn oauth_state_expires_at_ttl() { + let ttl = Duration::from_secs(10); + let store = OAuthStateStore::with_limits(ttl, 2); + let now = Instant::now(); + let oauth_state = state("expired", 1); + + store + .store_at(&oauth_state.csrf_token, oauth_state.clone(), now) + .await; + + assert!(!store.consume_at(&oauth_state, now + ttl).await); + assert_eq!(store.len_at(now + ttl).await, 0); + } + + #[tokio::test] + async fn oauth_state_store_evicts_oldest_entry_at_capacity() { + let store = OAuthStateStore::with_limits(Duration::from_secs(60), 2); + let now = Instant::now(); + let first = state("first", 1); + let second = state("second", 2); + let third = state("third", 3); + + store.store_at(&first.csrf_token, first.clone(), now).await; + store + .store_at(&second.csrf_token, second.clone(), now) + .await; + store.store_at(&third.csrf_token, third.clone(), now).await; + + assert_eq!(store.len_at(now).await, 2); + assert!(!store.consume_at(&first, now).await); + assert!(store.consume_at(&second, now).await); + assert!(store.consume_at(&third, now).await); + } + + #[tokio::test] + async fn oauth_state_store_prunes_expired_entries_before_eviction() { + let ttl = Duration::from_secs(10); + let store = OAuthStateStore::with_limits(ttl, 2); + let now = Instant::now(); + let expired = state("expired", 1); + let live = state("live", 2); + let newest = state("newest", 3); + + store + .store_at(&expired.csrf_token, expired.clone(), now) + .await; + store + .store_at(&live.csrf_token, live.clone(), now + Duration::from_secs(5)) + .await; + store + .store_at(&newest.csrf_token, newest.clone(), now + ttl) + .await; + + assert_eq!(store.len_at(now + ttl).await, 2); + assert!(!store.consume_at(&expired, now + ttl).await); + assert!(store.consume_at(&live, now + ttl).await); + assert!(store.consume_at(&newest, now + ttl).await); + } + + #[tokio::test] + async fn mismatched_oauth_callback_does_not_consume_valid_state() { + let store = OAuthStateStore::with_limits(Duration::from_secs(60), 2); + let now = Instant::now(); + let valid = state("csrf-token", 1); + let mismatched = state("csrf-token", 2); + + store.store_at(&valid.csrf_token, valid.clone(), now).await; + + assert!(!store.consume_at(&mismatched, now).await); + assert!(store.consume_at(&valid, now).await); + } + + #[tokio::test] + async fn valid_oauth_state_can_only_be_consumed_once() { + let store = OAuthStateStore::with_limits(Duration::from_secs(60), 2); + let now = Instant::now(); + let oauth_state = state("one-time", 1); + + store + .store_at(&oauth_state.csrf_token, oauth_state.clone(), now) + .await; + + assert!(store.consume_at(&oauth_state, now).await); + assert!(!store.consume_at(&oauth_state, now).await); + } + + #[tokio::test] + async fn concurrent_oauth_callbacks_have_one_winner() { + let store = OAuthStateStore::with_limits(Duration::from_secs(60), 2); + let now = Instant::now(); + let oauth_state = state("raced", 1); + store + .store_at(&oauth_state.csrf_token, oauth_state.clone(), now) + .await; + + let barrier = Arc::new(Barrier::new(3)); + let first = { + let store = store.clone(); + let barrier = barrier.clone(); + let oauth_state = oauth_state.clone(); + tokio::spawn(async move { + barrier.wait().await; + store.consume_at(&oauth_state, now).await + }) + }; + let second = { + let store = store.clone(); + let barrier = barrier.clone(); + let oauth_state = oauth_state.clone(); + tokio::spawn(async move { + barrier.wait().await; + store.consume_at(&oauth_state, now).await + }) + }; + + barrier.wait().await; + let consumed = [first.await.unwrap(), second.await.unwrap()] + .into_iter() + .filter(|was_consumed| *was_consumed) + .count(); + + assert_eq!(consumed, 1); + } +} diff --git a/src/web/oauth_routes.rs b/src/web/oauth_routes.rs index ef9a8bd8..39e1e0fc 100644 --- a/src/web/oauth_routes.rs +++ b/src/web/oauth_routes.rs @@ -505,8 +505,9 @@ pub async fn oauth_callback( ApiError::InternalServerError })?; - // Validate the complete state (both CSRF token and client_id) - let is_valid = oauth_provider.validate_state(&state).await; + // Validate and atomically consume the complete state so a successful match is one-time. + // A mismatched state is left in the store for the legitimate callback. + let is_valid = oauth_provider.consume_state(&state).await; if !is_valid { error!("Invalid state in {} callback", provider_name); From 6dd1f7fb9a6680930d6847e0bb1c4c4859a731e4 Mon Sep 17 00:00:00 2001 From: Anthony Ronning <101225832+AnthonyRonning@users.noreply.github.com> Date: Tue, 14 Jul 2026 18:44:24 +0000 Subject: [PATCH 2/2] Reject OAuth state when capacity is full --- src/main.rs | 4 ++ src/oauth.rs | 148 +++++++++++++++++++++++++++++----------- src/web/oauth_routes.rs | 3 +- 3 files changed, 114 insertions(+), 41 deletions(-) diff --git a/src/main.rs b/src/main.rs index 86cc5ff6..30142ee5 100644 --- a/src/main.rs +++ b/src/main.rs @@ -327,6 +327,9 @@ pub enum ApiError { #[error("Payload too large")] PayloadTooLarge, + + #[error("Too many requests")] + TooManyRequests, } impl IntoResponse for ApiError { @@ -352,6 +355,7 @@ impl IntoResponse for ApiError { ApiError::NotFound => StatusCode::NOT_FOUND, ApiError::UnprocessableEntity => StatusCode::UNPROCESSABLE_ENTITY, ApiError::PayloadTooLarge => StatusCode::PAYLOAD_TOO_LARGE, + ApiError::TooManyRequests => StatusCode::TOO_MANY_REQUESTS, }; ( status, diff --git a/src/oauth.rs b/src/oauth.rs index 2fcdb566..9586a06f 100644 --- a/src/oauth.rs +++ b/src/oauth.rs @@ -35,16 +35,18 @@ struct OAuthStateStore { #[derive(Debug, Default)] struct OAuthStateStoreInner { entries: HashMap, - next_sequence: u64, } #[derive(Debug)] struct StoredOAuthState { state: OAuthState, expires_at: Instant, - sequence: u64, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +#[error("OAuth state capacity reached")] +pub struct OAuthStateCapacityError; + impl OAuthStateStore { fn new() -> Self { Self::with_limits(OAUTH_STATE_TTL, OAUTH_STATE_CAPACITY) @@ -59,29 +61,29 @@ impl OAuthStateStore { } } - async fn store(&self, csrf_token: &str, state: OAuthState) { - self.store_at(csrf_token, state, Instant::now()).await; + async fn store( + &self, + csrf_token: &str, + state: OAuthState, + ) -> Result<(), OAuthStateCapacityError> { + self.store_at(csrf_token, state, Instant::now()).await } - async fn store_at(&self, csrf_token: &str, state: OAuthState, now: Instant) { + async fn store_at( + &self, + csrf_token: &str, + state: OAuthState, + now: Instant, + ) -> Result<(), OAuthStateCapacityError> { let mut inner = self.inner.lock().await; inner .entries .retain(|_, stored_state| stored_state.expires_at > now); if !inner.entries.contains_key(csrf_token) && inner.entries.len() >= self.capacity { - let oldest_key = inner - .entries - .iter() - .min_by_key(|(_, stored_state)| stored_state.sequence) - .map(|(key, _)| key.clone()); - if let Some(oldest_key) = oldest_key { - inner.entries.remove(&oldest_key); - } + return Err(OAuthStateCapacityError); } - let sequence = inner.next_sequence; - inner.next_sequence = inner.next_sequence.saturating_add(1); inner.entries.insert( csrf_token.to_string(), StoredOAuthState { @@ -89,9 +91,9 @@ impl OAuthStateStore { expires_at: now .checked_add(self.ttl) .expect("OAuth state TTL must fit in Instant"), - sequence, }, ); + Ok(()) } async fn consume(&self, state: &OAuthState) -> bool { @@ -184,8 +186,12 @@ impl GithubProvider { (auth_url.to_string(), csrf_token) } - pub async fn store_state(&self, csrf_token: &str, state: OAuthState) { - self.state_store.store(csrf_token, state).await; + pub async fn store_state( + &self, + csrf_token: &str, + state: OAuthState, + ) -> Result<(), OAuthStateCapacityError> { + self.state_store.store(csrf_token, state).await } pub async fn consume_state(&self, state: &OAuthState) -> bool { @@ -287,8 +293,12 @@ impl GoogleProvider { (auth_url.to_string(), csrf_token) } - pub async fn store_state(&self, csrf_token: &str, state: OAuthState) { - self.state_store.store(csrf_token, state).await; + pub async fn store_state( + &self, + csrf_token: &str, + state: OAuthState, + ) -> Result<(), OAuthStateCapacityError> { + self.state_store.store(csrf_token, state).await } pub async fn consume_state(&self, state: &OAuthState) -> bool { @@ -395,8 +405,12 @@ impl AppleProvider { (auth_url.to_string(), csrf_token) } - pub async fn store_state(&self, csrf_token: &str, state: OAuthState) { - self.state_store.store(csrf_token, state).await; + pub async fn store_state( + &self, + csrf_token: &str, + state: OAuthState, + ) -> Result<(), OAuthStateCapacityError> { + self.state_store.store(csrf_token, state).await } pub async fn consume_state(&self, state: &OAuthState) -> bool { @@ -450,7 +464,11 @@ pub trait OAuthProvider: Send + Sync + 'static { } async fn generate_authorize_url(&self, client: &BasicClient) -> (String, CsrfToken); - async fn store_state(&self, csrf_token: &str, state: OAuthState); + async fn store_state( + &self, + csrf_token: &str, + state: OAuthState, + ) -> Result<(), OAuthStateCapacityError>; async fn consume_state(&self, state: &OAuthState) -> bool; async fn build_client( &self, @@ -490,7 +508,11 @@ impl OAuthProvider for GithubProvider { self.generate_authorize_url(client).await } - async fn store_state(&self, csrf_token: &str, state: OAuthState) { + async fn store_state( + &self, + csrf_token: &str, + state: OAuthState, + ) -> Result<(), OAuthStateCapacityError> { self.store_state(csrf_token, state).await } @@ -519,7 +541,11 @@ impl OAuthProvider for GoogleProvider { self.generate_authorize_url(client).await } - async fn store_state(&self, csrf_token: &str, state: OAuthState) { + async fn store_state( + &self, + csrf_token: &str, + state: OAuthState, + ) -> Result<(), OAuthStateCapacityError> { self.store_state(csrf_token, state).await } @@ -548,7 +574,11 @@ impl OAuthProvider for AppleProvider { self.generate_authorize_url(client).await } - async fn store_state(&self, csrf_token: &str, state: OAuthState) { + async fn store_state( + &self, + csrf_token: &str, + state: OAuthState, + ) -> Result<(), OAuthStateCapacityError> { self.store_state(csrf_token, state).await } @@ -588,34 +618,64 @@ mod tests { store .store_at(&oauth_state.csrf_token, oauth_state.clone(), now) - .await; + .await + .unwrap(); assert!(!store.consume_at(&oauth_state, now + ttl).await); assert_eq!(store.len_at(now + ttl).await, 0); } #[tokio::test] - async fn oauth_state_store_evicts_oldest_entry_at_capacity() { + async fn oauth_state_store_rejects_new_entry_at_capacity() { let store = OAuthStateStore::with_limits(Duration::from_secs(60), 2); let now = Instant::now(); let first = state("first", 1); let second = state("second", 2); let third = state("third", 3); - store.store_at(&first.csrf_token, first.clone(), now).await; + store + .store_at(&first.csrf_token, first.clone(), now) + .await + .unwrap(); store .store_at(&second.csrf_token, second.clone(), now) - .await; - store.store_at(&third.csrf_token, third.clone(), now).await; + .await + .unwrap(); + + assert_eq!( + store.store_at(&third.csrf_token, third.clone(), now).await, + Err(OAuthStateCapacityError) + ); assert_eq!(store.len_at(now).await, 2); - assert!(!store.consume_at(&first, now).await); + assert!(store.consume_at(&first, now).await); assert!(store.consume_at(&second, now).await); - assert!(store.consume_at(&third, now).await); + assert!(!store.consume_at(&third, now).await); } #[tokio::test] - async fn oauth_state_store_prunes_expired_entries_before_eviction() { + async fn oauth_state_store_replaces_existing_entry_at_capacity() { + let store = OAuthStateStore::with_limits(Duration::from_secs(60), 1); + let now = Instant::now(); + let original = state("shared", 1); + let replacement = state("shared", 2); + + store + .store_at(&original.csrf_token, original.clone(), now) + .await + .unwrap(); + store + .store_at(&replacement.csrf_token, replacement.clone(), now) + .await + .unwrap(); + + assert_eq!(store.len_at(now).await, 1); + assert!(!store.consume_at(&original, now).await); + assert!(store.consume_at(&replacement, now).await); + } + + #[tokio::test] + async fn oauth_state_store_prunes_expired_entries_before_capacity_check() { let ttl = Duration::from_secs(10); let store = OAuthStateStore::with_limits(ttl, 2); let now = Instant::now(); @@ -625,13 +685,16 @@ mod tests { store .store_at(&expired.csrf_token, expired.clone(), now) - .await; + .await + .unwrap(); store .store_at(&live.csrf_token, live.clone(), now + Duration::from_secs(5)) - .await; + .await + .unwrap(); store .store_at(&newest.csrf_token, newest.clone(), now + ttl) - .await; + .await + .unwrap(); assert_eq!(store.len_at(now + ttl).await, 2); assert!(!store.consume_at(&expired, now + ttl).await); @@ -646,7 +709,10 @@ mod tests { let valid = state("csrf-token", 1); let mismatched = state("csrf-token", 2); - store.store_at(&valid.csrf_token, valid.clone(), now).await; + store + .store_at(&valid.csrf_token, valid.clone(), now) + .await + .unwrap(); assert!(!store.consume_at(&mismatched, now).await); assert!(store.consume_at(&valid, now).await); @@ -660,7 +726,8 @@ mod tests { store .store_at(&oauth_state.csrf_token, oauth_state.clone(), now) - .await; + .await + .unwrap(); assert!(store.consume_at(&oauth_state, now).await); assert!(!store.consume_at(&oauth_state, now).await); @@ -673,7 +740,8 @@ mod tests { let oauth_state = state("raced", 1); store .store_at(&oauth_state.csrf_token, oauth_state.clone(), now) - .await; + .await + .unwrap(); let barrier = Arc::new(Barrier::new(3)); let first = { diff --git a/src/web/oauth_routes.rs b/src/web/oauth_routes.rs index 39e1e0fc..2e4eda66 100644 --- a/src/web/oauth_routes.rs +++ b/src/web/oauth_routes.rs @@ -455,7 +455,8 @@ pub async fn initiate_oauth( // Store the complete state in the provider oauth_provider .store_state(csrf_token.secret(), state.clone()) - .await; + .await + .map_err(|_| ApiError::TooManyRequests)?; let state_json = serde_json::to_string(&state).map_err(|_| ApiError::InternalServerError)?; let state_base64 = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(state_json);