vaultwarden/src/db/dsql.rs

use std::{
    collections::HashMap,
    sync::{Arc, LazyLock, Mutex},
    time::Duration,
};

use diesel::ConnectionError;
use url::Url;

struct PsqlUrl {
    timestamp: std::time::Instant,
    url: String,
}

type PsqlUrlLock = Arc<Mutex<Option<PsqlUrl>>>;

static PSQL_URLS: LazyLock<Mutex<HashMap<String, PsqlUrlLock>>> = LazyLock::new(|| Mutex::new(HashMap::new()));
static DSQL_REGION_FROM_HOST_RE: LazyLock<regex::Regex> = LazyLock::new(|| {
    regex::Regex::new(r"^[a-z0-9]+\.dsql\.(?P<region>[a-z0-9-]+)\.on\.aws$")
        .expect("Failed to compile DSQL region regex")
});

// Generate a Postgres libpq connection string. The input connection string has
// the following format:
//
// dsql://<dsql-id>.dsql.<aws-region>.on.aws
//
// The generated connection string has the form:
//
// postgresql://<dsql-id>.dsql.<aws-region>.on.aws/postgres?sslmode=require&user=admin&password=<auth-token>
//
// The auth token is generated by the AWS SDK for DSQL and is valid for up to
// 15 minutes. Cache each unique DSQL URL for 14 minutes to avoid regenerating
// a token for every pooled connection.
pub(crate) fn psql_url(url: &str) -> Result<String, ConnectionError> {
    let mut psql_urls =
        PSQL_URLS.lock().map_err(|e| ConnectionError::BadConnection(format!("Failed to lock DSQL URLs: {e}")))?;

    let psql_url_lock = if let Some(existing_psql_url_lock) = psql_urls.get(url) {
        Arc::clone(existing_psql_url_lock)
    } else {
        let psql_url_lock = Arc::new(Mutex::new(None));
        psql_urls.insert(url.to_owned(), Arc::clone(&psql_url_lock));
        psql_url_lock
    };

    let mut psql_url_lock_guard =
        psql_url_lock.lock().map_err(|e| ConnectionError::BadConnection(format!("Failed to lock DSQL URL: {e}")))?;

    drop(psql_urls);

    if let Some(ref psql_url) = *psql_url_lock_guard {
        if psql_url.timestamp.elapsed() < Duration::from_mins(14) {
            debug!("Reusing DSQL auth token for connection '{url}'");
            return Ok(psql_url.url.clone());
        }

        info!("Refreshing DSQL auth token for connection '{url}'");
    } else {
        info!("Generating new DSQL auth token for connection '{url}'");
    }

    let mut psql_url = Url::parse(url).map_err(|e| ConnectionError::InvalidConnectionUrl(e.to_string()))?;

    let host = psql_url
        .host_str()
        .ok_or_else(|| ConnectionError::InvalidConnectionUrl("Missing hostname in DSQL URL".to_owned()))?
        .to_owned();

    let region = DSQL_REGION_FROM_HOST_RE
        .captures(&host)
        .and_then(|captures| captures.name("region"))
        .ok_or_else(|| ConnectionError::InvalidConnectionUrl("Failed to find AWS region in DSQL hostname".to_owned()))?
        .as_str()
        .to_owned();

    let auth_config = aws_sdk_dsql::auth_token::Config::builder()
        .hostname(host)
        .region(aws_config::Region::new(region))
        .build()
        .map_err(|e| ConnectionError::BadConnection(format!("Failed to build DSQL auth token signer config: {e}")))?;

    let signer = aws_sdk_dsql::auth_token::AuthTokenGenerator::new(auth_config);
    let sdk_config = crate::aws::aws_sdk_config_blocking()
        .map_err(|e| ConnectionError::BadConnection(format!("Failed to load AWS SDK config: {e}")))?;
    let now = std::time::Instant::now();

    let auth_token = std::thread::spawn(move || {
        let rt = tokio::runtime::Builder::new_current_thread().enable_all().build()?;
        rt.block_on(signer.db_connect_admin_auth_token(sdk_config))
    })
    .join()
    .map_err(|e| ConnectionError::BadConnection(format!("Failed to generate DSQL auth token: {e:?}")))?
    .map_err(|e| ConnectionError::BadConnection(format!("Failed to generate DSQL auth token: {e}")))?;

    psql_url.set_scheme("postgresql").expect("Failed to set 'postgresql' as scheme for DSQL connection URL");
    psql_url.set_path("postgres");
    psql_url
        .query_pairs_mut()
        .append_pair("sslmode", "require")
        .append_pair("user", "admin")
        .append_pair("password", auth_token.as_str());

    psql_url_lock_guard.replace(PsqlUrl {
        timestamp: now,
        url: psql_url.to_string(),
    });

    Ok(psql_url.to_string())
}

#[cfg(test)]
mod tests {
    use chrono::Utc;
    use diesel::{Connection, QueryDsl, RunQueryDsl, connection::SimpleConnection};

    use crate::db::schema::{org_policies, users};

    #[test]
    #[ignore = "requires DSQL_MIGRATION_TEST_DATABASE_URL pointing to PostgreSQL"]
    fn release_1_37_5_migrations_preserve_existing_rows() {
        let url = std::env::var("DSQL_MIGRATION_TEST_DATABASE_URL").expect("Missing test database URL");
        let mut conn = diesel::pg::PgConnection::establish(&url).expect("Connect to test database");
        // Temporary tables shadow any real tables and disappear with this connection.
        conn.batch_execute(
            "SET TIME ZONE 'Pacific/Honolulu';
             CREATE TEMP TABLE users (uuid TEXT PRIMARY KEY);
             CREATE TEMP TABLE org_policies (uuid TEXT PRIMARY KEY);
             INSERT INTO users VALUES ('existing-user');
             INSERT INTO org_policies VALUES ('existing-policy-1'), ('existing-policy-2'), ('existing-policy-3');",
        )
        .unwrap();

        conn.batch_execute(include_str!("../../migrations/dsql/2026-09-02-120000_add_key_id/up.sql")).unwrap();
        let keys = users::table.select(users::key_id).load::<Option<String>>(&mut conn).unwrap();
        assert_eq!(keys, vec![None]);

        conn.batch_execute(include_str!("../../migrations/dsql/2026-10-07-120000_add_org_policy_revision_date/up.sql"))
            .unwrap();
        conn.batch_execute(include_str!(
            "../../migrations/dsql/2026-10-07-120050_set_org_policy_revision_default/up.sql"
        ))
        .unwrap();
        conn.batch_execute("INSERT INTO org_policies VALUES ('preserved-policy', TIMESTAMP '2020-01-01 00:00:00');")
            .unwrap();

        let before = Utc::now().naive_utc();
        let backfill = include_str!("../../migrations/dsql/2026-10-07-120100_backfill_org_policy_revision_date/up.sql");
        conn.batch_execute(backfill).unwrap();
        let mut dates = org_policies::table
            .select(org_policies::revision_date)
            .order(org_policies::uuid)
            .load::<chrono::NaiveDateTime>(&mut conn)
            .unwrap();
        let after = Utc::now().naive_utc();
        assert_eq!(dates.len(), 4);
        assert!(dates[..3].iter().all(|date| *date >= before && *date <= after));
        assert_eq!(dates[3], "2020-01-01T00:00:00".parse::<chrono::NaiveDateTime>().unwrap());

        // A still-running old instance must not introduce NULL after the backfill.
        conn.batch_execute("INSERT INTO org_policies (uuid) VALUES ('rolled-back-policy');").unwrap();
        dates.push("1970-01-01T00:00:00".parse::<chrono::NaiveDateTime>().unwrap());
        conn.batch_execute(backfill).unwrap();
        let retried = org_policies::table
            .select(org_policies::revision_date)
            .order(org_policies::uuid)
            .load::<chrono::NaiveDateTime>(&mut conn)
            .unwrap();
        assert_eq!(retried, dates);
    }
}