WIP

Cargo.lock

···
 checksum = "778e2ac28f6c47af28e4907f13ffd1e1ddbd400980a9abd7c8df189bf578a5ad"
 dependencies = [
  "libc",
- "windows-sys 0.52.0",
+ "windows-sys 0.59.0",
 ]
 
 [[package]]
···
 checksum = "07033963ba89ebaf1584d767badaa2e8fcec21aedea6b8c0346d487d49c28667"
 dependencies = [
  "cfg-if",
- "windows-targets 0.48.5",
+ "windows-targets 0.52.6",
 ]
 
 [[package]]
···
  "errno",
  "libc",
  "linux-raw-sys",
- "windows-sys 0.52.0",
+ "windows-sys 0.59.0",
 ]
 
 [[package]]

src/main.rs

···
 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;
···
         .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");
+        .expect(
+            "failed to create governor config for sign in. this should not happen and is a bug",
+        );
+
+    let create_account_limiter_time: Option<String> =
+        env::var("GATEKEEPER_CREATE_ACCOUNT_PER_SECOND").ok();
+    let create_account_limiter_burst: Option<String> =
+        env::var("GATEKEEPER_CREATE_ACCOUNT_BURST").ok();
+    let mut create_account_governor_conf = None;
 
-    // let create_account_limiter_time: Option<String> =
-    //     env::var("GATEKEEPER_CREATE_ACCOUNT_LIMITER_WINDOW").unwrap_or_else(|_| 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::<u64>()
+            .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::<u32>()
+            .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 || {
···
             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();
+            }
         }
     });
 
···
 
     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)),
···
             "/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);

src/middleware.rs

···
                 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<Claims<TokenClaims>, 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<Claims<TokenClaims>, 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<Token<TokenClaims>, 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<Token<TokenClaims>, 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
                 }

src/xrpc/com_atproto_server.rs

···
     //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)?;
+    //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();
 
-        //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'",
+    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!(
         "{}{}",