totp: let legacy sessions disable totp

Lewis: May this revision serve well! <lu5a@proton.me>

crates/tranquil-api/src/server/totp.rs

@@ -6,8 +6,8 @@ use tranquil_pds::api::error::{ApiError, DbResultExt};
 use tranquil_pds::auth::{
     Active, Auth, decrypt_totp_secret, encrypt_totp_secret, generate_backup_codes,
     generate_qr_png_base64, generate_totp_secret, generate_totp_uri, hash_backup_code,
-    is_backup_code_format, require_legacy_session_mfa, verify_backup_code, verify_password_mfa,
-    verify_totp_code, verify_totp_mfa,
+    is_backup_code_format, verify_backup_code, verify_password_mfa, verify_totp_code,
+    verify_totp_mfa,
 };
 use tranquil_pds::rate_limit::{TotpVerifyLimit, check_user_rate_limit_with_message};
 use tranquil_pds::state::AppState;
@@ -163,11 +163,9 @@ pub async fn disable_totp(
     auth: Auth<Active>,
     Json(input): Json<DisableTotpInput>,
 ) -> Result<Json<EmptyResponse>, ApiError> {
-    let session_mfa = require_legacy_session_mfa(&state, &auth).await?;
-
     let _rate_limit = check_user_rate_limit_with_message::<TotpVerifyLimit>(
         &state,
-        session_mfa.did(),
+        &auth.did,
         "Too many verification attempts. Please try again in a few minutes.",
     )
     .await?;
@@ -184,7 +182,7 @@ pub async fn disable_totp(
 
     tranquil_pds::auth::legacy_2fa::clear_challenge(state.cache.as_ref(), &auth.did).await;
 
-    info!(did = %session_mfa.did(), "TOTP disabled (verified via {} and {})", password_mfa.method(), totp_mfa.method());
+    info!(did = %password_mfa.did(), "TOTP disabled (verified via {} and {})", password_mfa.method(), totp_mfa.method());
 
     Ok(Json(EmptyResponse {}))
 }

crates/tranquil-pds/Cargo.toml

@@ -95,6 +95,7 @@ native-tls-roots = ["tranquil-oauth/native-tls-roots"]
 
 [dev-dependencies]
 tempfile = "3"
+totp-rs = { workspace = true }
 ciborium = { workspace = true }
 ctor = { workspace = true }
 testcontainers = { workspace = true }

crates/tranquil-pds/tests/legacy_2fa.rs

@@ -15,6 +15,63 @@ async fn enable_totp_for_user(did: &str) {
         .unwrap();
 }
 
+const KNOWN_TOTP_SECRET: [u8; 20] = [0u8; 20];
+
+async fn enable_totp_encrypted(did: &str) {
+    let encrypted = tranquil_pds::auth::encrypt_totp_secret(&KNOWN_TOTP_SECRET)
+        .expect("encrypt totp secret");
+    let repos = get_test_repos().await;
+    repos
+        .user
+        .enable_totp_verified(&Did::new(did.to_string()).unwrap(), &encrypted)
+        .await
+        .unwrap();
+}
+
+fn current_totp_code() -> String {
+    use totp_rs::{Algorithm, TOTP};
+    let totp = TOTP::new(
+        Algorithm::SHA1,
+        6,
+        1,
+        30,
+        KNOWN_TOTP_SECRET.to_vec(),
+        None,
+        String::new(),
+    )
+    .expect("valid totp params");
+    totp.generate_current().expect("generate totp code")
+}
+
+async fn obtain_legacy_session(client: &reqwest::Client, handle: &str, did: &str) -> String {
+    let base = base_url().await;
+    let resp = client
+        .post(format!("{}/xrpc/com.atproto.server.createSession", base))
+        .json(&json!({ "identifier": handle, "password": "Testpass123!" }))
+        .send()
+        .await
+        .unwrap();
+    assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
+
+    tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+    let code = get_2fa_code_from_queue(did)
+        .await
+        .expect("2FA code should be in queue");
+
+    let resp = client
+        .post(format!("{}/xrpc/com.atproto.server.createSession", base))
+        .json(&json!({ "identifier": handle, "password": "Testpass123!", "authFactorToken": code }))
+        .send()
+        .await
+        .unwrap();
+    assert_eq!(resp.status(), StatusCode::OK);
+    let body: Value = resp.json().await.unwrap();
+    body["accessJwt"]
+        .as_str()
+        .expect("accessJwt in legacy session")
+        .to_string()
+}
+
 async fn set_allow_legacy_login(did: &str, allow: bool) {
     let repos = get_test_repos().await;
     repos
@@ -419,6 +476,60 @@ async fn test_email_auth_factor_requires_code() {
     assert_eq!(body["emailAuthFactor"], true);
 }
 
+#[tokio::test]
+async fn test_disable_totp_allowed_with_legacy_session() {
+    let client = client();
+    let base = base_url().await;
+    let (_token, did) = create_account_and_login(&client).await;
+
+    enable_totp_encrypted(&did).await;
+    set_allow_legacy_login(&did, true).await;
+    clear_2fa_challenges_for_user(&did).await;
+
+    let handle = get_handle(&did).await;
+    let legacy_token = obtain_legacy_session(&client, &handle, &did).await;
+
+    let bad = client
+        .post(format!("{}/xrpc/com.atproto.server.disableTotp", base))
+        .bearer_auth(&legacy_token)
+        .json(&json!({ "password": "nope", "code": "123456" }))
+        .send()
+        .await
+        .unwrap();
+    let bad_body: Value = bad.json().await.unwrap();
+    assert_ne!(
+        bad_body["error"], "MfaVerificationRequired",
+        "legacy session must not be blocked by the MFA gate before credential verification: {:?}",
+        bad_body
+    );
+    assert_eq!(bad_body["error"], "InvalidPassword");
+
+    let ok = client
+        .post(format!("{}/xrpc/com.atproto.server.disableTotp", base))
+        .bearer_auth(&legacy_token)
+        .json(&json!({ "password": "Testpass123!", "code": current_totp_code() }))
+        .send()
+        .await
+        .unwrap();
+    let ok_status = ok.status();
+    let ok_body: Value = ok.json().await.unwrap();
+    assert_eq!(
+        ok_status,
+        StatusCode::OK,
+        "disabling TOTP from a legacy session with valid password and code should succeed: {:?}",
+        ok_body
+    );
+
+    let status = client
+        .get(format!("{}/xrpc/com.atproto.server.getTotpStatus", base))
+        .bearer_auth(&legacy_token)
+        .send()
+        .await
+        .unwrap();
+    let status_body: Value = status.json().await.unwrap();
+    assert_eq!(status_body["enabled"], false);
+}
+
 #[tokio::test]
 async fn test_email_auth_factor_disabled_no_challenge() {
     let client = client();