token create and all that

Cargo.lock

···
 dependencies = [
  "android-tzdata",
  "iana-time-zone",
+ "js-sys",
  "num-traits",
+ "wasm-bindgen",
  "windows-link",
 ]
 
···
  "anyhow",
  "axum",
  "axum-template",
+ "chrono",
  "dotenvy",
  "handlebars",
  "hex",
···
 dependencies = [
  "base64",
  "bytes",
+ "chrono",
  "crc",
  "crossbeam-queue",
  "either",
···
  "bitflags",
  "byteorder",
  "bytes",
+ "chrono",
  "crc",
  "digest",
  "dotenvy",
···
  "base64",
  "bitflags",
  "byteorder",
+ "chrono",
  "crc",
  "dotenvy",
  "etcetera",
···
 checksum = "c2d12fe70b2c1b4401038055f90f151b78208de1f9f89a7dbfd41587a10c3eea"
 dependencies = [
  "atoi",
+ "chrono",
  "flume",
  "futures-channel",
  "futures-core",

Cargo.toml

···
 [dependencies]
 axum = { version = "0.8.4", features = ["macros", "json"] }
 tokio = { version = "1.47.1", features = ["rt-multi-thread", "macros", "signal"] }
-sqlx = { version = "0.8.6", features = ["runtime-tokio-rustls", "sqlite", "migrate"] }
+sqlx = { version = "0.8.6", features = ["runtime-tokio-rustls", "sqlite", "migrate", "chrono"] }
 dotenvy = "0.15.7"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
···
 axum-template = { version = "3.0.0", features = ["handlebars"] }
 rand = "0.9.2"
 anyhow = "1.0.99"
+chrono = "0.4.41"

src/xrpc/com_atproto_server.rs

···
 use crate::AppState;
 use crate::middleware::Did;
-use crate::xrpc::helpers::{ProxiedResult, json_error_response, proxy_get_json};
+use crate::xrpc::helpers::{
+    AuthResult, ProxiedResult, TokenCheckError, json_error_response, preauth_check, proxy_get_json,
+};
 use axum::body::Body;
 use axum::extract::State;
 use axum::http::{HeaderMap, StatusCode};
···
 pub struct CreateSessionRequest {
     identifier: String,
     password: String,
-    auth_factor_token: String,
-    allow_takendown: bool,
-}
-
-pub enum AuthResult {
-    WrongIdentityOrPassword,
-    TwoFactorRequired,
-    TwoFactorFailed,
-    /// User does not have 2FA enabled, or passes it
-    ProxyThrough,
-}
-
-pub enum IdentifierType {
-    Email,
-    DID,
-    Handle,
-}
-
-impl IdentifierType {
-    fn what_is_it(identifier: String) -> Self {
-        if identifier.contains("@") {
-            IdentifierType::Email
-        } else if identifier.contains("did:") {
-            IdentifierType::DID
-        } else {
-            IdentifierType::Handle
-        }
-    }
-}
-
-async fn verify_password(password: &str, password_scrypt: &str) -> Result<bool, StatusCode> {
-    // Expected format: "salt:hash" where hash is hex of scrypt(password, salt, 64 bytes)
-    let mut parts = password_scrypt.splitn(2, ':');
-    let salt = match parts.next() {
-        Some(s) if !s.is_empty() => s,
-        _ => return Ok(false),
-    };
-    let stored_hash_hex = match parts.next() {
-        Some(h) if !h.is_empty() => h,
-        _ => return Ok(false),
-    };
-
-    //Sets up scrypt to mimic node's scrypt
-    let params = match scrypt::Params::new(14, 8, 1, 64) {
-        Ok(p) => p,
-        Err(_) => return Ok(false),
-    };
-    let mut derived = [0u8; 64];
-    if scrypt::scrypt(password.as_bytes(), salt.as_bytes(), &params, &mut derived).is_err() {
-        return Ok(false);
-    }
-
-    let stored_bytes = match hex::decode(stored_hash_hex) {
-        Ok(b) => b,
-        Err(e) => {
-            log::error!("Error decoding stored hash: {}", e);
-            return Ok(false);
-        }
-    };
-
-    Ok(derived.as_slice() == stored_bytes.as_slice())
-}
-
-async fn preauth_check(
-    state: &AppState,
-    identifier: &str,
-    password: &str,
-) -> Result<AuthResult, StatusCode> {
-    // Determine identifier type
-    let id_type = IdentifierType::what_is_it(identifier.to_string());
-
-    // Query account DB for did and passwordScrypt based on identifier type
-    let account_row: Option<(String, String, String)> = match id_type {
-        IdentifierType::Email => sqlx::query_as::<_, (String, String, String)>(
-            "SELECT did, passwordScrypt, account.email FROM account WHERE email = ? LIMIT 1",
-        )
-        .bind(identifier)
-        .fetch_optional(&state.account_pool)
-        .await
-        .map_err(|_| StatusCode::BAD_REQUEST)?,
-        IdentifierType::Handle => sqlx::query_as::<_, (String, String, String)>(
-            "SELECT account.did, account.passwordScrypt, account.email
-                 FROM actor
-                 LEFT JOIN account ON actor.did = account.did
-                 where actor.handle =? LIMIT 1",
-        )
-        .bind(identifier)
-        .fetch_optional(&state.account_pool)
-        .await
-        .map_err(|_| StatusCode::BAD_REQUEST)?,
-        IdentifierType::DID => sqlx::query_as::<_, (String, String, String)>(
-            "SELECT did, passwordScrypt, account.email FROM account WHERE did = ? LIMIT 1",
-        )
-        .bind(identifier)
-        .fetch_optional(&state.account_pool)
-        .await
-        .map_err(|_| StatusCode::BAD_REQUEST)?,
-    };
-
-    if let Some((did, password_scrypt, email)) = account_row {
-        // Check two-factor requirement for this DID in the gatekeeper DB
-        let required_opt = sqlx::query_as::<_, (u8,)>(
-            "SELECT required FROM two_factor_accounts WHERE did = ? LIMIT 1",
-        )
-        .bind(&did)
-        .fetch_optional(&state.pds_gatekeeper_pool)
-        .await
-        .map_err(|_| StatusCode::BAD_REQUEST)?;
-
-        let two_factor_required = match required_opt {
-            Some(row) => row.0 != 0,
-            None => false,
-        };
-
-        if two_factor_required {
-            // Verify password before proceeding to 2FA email step
-            let verified = verify_password(password, &password_scrypt).await?;
-            if !verified {
-                return Ok(AuthResult::WrongIdentityOrPassword);
-            }
-            let mut email_data = Map::new();
-            //TODO these need real values
-            let token = "test".to_string();
-            let handle = "baileytownsend.dev".to_string();
-            email_data.insert("token".to_string(), Value::from(token.clone()));
-            email_data.insert("handle".to_string(), Value::from(handle.clone()));
-            //TODO bad unwrap
-            let email_body = state
-                .template_engine
-                .render("two_factor_code.hbs", email_data)
-                .unwrap();
-
-            let email = Message::builder()
-                //TODO prob get the proper type in the state
-                .from(state.mailer_from.parse().unwrap())
-                .to(email.parse().unwrap())
-                .subject("Sign in to Bluesky")
-                .multipart(
-                    MultiPart::alternative() // This is composed of two parts.
-                        .singlepart(
-                            SinglePart::builder()
-                                .header(header::ContentType::TEXT_PLAIN)
-                                .body(format!("We received a sign-in request for the account @{}. Use the code: {} to sign in. If this wasn't you, we recommend taking steps to protect your account by changing your password at https://bsky.app/settings.", handle, token)), // Every message should have a plain text fallback.
-                        )
-                        .singlepart(
-                            SinglePart::builder()
-                                .header(header::ContentType::TEXT_HTML)
-                                .body(email_body),
-                        ),
-                )
-                //TODO bad
-                .unwrap();
-            return match state.mailer.send(email).await {
-                Ok(_) => Ok(AuthResult::TwoFactorRequired),
-                Err(err) => {
-                    log::error!("Error sending the 2FA email: {}", err);
-                    Err(StatusCode::BAD_REQUEST)
-                }
-            };
-        }
-    }
-
-    // No local 2FA requirement (or account not found)
-    Ok(AuthResult::ProxyThrough)
+    auth_factor_token: Option<String>,
+    allow_takendown: Option<bool>,
 }
 
 pub async fn create_session(
···
 ) -> Result<Response<Body>, StatusCode> {
     let identifier = payload.identifier.clone();
     let password = payload.password.clone();
+    let auth_factor_token = payload.auth_factor_token.clone();
 
     // Run the shared pre-auth logic to validate and check 2FA requirement
-    match preauth_check(&state, &identifier, &password).await? {
+    match preauth_check(&state, &identifier, &password, auth_factor_token).await? {
         AuthResult::WrongIdentityOrPassword => json_error_response(
             StatusCode::UNAUTHORIZED,
             "AuthenticationRequired",
···
 
             Ok(proxied)
         }
+        AuthResult::TokenCheckFailed(err) => match err {
+            TokenCheckError::InvalidToken => json_error_response(
+                StatusCode::BAD_REQUEST,
+                "InvalidToken",
+                "Hey this token is invalid and this is a custom message to show it's bot a normal PDS",
+            ),
+            TokenCheckError::ExpiredToken => {
+                json_error_response(StatusCode::BAD_REQUEST, "ExpiredToken", "Token is expired")
+            }
+        },
     }
 }
 

src/xrpc/helpers.rs

···
 use crate::AppState;
+use crate::xrpc::helpers::TokenCheckError::InvalidToken;
 use axum::body::{Body, to_bytes};
 use axum::extract::Request;
 use axum::http::header::CONTENT_TYPE;
 use axum::http::{HeaderMap, Method, StatusCode, Uri};
 use axum::response::{IntoResponse, Response};
+use axum_template::TemplateEngine;
+use chrono::Utc;
+use lettre::message::{MultiPart, SinglePart, header};
+use lettre::{AsyncTransport, Message};
 use rand::distr::{Alphanumeric, SampleString};
 use serde::de::DeserializeOwned;
+use serde_json::{Map, Value};
 use sqlx::SqlitePool;
-use sqlx::sqlite::SqliteError;
-use tracing::error;
+use tracing::{error, log};
 
 /// The result of a proxied call that attempts to parse JSON.
 pub enum ProxiedResult<T> {
···
     format!("{}-{}", slice_one, slice_two)
 }
 
-pub fn create_two_factor_token(account_db: &SqlitePool, did: String) -> anyhow::Result<String> {}
+pub enum TokenCheckError {
+    InvalidToken,
+    ExpiredToken,
+}
+pub enum AuthResult {
+    WrongIdentityOrPassword,
+    TwoFactorRequired,
+    TwoFactorFailed,
+    /// User does not have 2FA enabled, or passes it
+    ProxyThrough,
+    TokenCheckFailed(TokenCheckError),
+}
+
+pub enum IdentifierType {
+    Email,
+    DID,
+    Handle,
+}
+
+impl IdentifierType {
+    fn what_is_it(identifier: String) -> Self {
+        if identifier.contains("@") {
+            IdentifierType::Email
+        } else if identifier.contains("did:") {
+            IdentifierType::DID
+        } else {
+            IdentifierType::Handle
+        }
+    }
+}
+
+async fn verify_password(password: &str, password_scrypt: &str) -> Result<bool, StatusCode> {
+    // Expected format: "salt:hash" where hash is hex of scrypt(password, salt, 64 bytes)
+    let mut parts = password_scrypt.splitn(2, ':');
+    let salt = match parts.next() {
+        Some(s) if !s.is_empty() => s,
+        _ => return Ok(false),
+    };
+    let stored_hash_hex = match parts.next() {
+        Some(h) if !h.is_empty() => h,
+        _ => return Ok(false),
+    };
+
+    //Sets up scrypt to mimic node's scrypt
+    let params = match scrypt::Params::new(14, 8, 1, 64) {
+        Ok(p) => p,
+        Err(_) => return Ok(false),
+    };
+    let mut derived = [0u8; 64];
+    if scrypt::scrypt(password.as_bytes(), salt.as_bytes(), &params, &mut derived).is_err() {
+        return Ok(false);
+    }
+
+    let stored_bytes = match hex::decode(stored_hash_hex) {
+        Ok(b) => b,
+        Err(e) => {
+            log::error!("Error decoding stored hash: {}", e);
+            return Ok(false);
+        }
+    };
+
+    Ok(derived.as_slice() == stored_bytes.as_slice())
+}
+
+pub async fn preauth_check(
+    state: &AppState,
+    identifier: &str,
+    password: &str,
+    two_factor_code: Option<String>,
+) -> Result<AuthResult, StatusCode> {
+    // Determine identifier type
+    let id_type = IdentifierType::what_is_it(identifier.to_string());
+
+    // Query account DB for did and passwordScrypt based on identifier type
+    let account_row: Option<(String, String, String, String)> = match id_type {
+        IdentifierType::Email => sqlx::query_as::<_, (String, String, String, String)>(
+            "SELECT account.did, account.passwordScrypt, account.email, actor.handle
+                 FROM actor
+                 LEFT JOIN account ON actor.did = account.did
+                 where account.email = ? LIMIT 1",
+        )
+        .bind(identifier)
+        .fetch_optional(&state.account_pool)
+        .await
+        .map_err(|_| StatusCode::BAD_REQUEST)?,
+        IdentifierType::Handle => sqlx::query_as::<_, (String, String, String, String)>(
+            "SELECT account.did, account.passwordScrypt, account.email, actor.handle
+                 FROM actor
+                 LEFT JOIN account ON actor.did = account.did
+                 where actor.handle = ? LIMIT 1",
+        )
+        .bind(identifier)
+        .fetch_optional(&state.account_pool)
+        .await
+        .map_err(|_| StatusCode::BAD_REQUEST)?,
+        IdentifierType::DID => sqlx::query_as::<_, (String, String, String, String)>(
+            "SELECT account.did, account.passwordScrypt, account.email, actor.handle
+                 FROM actor
+                 LEFT JOIN account ON actor.did = account.did
+                 where account.did = ? LIMIT 1",
+        )
+        .bind(identifier)
+        .fetch_optional(&state.account_pool)
+        .await
+        .map_err(|_| StatusCode::BAD_REQUEST)?,
+    };
+
+    if let Some((did, password_scrypt, email, handle)) = account_row {
+        // Check two-factor requirement for this DID in the gatekeeper DB
+        let required_opt = sqlx::query_as::<_, (u8,)>(
+            "SELECT required FROM two_factor_accounts WHERE did = ? LIMIT 1",
+        )
+        .bind(&did)
+        .fetch_optional(&state.pds_gatekeeper_pool)
+        .await
+        .map_err(|_| StatusCode::BAD_REQUEST)?;
+
+        let two_factor_required = match required_opt {
+            Some(row) => row.0 != 0,
+            None => false,
+        };
+
+        if two_factor_required {
+            // Verify password before proceeding to 2FA email step
+            let verified = verify_password(password, &password_scrypt).await?;
+            if !verified {
+                return Ok(AuthResult::WrongIdentityOrPassword);
+            }
+            //Two factor is required and a taken was provided
+            if let Some(two_factor_code) = two_factor_code {
+                //Seems it sends over a empty on login without it set? As in no input is shown on the ui for first login try
+                if two_factor_code != "" {
+                    return match assert_valid_token(
+                        &state.account_pool,
+                        did.clone(),
+                        two_factor_code,
+                    )
+                    .await
+                    {
+                        Ok(_) => {
+                            let _ = delete_all_email_tokens(&state.account_pool, did.clone()).await;
+                            Ok(AuthResult::ProxyThrough)
+                        }
+                        Err(err) => Ok(AuthResult::TokenCheckFailed(err)),
+                    };
+                }
+            }
+
+            return match create_two_factor_token(&state.account_pool, did).await {
+                Ok(code) => {
+                    let mut email_data = Map::new();
+                    email_data.insert("token".to_string(), Value::from(code.clone()));
+                    email_data.insert("handle".to_string(), Value::from(handle.clone()));
+                    //TODO bad unwrap
+                    let email_body = state
+                        .template_engine
+                        .render("two_factor_code.hbs", email_data)
+                        .unwrap();
+
+                    let email = Message::builder()
+                        //TODO prob get the proper type in the state
+                        .from(state.mailer_from.parse().unwrap())
+                        .to(email.parse().unwrap())
+                        .subject("Sign in to Bluesky")
+                        .multipart(
+                            MultiPart::alternative() // This is composed of two parts.
+                                .singlepart(
+                                    SinglePart::builder()
+                                        .header(header::ContentType::TEXT_PLAIN)
+                                        .body(format!("We received a sign-in request for the account @{}. Use the code: {} to sign in. If this wasn't you, we recommend taking steps to protect your account by changing your password at https://bsky.app/settings.", handle, code)), // Every message should have a plain text fallback.
+                                )
+                                .singlepart(
+                                    SinglePart::builder()
+                                        .header(header::ContentType::TEXT_HTML)
+                                        .body(email_body),
+                                ),
+                        )
+                        //TODO bad
+                        .unwrap();
+                    match state.mailer.send(email).await {
+                        Ok(_) => Ok(AuthResult::TwoFactorRequired),
+                        Err(err) => {
+                            log::error!("Error sending the 2FA email: {}", err);
+                            Err(StatusCode::BAD_REQUEST)
+                        }
+                    }
+                }
+                Err(err) => {
+                    log::error!("error on creating a 2fa token: {}", err);
+                    Err(StatusCode::BAD_REQUEST)
+                }
+            };
+        }
+    }
+
+    // No local 2FA requirement (or account not found)
+    Ok(AuthResult::ProxyThrough)
+}
+
+pub async fn create_two_factor_token(
+    account_db: &SqlitePool,
+    did: String,
+) -> anyhow::Result<String> {
+    let purpose = "2fa_code";
+
+    loop {
+        let token = get_random_token();
+        let right_now = Utc::now();
+        let query = "INSERT INTO email_token (purpose, did, token, requestedAt)
+                     VALUES (?, ?, ?, ?)
+                     ON CONFLICT(purpose, did) DO UPDATE SET
+                       token=excluded.token,
+                       requestedAt=excluded.requestedAt";
+
+        let res = sqlx::query(query)
+            .bind(purpose)
+            .bind(&did)
+            .bind(&token)
+            .bind(right_now)
+            .execute(account_db)
+            .await;
+
+        return match res {
+            Ok(_) => Ok(token),
+            Err(e) => {
+                log::error!("Error creating a two factor token: {}", e);
+                Err(anyhow::anyhow!(e))
+            }
+        };
+    }
+}
+
+pub async fn delete_all_email_tokens(account_db: &SqlitePool, did: String) -> anyhow::Result<()> {
+    sqlx::query("DELETE FROM email_token WHERE did = ?")
+        .bind(did)
+        .execute(account_db)
+        .await?;
+    Ok(())
+}
+
+pub async fn assert_valid_token(
+    account_db: &SqlitePool,
+    did: String,
+    token: String,
+) -> Result<(), TokenCheckError> {
+    let token_upper = token.to_ascii_uppercase();
+    let purpose = "2fa_code";
+
+    let row: Option<(String,)> = sqlx::query_as(
+        "SELECT requestedAt FROM email_token WHERE purpose = ? AND did = ? AND token = ? LIMIT 1",
+    )
+    .bind(purpose)
+    .bind(did)
+    .bind(token_upper)
+    .fetch_optional(account_db)
+    .await
+    .map_err(|err| {
+        log::error!("Error getting the 2fa token: {}", err);
+        TokenCheckError::InvalidToken
+    })?;
+
+    match row {
+        None => Err(TokenCheckError::InvalidToken),
+        Some(row) => {
+            // Token lives for 15 minutes
+            let expiration_ms = 15 * 60_000;
+
+            // Parse requestedAt; assume RFC3339-like string (as created_by PDS or by our code)
+            let requested_at_utc = match chrono::DateTime::parse_from_rfc3339(&row.0) {
+                Ok(dt) => dt.with_timezone(&Utc),
+                Err(_) => {
+                    return Err(TokenCheckError::InvalidToken);
+                }
+            };
+
+            let now = Utc::now();
+            let age_ms = (now - requested_at_utc).num_milliseconds();
+            let expired = age_ms > expiration_ms;
+            if expired {
+                return Err(TokenCheckError::ExpiredToken);
+            }
+
+            Ok(())
+        }
+    }
+}