diff --git a/Cargo.lock b/Cargo.lock index 0f2cee2..0ba727b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -656,7 +656,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "778e2ac28f6c47af28e4907f13ffd1e1ddbd400980a9abd7c8df189bf578a5ad" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -1392,7 +1392,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07033963ba89ebaf1584d767badaa2e8fcec21aedea6b8c0346d487d49c28667" dependencies = [ "cfg-if", - "windows-targets 0.48.5", + "windows-targets 0.52.6", ] [[package]] @@ -2136,7 +2136,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] diff --git a/src/main.rs b/src/main.rs index abf3306..42fcb1a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -19,7 +19,8 @@ use std::path::Path; use std::time::Duration; use std::{env, net::SocketAddr}; use tower_governor::GovernorLayer; -use tower_governor::governor::GovernorConfigBuilder; +use tower_governor::governor::{GovernorConfig, GovernorConfigBuilder}; +use tower_governor::key_extractor::PeerIpKeyExtractor; use tower_http::compression::CompressionLayer; use tower_http::cors::{Any, CorsLayer}; use tracing::log; @@ -166,20 +167,49 @@ async fn main() -> Result<(), Box> { .per_second(60) .burst_size(5) .finish() - .expect("failed to create governor config. this should not happen and is a bug"); + .expect("failed to create governor config for create session. this should not happen and is a bug"); // Create a second config with the same settings for the other endpoint let sign_in_governor_conf = GovernorConfigBuilder::default() .per_second(60) .burst_size(5) .finish() - .expect("failed to create governor config. this should not happen and is a bug"); - - // let create_account_limiter_time: Option = - // env::var("GATEKEEPER_CREATE_ACCOUNT_LIMITER_WINDOW").unwrap_or_else(|_| None); + .expect( + "failed to create governor config for sign in. this should not happen and is a bug", + ); + + let create_account_limiter_time: Option = + env::var("GATEKEEPER_CREATE_ACCOUNT_PER_SECOND").ok(); + let create_account_limiter_burst: Option = + env::var("GATEKEEPER_CREATE_ACCOUNT_BURST").ok(); + let mut create_account_governor_conf = None; + + if create_account_governor_conf.is_some() && create_account_limiter_time.is_some() { + let time = create_account_limiter_time + .expect("GATEKEEPER_CREATE_ACCOUNT_PER_SECOND not set") + .parse::() + .expect("GATEKEEPER_CREATE_ACCOUNT_PER_SECOND must be a valid integer"); + let burst = create_account_limiter_burst + .expect("GATEKEEPER_CREATE_ACCOUNT_BURST not set") + .parse::() + .expect("GATEKEEPER_CREATE_ACCOUNT_BURST must be a valid integer"); + + create_account_governor_conf = Some( + GovernorConfigBuilder::default() + .per_second(time) + .burst_size(burst) + .finish() + .expect("failed to create governor config for create account. this should not happen and is a bug"), + ) + } let create_session_governor_limiter = create_session_governor_conf.limiter().clone(); let sign_in_governor_limiter = sign_in_governor_conf.limiter().clone(); + let create_account_governor_limiter = match create_account_governor_conf { + None => None, + Some(conf) => Some(conf.limiter().clone()), + }; + let interval = Duration::from_secs(60); // a separate background task to clean up std::thread::spawn(move || { @@ -187,6 +217,9 @@ async fn main() -> Result<(), Box> { std::thread::sleep(interval); create_session_governor_limiter.retain_recent(); sign_in_governor_limiter.retain_recent(); + if let Some(ref limiter) = create_account_governor_limiter { + limiter.retain_recent(); + } } }); @@ -197,10 +230,7 @@ async fn main() -> Result<(), Box> { let app = Router::new() .route("/", get(root_handler)) - .route( - "/xrpc/com.atproto.server.getSession", - get(get_session).layer(ax_middleware::from_fn(middleware::extract_did)), - ) + .route("/xrpc/com.atproto.server.getSession", get(get_session)) .route( "/xrpc/com.atproto.server.updateEmail", post(update_email).layer(ax_middleware::from_fn(middleware::extract_did)), @@ -213,6 +243,7 @@ async fn main() -> Result<(), Box> { "/xrpc/com.atproto.server.createSession", post(create_session.layer(GovernorLayer::new(create_session_governor_conf))), ) + .route("/xrpc/com.atproto.server.createAccount") .layer(CompressionLayer::new()) .layer(cors) .with_state(state); diff --git a/src/middleware.rs b/src/middleware.rs index 3025c38..4b8044a 100644 --- a/src/middleware.rs +++ b/src/middleware.rs @@ -35,40 +35,53 @@ pub async fn extract_did(mut req: Request, next: Next) -> impl IntoResponse { Some((scheme, token_str)) => { // For Bearer, validate JWT and extract DID from `sub`. // For DPoP, we currently only pass through and do not validate here; insert None DID. - // match scheme { - // AuthScheme::Bearer => { - let token = UntrustedToken::new(&token_str); - if token.is_err() { - return json_error_response(StatusCode::BAD_REQUEST, "TokenRequired", "") - .expect("Error creating an error response"); - } - let parsed_token = token.expect("Already checked for error"); - let claims: Result, ValidationError> = - parsed_token.deserialize_claims_unchecked(); - if claims.is_err() { - return json_error_response(StatusCode::BAD_REQUEST, "TokenRequired", "") - .expect("Error creating an error response"); - } + match scheme { + AuthScheme::Bearer => { + let token = UntrustedToken::new(&token_str); + if token.is_err() { + return json_error_response( + StatusCode::BAD_REQUEST, + "TokenRequired", + "", + ) + .expect("Error creating an error response"); + } + let parsed_token = token.expect("Already checked for error"); + let claims: Result, ValidationError> = + parsed_token.deserialize_claims_unchecked(); + if claims.is_err() { + return json_error_response( + StatusCode::BAD_REQUEST, + "TokenRequired", + "", + ) + .expect("Error creating an error response"); + } - let key = Hs256Key::new( - env::var("PDS_JWT_SECRET").expect("PDS_JWT_SECRET not set in the pds.env"), - ); - let token: Result, ValidationError> = - Hs256.validator(&key).validate(&parsed_token); - if token.is_err() { - return json_error_response(StatusCode::BAD_REQUEST, "InvalidToken", "") - .expect("Error creating an error response"); + let key = Hs256Key::new( + env::var("PDS_JWT_SECRET") + .expect("PDS_JWT_SECRET not set in the pds.env"), + ); + let token: Result, ValidationError> = + Hs256.validator(&key).validate(&parsed_token); + if token.is_err() { + return json_error_response( + StatusCode::BAD_REQUEST, + "InvalidToken", + "", + ) + .expect("Error creating an error response"); + } + let token = token.expect("Already checked for error,"); + // Not going to worry about expiration since it still goes to the PDS + req.extensions_mut() + .insert(Did(Some(token.claims().custom.sub.clone()))); + } + AuthScheme::DPoP => { + //Not going to worry about oauth email update for now, just always forward to the PDS + req.extensions_mut().insert(Did(None)); + } } - let token = token.expect("Already checked for error,"); - // Not going to worry about expiration since it still goes to the PDS - req.extensions_mut() - .insert(Did(Some(token.claims().custom.sub.clone()))); - // } - // AuthScheme::DPoP => { - // // No DID extraction from DPoP here; leave None - // req.extensions_mut().insert(Did(None)); - // } - // } next.run(req).await } diff --git a/src/xrpc/com_atproto_server.rs b/src/xrpc/com_atproto_server.rs index 0f40de4..b05676f 100644 --- a/src/xrpc/com_atproto_server.rs +++ b/src/xrpc/com_atproto_server.rs @@ -147,63 +147,68 @@ pub async fn update_email( //If email auth is set it is to either turn on or off 2fa let email_auth_update = payload.email_auth_factor.unwrap_or(false); - // Email update asked for - if email_auth_update { - let email = payload.email.clone(); - let email_confirmed = sqlx::query_as::<_, (String,)>( - "SELECT did FROM account WHERE emailConfirmedAt IS NOT NULL AND email = ?", - ) - .bind(&email) - .fetch_optional(&state.account_pool) - .await - .map_err(|_| StatusCode::BAD_REQUEST)?; - - //Since the email is already confirmed we can enable 2fa - return match email_confirmed { - None => Err(StatusCode::BAD_REQUEST), - Some(did_row) => { - let _ = sqlx::query( - "INSERT INTO two_factor_accounts (did, required) VALUES (?, 1) ON CONFLICT(did) DO UPDATE SET required = 1", - ) - .bind(&did_row.0) - .execute(&state.pds_gatekeeper_pool) - .await - .map_err(|_| StatusCode::BAD_REQUEST)?; + //This means the middleware successfully extracted a did from the request, if not it just needs to be forward to the PDS + //This is also empty if it is an oauth request, which is not supported by gatekeeper turning on 2fa since the dpop stuff needs to be implemented + let did_is_not_empty = did.0.is_some(); - Ok(StatusCode::OK.into_response()) - } - }; - } - - // User wants auth turned off - if !email_auth_update && !email_auth_not_set { - //User wants auth turned off and has a token - if let Some(token) = &payload.token { - let token_found = sqlx::query_as::<_, (String,)>( - "SELECT token FROM email_token WHERE token = ? AND did = ? AND purpose = 'update_email'", + if did_is_not_empty { + // Email update asked for + if email_auth_update { + let email = payload.email.clone(); + let email_confirmed = sqlx::query_as::<_, (String,)>( + "SELECT did FROM account WHERE emailConfirmedAt IS NOT NULL AND email = ?", ) - .bind(token) - .bind(&did.0) + .bind(&email) .fetch_optional(&state.account_pool) .await .map_err(|_| StatusCode::BAD_REQUEST)?; - if token_found.is_some() { - let _ = sqlx::query( - "INSERT INTO two_factor_accounts (did, required) VALUES (?, 0) ON CONFLICT(did) DO UPDATE SET required = 0", + //Since the email is already confirmed we can enable 2fa + return match email_confirmed { + None => Err(StatusCode::BAD_REQUEST), + Some(did_row) => { + let _ = sqlx::query( + "INSERT INTO two_factor_accounts (did, required) VALUES (?, 1) ON CONFLICT(did) DO UPDATE SET required = 1", + ) + .bind(&did_row.0) + .execute(&state.pds_gatekeeper_pool) + .await + .map_err(|_| StatusCode::BAD_REQUEST)?; + + Ok(StatusCode::OK.into_response()) + } + }; + } + + // User wants auth turned off + if !email_auth_update && !email_auth_not_set { + //User wants auth turned off and has a token + if let Some(token) = &payload.token { + let token_found = sqlx::query_as::<_, (String,)>( + "SELECT token FROM email_token WHERE token = ? AND did = ? AND purpose = 'update_email'", ) - .bind(&did.0) - .execute(&state.pds_gatekeeper_pool) - .await - .map_err(|_| StatusCode::BAD_REQUEST)?; + .bind(token) + .bind(&did.0) + .fetch_optional(&state.account_pool) + .await + .map_err(|_| StatusCode::BAD_REQUEST)?; - return Ok(StatusCode::OK.into_response()); - } else { - return Err(StatusCode::BAD_REQUEST); + return if token_found.is_some() { + let _ = sqlx::query( + "INSERT INTO two_factor_accounts (did, required) VALUES (?, 0) ON CONFLICT(did) DO UPDATE SET required = 0", + ) + .bind(&did.0) + .execute(&state.pds_gatekeeper_pool) + .await + .map_err(|_| StatusCode::BAD_REQUEST)?; + + Ok(StatusCode::OK.into_response()) + } else { + Err(StatusCode::BAD_REQUEST) + }; } } } - // Updating the actual email address by sending it on to the PDS let uri = format!( "{}{}",