diff --git a/src/auth/oauth_store.rs b/src/auth/oauth_store.rs index 79dffcd..89315e0 100644 --- a/src/auth/oauth_store.rs +++ b/src/auth/oauth_store.rs @@ -4,6 +4,7 @@ use atrium_oauth::store::session::{Session, SessionStore}; use atrium_oauth::store::state::{InternalStateData, StateStore}; use sqlx::AnyPool; use std::fmt; +use std::sync::{Arc, Mutex}; use crate::db::{DatabaseBackend, adapt_sql}; @@ -116,11 +117,27 @@ impl SessionStore for DbSessionStore {} pub struct DbStateStore { pool: AnyPool, backend: DatabaseBackend, + /// Captures the most recently stored state key so callers can associate + /// additional data (e.g., redirect URIs) with the OAuth state. + last_state_key: Arc>>, + /// Serializes authorize() + take_last_state_key() pairs so concurrent + /// logins cannot interleave and swap each other's state keys. + pub authorize_lock: Arc>, } impl DbStateStore { pub fn new(pool: AnyPool, backend: DatabaseBackend) -> Self { - Self { pool, backend } + Self { + pool, + backend, + last_state_key: Arc::new(Mutex::new(None)), + authorize_lock: Arc::new(tokio::sync::Mutex::new(())), + } + } + + /// Returns the state key from the most recent `set()` call, clearing it. + pub fn take_last_state_key(&self) -> Option { + self.last_state_key.lock().unwrap().take() } } @@ -153,6 +170,7 @@ impl Store for DbStateStore { .bind(&json) .execute(&self.pool) .await?; + *self.last_state_key.lock().unwrap() = Some(key); Ok(()) } diff --git a/src/auth/routes.rs b/src/auth/routes.rs index 105f6e7..2312291 100644 --- a/src/auth/routes.rs +++ b/src/auth/routes.rs @@ -41,25 +41,31 @@ async fn login( jar: SignedCookieJar, Query(query): Query, ) -> Result<(SignedCookieJar, Json), AppError> { + tracing::debug!(handle = %query.handle, redirect_uri = ?query.redirect_uri, "login request"); + + // Hold the authorize lock so that authorize() + take_last_state_key() are atomic. + // This prevents concurrent logins from swapping each other's state keys. + let _authorize_guard = state.oauth_state_store.authorize_lock.lock().await; + let url = state .oauth .authorize(&query.handle, Default::default()) .await .map_err(|e| AppError::Internal(format!("OAuth authorize failed: {e}")))?; + // Capture the state key immediately after authorize(). We can't parse it from the URL + // because atrium uses PAR (Pushed Authorization Requests), so the state is embedded + // in the pushed request, not visible in the URL. + let oauth_state = state.oauth_state_store.take_last_state_key(); + + drop(_authorize_guard); + + tracing::debug!(authorize_url = %url, "authorize URL generated"); + // Store the redirect URI in the database, keyed by the OAuth state parameter. // This avoids third-party cookie issues when Pentaract (cross-origin) calls this endpoint. if let Some(redirect_uri) = &query.redirect_uri { - // Extract the state param from the authorize URL query string - let oauth_state = url - .split('?') - .nth(1) - .and_then(|qs| qs.split('&').find_map(|pair| pair.strip_prefix("state="))) - .map(|s| { - urlencoding::decode(s) - .unwrap_or_else(|_| s.into()) - .to_string() - }); + tracing::debug!(oauth_state = ?oauth_state, redirect_uri = %redirect_uri, "storing redirect for state"); if let Some(oauth_state) = oauth_state { let now = now_rfc3339(); @@ -75,6 +81,8 @@ async fn login( .bind(&expires_at) .execute(&state.db) .await; + } else { + tracing::warn!("no state key captured from OAuth authorize — redirect will be lost"); } } @@ -86,6 +94,8 @@ async fn callback( jar: SignedCookieJar, Query(query): Query, ) -> Result<(SignedCookieJar, Redirect), AppError> { + tracing::debug!(state = ?query.state, "callback received"); + // Look up the redirect URI from the database before the OAuth library consumes the state let redirect_url = if let Some(oauth_state) = &query.state { let sql = adapt_sql( @@ -112,8 +122,10 @@ async fn callback( .await; } + tracing::debug!(found_redirect = ?row, "redirect lookup result"); row.map(|(uri,)| uri) } else { + tracing::debug!("no state in callback query"); None }; @@ -137,6 +149,7 @@ async fn callback( // Use DB-stored redirect, or default to "/" let redirect_url = redirect_url.unwrap_or_else(|| "/".to_string()); + tracing::debug!(redirect_url = %redirect_url, "redirecting after callback"); // Set the session cookie // Must use SameSite=None for cross-origin requests (e.g., Pentaract calling HappyView) diff --git a/src/lib.rs b/src/lib.rs index 691aeaf..b69c398 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -58,6 +58,7 @@ pub struct AppState { pub labeler_subscriptions_tx: watch::Sender<()>, pub rate_limiter: Arc, pub oauth: Arc, + pub oauth_state_store: DbStateStore, pub cookie_key: axum_extra::extract::cookie::Key, pub plugin_registry: Arc, pub wasm_runtime: Arc, diff --git a/src/lua/atproto_api.rs b/src/lua/atproto_api.rs index f2e2d3a..4b5a671 100644 --- a/src/lua/atproto_api.rs +++ b/src/lua/atproto_api.rs @@ -267,7 +267,7 @@ mod tests { AppState { config, http: reqwest::Client::new(), - db: test_db, + db: test_db.clone(), db_backend: DatabaseBackend::Sqlite, lexicons: LexiconRegistry::new(), collections_tx: tx, @@ -284,6 +284,10 @@ mod tests { vec![], ), oauth: std::sync::Arc::new(oauth), + oauth_state_store: crate::auth::oauth_store::DbStateStore::new( + test_db.clone(), + crate::db::DatabaseBackend::Sqlite, + ), cookie_key: axum_extra::extract::cookie::Key::derive_from( b"test-secret-for-tests-only-not-production", ), diff --git a/src/lua/db_api.rs b/src/lua/db_api.rs index 037364d..a0614c8 100644 --- a/src/lua/db_api.rs +++ b/src/lua/db_api.rs @@ -683,7 +683,7 @@ mod tests { AppState { config, http: reqwest::Client::new(), - db: test_db, + db: test_db.clone(), db_backend: DatabaseBackend::Sqlite, lexicons: LexiconRegistry::new(), collections_tx: tx, @@ -700,6 +700,10 @@ mod tests { vec![], ), oauth: std::sync::Arc::new(oauth), + oauth_state_store: crate::auth::oauth_store::DbStateStore::new( + test_db.clone(), + crate::db::DatabaseBackend::Sqlite, + ), cookie_key: axum_extra::extract::cookie::Key::derive_from( b"test-secret-for-tests-only-not-production", ), diff --git a/src/lua/execute.rs b/src/lua/execute.rs index 637cc0f..a5e4d69 100644 --- a/src/lua/execute.rs +++ b/src/lua/execute.rs @@ -1009,7 +1009,7 @@ mod tests { AppState { config, http: reqwest::Client::new(), - db: test_db, + db: test_db.clone(), db_backend: DatabaseBackend::Sqlite, lexicons: LexiconRegistry::new(), collections_tx: tx, @@ -1026,6 +1026,10 @@ mod tests { vec![], ), oauth: std::sync::Arc::new(oauth), + oauth_state_store: crate::auth::oauth_store::DbStateStore::new( + test_db.clone(), + crate::db::DatabaseBackend::Sqlite, + ), cookie_key: axum_extra::extract::cookie::Key::derive_from( b"test-secret-for-tests-only-not-production", ), diff --git a/src/lua/http_api.rs b/src/lua/http_api.rs index 32b9e2a..af8c324 100644 --- a/src/lua/http_api.rs +++ b/src/lua/http_api.rs @@ -149,7 +149,7 @@ mod tests { AppState { config, http: reqwest::Client::new(), - db: test_db, + db: test_db.clone(), db_backend: crate::db::DatabaseBackend::Sqlite, lexicons: LexiconRegistry::new(), collections_tx: tx, @@ -166,6 +166,10 @@ mod tests { vec![], ), oauth: std::sync::Arc::new(oauth), + oauth_state_store: crate::auth::oauth_store::DbStateStore::new( + test_db.clone(), + crate::db::DatabaseBackend::Sqlite, + ), cookie_key: axum_extra::extract::cookie::Key::derive_from( b"test-secret-for-tests-only-not-production", ), diff --git a/src/main.rs b/src/main.rs index d170884..c77a621 100644 --- a/src/main.rs +++ b/src/main.rs @@ -282,8 +282,9 @@ async fn main() { http_client: Arc::clone(&atrium_http), }); - let is_loopback = - config.public_url.contains("127.0.0.1") || config.public_url.contains("[::1]"); + let is_loopback = config.public_url.contains("127.0.0.1") + || config.public_url.contains("[::1]") + || config.public_url.contains("localhost"); let resolver_config = OAuthResolverConfig { did_resolver, @@ -292,6 +293,8 @@ async fn main() { protected_resource_metadata: Default::default(), }; + let oauth_state_store = DbStateStore::new(db_pool.clone(), db_backend); + let oauth_client = if is_loopback { info!("Using loopback OAuth client metadata (local development)"); atrium_oauth::OAuthClient::new(OAuthClientConfig { @@ -303,7 +306,7 @@ async fn main() { ]), }, keys: None, - state_store: DbStateStore::new(db_pool.clone(), db_backend), + state_store: oauth_state_store.clone(), session_store: DbSessionStore::new(db_pool.clone(), db_backend), resolver: resolver_config, }) @@ -327,7 +330,7 @@ async fn main() { token_endpoint_auth_signing_alg: None, }, keys: None, - state_store: DbStateStore::new(db_pool.clone(), db_backend), + state_store: oauth_state_store.clone(), session_store: DbSessionStore::new(db_pool.clone(), db_backend), resolver: resolver_config, }) @@ -366,6 +369,7 @@ async fn main() { labeler_subscriptions_tx, rate_limiter, oauth: Arc::new(oauth_client), + oauth_state_store, cookie_key, plugin_registry, wasm_runtime, diff --git a/tests/common/app.rs b/tests/common/app.rs index 9fad490..4033e30 100644 --- a/tests/common/app.rs +++ b/tests/common/app.rs @@ -109,7 +109,7 @@ impl TestApp { let state = AppState { config, http: reqwest::Client::new(), - db: pool, + db: pool.clone(), db_backend: backend, lexicons, collections_tx, @@ -126,6 +126,10 @@ impl TestApp { vec![], ), oauth: std::sync::Arc::new(oauth), + oauth_state_store: happyview::auth::oauth_store::DbStateStore::new( + pool.clone(), + backend, + ), cookie_key: axum_extra::extract::cookie::Key::derive_from( b"test-secret-that-is-at-least-32-bytes-long", ), diff --git a/tests/lua_atproto_api.rs b/tests/lua_atproto_api.rs index 518e6f4..fa0e5c8 100644 --- a/tests/lua_atproto_api.rs +++ b/tests/lua_atproto_api.rs @@ -65,7 +65,7 @@ async fn test_state_with_pool(pool: sqlx::AnyPool, backend: DatabaseBackend) -> AppState { config, http: reqwest::Client::new(), - db: pool, + db: pool.clone(), db_backend: backend, lexicons: LexiconRegistry::new(), collections_tx: tx, @@ -82,6 +82,7 @@ async fn test_state_with_pool(pool: sqlx::AnyPool, backend: DatabaseBackend) -> vec![], ), oauth: std::sync::Arc::new(oauth), + oauth_state_store: happyview::auth::oauth_store::DbStateStore::new(pool.clone(), backend), cookie_key: axum_extra::extract::cookie::Key::derive_from(b"test-secret"), plugin_registry: std::sync::Arc::new(happyview::plugin::PluginRegistry::new()), wasm_runtime: std::sync::Arc::new( diff --git a/tests/lua_db_api.rs b/tests/lua_db_api.rs index f3c461f..98af0ff 100644 --- a/tests/lua_db_api.rs +++ b/tests/lua_db_api.rs @@ -68,7 +68,7 @@ async fn test_state_with_pool(pool: sqlx::AnyPool, backend: DatabaseBackend) -> AppState { config, http: reqwest::Client::new(), - db: pool, + db: pool.clone(), db_backend: backend, lexicons: LexiconRegistry::new(), collections_tx: tx, @@ -85,6 +85,7 @@ async fn test_state_with_pool(pool: sqlx::AnyPool, backend: DatabaseBackend) -> vec![], ), oauth: std::sync::Arc::new(oauth), + oauth_state_store: happyview::auth::oauth_store::DbStateStore::new(pool.clone(), backend), cookie_key: axum_extra::extract::cookie::Key::derive_from(b"test-secret"), plugin_registry: std::sync::Arc::new(happyview::plugin::PluginRegistry::new()), wasm_runtime: std::sync::Arc::new(