Something went wrong. Try again.
Microservice to bring 2FA to self hosted PDSes
Something went wrong. Try again.
13 kB · 366 lines
Rust
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367//! One-shot hand over of the gatekeeper's per account 2FA flags to the PDS's own//! `account_email_auth_factor` table, for PDS versions with native email 2FA.
use chrono::{SecondsFormat, Utc};use sqlx::SqlitePool;use tracing::log;
/// Purpose the gatekeeper used for its own login codes in the PDS `email_token` table.pub const LEGACY_TOKEN_PURPOSE: &str = "2fa_code";
#[derive(Debug, Default, PartialEq, Eq)]pub struct MigrationSummary { /// Rows we turned on in the PDS. pub enabled: usize, /// Rows the PDS already had turned on (nothing to do). pub already_enabled: usize, /// Rows that had 2FA off in the gatekeeper (nothing to hand over). pub already_off: usize, /// DIDs with 2FA on, but whose email is not confirmed in the PDS. The PDS /// refuses to enable 2FA on an unconfirmed email, so these users need to /// confirm their email and turn 2FA back on in the app. pub skipped_unconfirmed: Vec<String>, /// DIDs with 2FA on that no longer exist in the PDS. pub skipped_missing: Vec<String>,}
impl MigrationSummary { pub fn total_processed(&self) -> usize { self.enabled + self.already_enabled + self.already_off + self.skipped_unconfirmed.len() + self.skipped_missing.len() }}
/// Moves every not yet migrated `two_factor_accounts` row into the PDS's/// `account_email_auth_factor` table, stamps it as migrated, and clears the/// gatekeeper's old login codes out of the PDS `email_token` table.////// Safe to call on every boot: rows are only ever handed over once.pub async fn migrate_two_factor_to_pds( gatekeeper_pool: &SqlitePool, account_pool: &SqlitePool,) -> anyhow::Result<MigrationSummary> { let pending: Vec<(String, i64)> = sqlx::query_as( "SELECT did, required FROM two_factor_accounts WHERE migrated_to_pds_at IS NULL", ) .fetch_all(gatekeeper_pool) .await?;
let mut summary = MigrationSummary::default();
for (did, required) in pending { if required != 0 { match hand_over_to_pds(account_pool, &did).await? { HandOver::Enabled => summary.enabled += 1, HandOver::AlreadyEnabled => summary.already_enabled += 1, HandOver::Unconfirmed => { log::warn!( "{did} had 2FA turned on in PDS Gatekeeper but its email is not confirmed in the PDS, so the PDS will not let it be enabled. They will need to confirm their email and turn 2FA back on in the app." ); summary.skipped_unconfirmed.push(did.clone()); } HandOver::Missing => { log::warn!( "{did} had 2FA turned on in PDS Gatekeeper but no longer exists in the PDS, skipping." ); summary.skipped_missing.push(did.clone()); } } } else { summary.already_off += 1; }
sqlx::query( "UPDATE two_factor_accounts SET migrated_to_pds_at = CURRENT_TIMESTAMP WHERE did = ?", ) .bind(&did) .execute(gatekeeper_pool) .await?; }
// The PDS issues its own codes now (purpose 'auth_factor'); ours are dead weight. let cleared = sqlx::query("DELETE FROM email_token WHERE purpose = ?") .bind(LEGACY_TOKEN_PURPOSE) .execute(account_pool) .await? .rows_affected();
if summary.total_processed() == 0 { log::info!("2FA hand over to the PDS: nothing to migrate"); } else { log::info!( "2FA hand over to the PDS: enabled {} account(s), {} already enabled in the PDS, {} had it off, skipped {} with unconfirmed email {:?}, skipped {} missing from the PDS {:?}, cleared {cleared} old login code(s)", summary.enabled, summary.already_enabled, summary.already_off, summary.skipped_unconfirmed.len(), summary.skipped_unconfirmed, summary.skipped_missing.len(), summary.skipped_missing, ); }
Ok(summary)}
enum HandOver { Enabled, AlreadyEnabled, Unconfirmed, Missing,}
/// Turns 2FA on for `did` in the PDS with the same guard the PDS itself uses:/// only if the email is confirmed, checked in the same statement as the insert/// so a concurrent email change cannot slip in between.async fn hand_over_to_pds(account_pool: &SqlitePool, did: &str) -> anyhow::Result<HandOver> { let account: Option<(Option<String>,)> = sqlx::query_as("SELECT emailConfirmedAt FROM account WHERE did = ?") .bind(did) .fetch_optional(account_pool) .await?;
match account { None => return Ok(HandOver::Missing), Some((None,)) => return Ok(HandOver::Unconfirmed), Some((Some(_),)) => {} }
// Same shape the PDS writes (Date.toISOString()), e.g. 2026-09-30T11:05:00.000Z let enabled_at = Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true);
let inserted = sqlx::query( "INSERT INTO account_email_auth_factor (did, emailAuthFactorEnabledAt) SELECT did, ? FROM account WHERE did = ? AND emailConfirmedAt IS NOT NULL ON CONFLICT(did) DO NOTHING", ) .bind(&enabled_at) .bind(did) .execute(account_pool) .await? .rows_affected();
if inserted == 1 { Ok(HandOver::Enabled) } else { Ok(HandOver::AlreadyEnabled) }}
#[cfg(test)]mod tests { #![allow(clippy::unwrap_used)] use super::*;
async fn memory_pool() -> SqlitePool { sqlx::sqlite::SqlitePoolOptions::new() .max_connections(1) .connect("sqlite::memory:") .await .expect("in memory sqlite") }
/// Gatekeeper DB with the real embedded migrations applied. async fn gatekeeper_db() -> SqlitePool { let pool = memory_pool().await; sqlx::migrate!("./migrations").run(&pool).await.unwrap(); pool }
/// Minimal copy of the PDS account.sqlite tables we touch (PDS >= 0.4.5036). async fn account_db() -> SqlitePool { let pool = memory_pool().await; for sql in [ "CREATE TABLE account (did TEXT PRIMARY KEY, email TEXT NOT NULL, passwordScrypt TEXT NOT NULL DEFAULT '', emailConfirmedAt TEXT NULL)", "CREATE TABLE email_token (purpose TEXT NOT NULL, did TEXT NOT NULL, token TEXT NOT NULL, requestedAt TEXT NOT NULL, PRIMARY KEY (purpose, did))", "CREATE TABLE account_email_auth_factor (did TEXT PRIMARY KEY, emailAuthFactorEnabledAt TEXT NOT NULL)", ] { sqlx::query(sql).execute(&pool).await.unwrap(); } pool }
async fn add_account(pool: &SqlitePool, did: &str, confirmed: bool) { sqlx::query("INSERT INTO account (did, email, emailConfirmedAt) VALUES (?, ?, ?)") .bind(did) .bind(format!("{did}@example.com")) .bind(confirmed.then_some("2026-01-01T00:00:00.000Z")) .execute(pool) .await .unwrap(); }
async fn add_gatekeeper_row(pool: &SqlitePool, did: &str, required: i32) { sqlx::query("INSERT INTO two_factor_accounts (did, required) VALUES (?, ?)") .bind(did) .bind(required) .execute(pool) .await .unwrap(); }
async fn native_enabled_at(pool: &SqlitePool, did: &str) -> Option<String> { sqlx::query_as::<_, (String,)>( "SELECT emailAuthFactorEnabledAt FROM account_email_auth_factor WHERE did = ?", ) .bind(did) .fetch_optional(pool) .await .unwrap() .map(|r| r.0) }
async fn migrated_at(pool: &SqlitePool, did: &str) -> Option<String> { sqlx::query_as::<_, (Option<String>,)>( "SELECT migrated_to_pds_at FROM two_factor_accounts WHERE did = ?", ) .bind(did) .fetch_one(pool) .await .unwrap() .0 }
#[tokio::test] async fn enables_native_two_factor_for_confirmed_accounts() { let gk = gatekeeper_db().await; let acct = account_db().await; add_account(&acct, "did:plc:alice", true).await; add_gatekeeper_row(&gk, "did:plc:alice", 1).await;
let summary = migrate_two_factor_to_pds(&gk, &acct).await.unwrap();
assert_eq!(summary.enabled, 1); assert_eq!(summary.total_processed(), 1); let enabled_at = native_enabled_at(&acct, "did:plc:alice") .await .expect("row in PDS table"); // Same shape the PDS writes: Date.toISOString(), e.g. 2026-09-30T11:05:00.000Z assert!( chrono::DateTime::parse_from_rfc3339(&enabled_at).is_ok(), "{enabled_at}" ); assert!(enabled_at.ends_with('Z'), "{enabled_at}"); assert!(migrated_at(&gk, "did:plc:alice").await.is_some()); }
#[tokio::test] async fn skips_unconfirmed_emails_but_still_stamps_them() { let gk = gatekeeper_db().await; let acct = account_db().await; add_account(&acct, "did:plc:bob", false).await; add_gatekeeper_row(&gk, "did:plc:bob", 1).await;
let summary = migrate_two_factor_to_pds(&gk, &acct).await.unwrap();
assert_eq!(summary.enabled, 0); assert_eq!(summary.skipped_unconfirmed, vec!["did:plc:bob".to_string()]); assert!(native_enabled_at(&acct, "did:plc:bob").await.is_none()); assert!(migrated_at(&gk, "did:plc:bob").await.is_some()); }
#[tokio::test] async fn skips_accounts_missing_from_the_pds() { let gk = gatekeeper_db().await; let acct = account_db().await; add_gatekeeper_row(&gk, "did:plc:gone", 1).await;
let summary = migrate_two_factor_to_pds(&gk, &acct).await.unwrap();
assert_eq!(summary.skipped_missing, vec!["did:plc:gone".to_string()]); assert!(native_enabled_at(&acct, "did:plc:gone").await.is_none()); assert!(migrated_at(&gk, "did:plc:gone").await.is_some()); }
#[tokio::test] async fn rows_with_two_factor_off_are_only_stamped() { let gk = gatekeeper_db().await; let acct = account_db().await; add_account(&acct, "did:plc:carol", true).await; add_gatekeeper_row(&gk, "did:plc:carol", 0).await;
let summary = migrate_two_factor_to_pds(&gk, &acct).await.unwrap();
assert_eq!(summary.already_off, 1); assert!(native_enabled_at(&acct, "did:plc:carol").await.is_none()); assert!(migrated_at(&gk, "did:plc:carol").await.is_some()); }
#[tokio::test] async fn keeps_the_pds_timestamp_when_already_enabled_natively() { let gk = gatekeeper_db().await; let acct = account_db().await; add_account(&acct, "did:plc:dave", true).await; sqlx::query( "INSERT INTO account_email_auth_factor (did, emailAuthFactorEnabledAt) VALUES (?, ?)", ) .bind("did:plc:dave") .bind("2020-01-01T00:00:00.000Z") .execute(&acct) .await .unwrap(); add_gatekeeper_row(&gk, "did:plc:dave", 1).await;
let summary = migrate_two_factor_to_pds(&gk, &acct).await.unwrap();
assert_eq!(summary.already_enabled, 1); assert_eq!(summary.enabled, 0); assert_eq!( native_enabled_at(&acct, "did:plc:dave").await.as_deref(), Some("2020-01-01T00:00:00.000Z") ); }
#[tokio::test] async fn second_run_is_a_no_op_and_does_not_re_enable() { let gk = gatekeeper_db().await; let acct = account_db().await; add_account(&acct, "did:plc:erin", true).await; add_gatekeeper_row(&gk, "did:plc:erin", 1).await; migrate_two_factor_to_pds(&gk, &acct).await.unwrap();
// User turns 2FA off natively in the PDS after the hand over sqlx::query("DELETE FROM account_email_auth_factor WHERE did = ?") .bind("did:plc:erin") .execute(&acct) .await .unwrap();
let summary = migrate_two_factor_to_pds(&gk, &acct).await.unwrap();
assert_eq!(summary, MigrationSummary::default()); assert!(native_enabled_at(&acct, "did:plc:erin").await.is_none()); }
#[tokio::test] async fn clears_only_the_gatekeepers_old_login_codes() { let gk = gatekeeper_db().await; let acct = account_db().await; for (purpose, did) in [ ("2fa_code", "did:plc:a"), ("2fa_code", "did:plc:b"), ("update_email", "did:plc:a"), ("auth_factor", "did:plc:a"), ] { sqlx::query("INSERT INTO email_token (purpose, did, token, requestedAt) VALUES (?, ?, 'ABCDE-FGHIJ', '2026-01-01T00:00:00.000Z')") .bind(purpose) .bind(did) .execute(&acct) .await .unwrap(); }
migrate_two_factor_to_pds(&gk, &acct).await.unwrap();
let purposes: Vec<(String,)> = sqlx::query_as("SELECT purpose FROM email_token ORDER BY purpose") .fetch_all(&acct) .await .unwrap(); let purposes: Vec<&str> = purposes.iter().map(|p| p.0.as_str()).collect(); assert_eq!(purposes, vec!["auth_factor", "update_email"]); }}