server/

provision.rs

//! Per-user database provisioning.
//!
//! Each user owns a dedicated Postgres database, reached via `userpool` from
//! the `users.db_name` DSN. This module creates that database, brings it up to
//! the (DDL-only) schema with the shared migration set, and returns its DSN so
//! the caller can persist it in `users.db_name`.
//!
//! Provisioning reuses the admin `DATABASE_URL` connection — Postgres allows a
//! connection to create a *different* database — so no separate maintenance DSN
//! is required. The only prerequisite is that the admin role holds `CREATEDB`.

use crate::db::{DBError, admin_database_url, execute_raw, migrate_url};
use rust_i18n::t;
use sqlx::types::Uuid;

/// Postgres SQLSTATE for `duplicate_database`.
const DUPLICATE_DATABASE: &str = "42P04";

/// Creates and migrates a dedicated database for `user_id`, returning its DSN.
///
/// Idempotent on the `CREATE DATABASE` step: a pre-existing database (SQLSTATE
/// `42P04`) is treated as success so a retried provision after a mid-way
/// failure reuses the same deterministically-named database rather than
/// erroring. The returned DSN is what the caller stores in `users.db_name`.
///
/// # Errors
///
/// - [`DBError::MissingUrl`] if `DATABASE_URL` is unset.
/// - [`DBError::NoCreateDb`] if the role lacks `CREATEDB`.
/// - [`DBError::Sqlx`] / [`DBError::Migration`] on other failures.
pub async fn create_user_database(user_id: Uuid) -> Result<String, DBError> {
    let base = admin_database_url()?;
    let db_name = database_name(user_id);
    let dsn = swap_database(&base, &db_name)?;

    let created = create_database(&db_name).await?;

    // If migration fails after *this call* created the database, drop it so a
    // transient failure doesn't leak an orphaned (un-migrated) per-user
    // database. We must NOT drop when we reused a pre-existing database
    // (42P04): a retry whose migration fails would otherwise destroy data we
    // didn't create.
    if let Err(err) = migrate_url(&dsn).await {
        if created && let Err(drop_err) = drop_user_database(user_id).await {
            log::error!(
                "{}",
                t!("Failed to drop partially-provisioned database after migration error: %{e}",
                       e = drop_err : {})
            );
        }
        return Err(err);
    }

    Ok(dsn)
}

/// Stores a user's RSA private signing key into their (already-provisioned)
/// per-user database at `dsn`. The private key never leaves that database.
///
/// # Errors
/// [`DBError::Sqlx`] on connect or insert failure.
pub async fn store_user_private_key(dsn: &str, private_pem_b64: &str) -> Result<(), DBError> {
    let pool = sqlx::postgres::PgPoolOptions::new()
        .max_connections(1)
        .acquire_timeout(std::time::Duration::from_secs(10))
        .connect(dsn)
        .await?;
    // Upsert on the singleton constraint: the per-user DB holds exactly one
    // signing key for its owner, enforced at the DB level (see 0005), so
    // concurrent calls collide on `singleton` rather than creating duplicates.
    let result = sqlx::query(
        "INSERT INTO user_auth_keys (singleton, private_key) VALUES (TRUE, $1) \
         ON CONFLICT (singleton) DO UPDATE SET private_key = excluded.private_key",
    )
    .bind(private_pem_b64)
    .execute(&pool)
    .await;
    pool.close().await;
    result.map(|_| ()).map_err(DBError::Sqlx)
}

/// Drops a user's database, used as compensating cleanup when a step *after*
/// provisioning (e.g. the `users` row insert) fails and would otherwise orphan
/// the just-created database. Best-effort and idempotent (`IF EXISTS`).
///
/// # Errors
/// [`DBError::Sqlx`] if the `DROP DATABASE` itself errors (e.g. the database is
/// still in use); callers typically log and continue rather than failing the
/// original error path.
pub async fn drop_user_database(user_id: Uuid) -> Result<(), DBError> {
    let db_name = database_name(user_id);
    // `db_name` is `nomi_user_<hex>` (only `[a-z0-9_]`), so the quoted
    // identifier is injection-safe.
    execute_raw(&format!("DROP DATABASE IF EXISTS \"{db_name}\"")).await
}

/// Deterministic, injection-safe database name for a user. Derived purely from
/// the UUID's hex (no hyphens), so a retry targets the same database.
fn database_name(user_id: Uuid) -> String {
    format!("nomi_user_{}", user_id.simple())
}

/// Postgres SQLSTATE for `insufficient_privilege`.
const INSUFFICIENT_PRIVILEGE: &str = "42501";

/// Issues `CREATE DATABASE` on the admin connection. Returns `true` if this call
/// created the database, `false` if it already existed (SQLSTATE `42P04`, the
/// retry-safe reuse path — the caller still migrates, which is idempotent). An
/// insufficient-privilege error maps to [`DBError::NoCreateDb`]. The boolean
/// lets the caller compensate-drop ONLY databases it actually created.
async fn create_database(db_name: &str) -> Result<bool, DBError> {
    // `db_name` is `nomi_user_<hex>` — only `[a-z0-9_]`, so the quoted
    // identifier cannot be abused. CREATE DATABASE forbids bind parameters.
    let stmt = format!("CREATE DATABASE \"{db_name}\"");
    let Err(DBError::Sqlx(err)) = execute_raw(&stmt).await else {
        return Ok(true);
    };
    match err.as_database_error().and_then(|db| db.code()).as_deref() {
        Some(DUPLICATE_DATABASE) => {
            log::debug!("{}", t!("Per-user database already exists, reusing"));
            Ok(false)
        }
        Some(INSUFFICIENT_PRIVILEGE) => Err(DBError::NoCreateDb),
        _ => Err(DBError::Sqlx(err)),
    }
}

/// Replaces the database (path) component of a Postgres DSN, preserving scheme,
/// credentials, host, port, and query parameters.
///
/// Parses `scheme://authority[/path][?query]`. The authority ends at the first
/// `/` or `?` (whichever comes first), so a DSN with no path but a query
/// (`postgres://user@host?sslmode=require`) is handled correctly — the query is
/// preserved and the new db name is inserted as the path.
fn swap_database(base: &str, db_name: &str) -> Result<String, DBError> {
    let (scheme, rest) = base.split_once("://").ok_or(DBError::MissingUrl)?;

    // Split off the query first so a query-only DSN (no path) doesn't fold the
    // query into the authority.
    let (before_query, query) = match rest.split_once('?') {
        Some((head, q)) => (head, format!("?{q}")),
        None => (rest, String::new()),
    };

    // Whatever path the original DSN had is dropped; the authority is the part
    // before the first '/'.
    let authority = before_query.split('/').next().unwrap_or(before_query);

    Ok(format!("{scheme}://{authority}/{db_name}{query}"))
}

#[cfg(test)]
mod tests {
    use super::swap_database;

    #[test]
    fn swap_database_replaces_path_component() {
        let out = swap_database("postgres://ray@localhost/maindb", "nomi_user_abc").unwrap();
        assert_eq!(out, "postgres://ray@localhost/nomi_user_abc");
    }

    #[test]
    fn swap_database_preserves_credentials_port_and_query() {
        let out = swap_database(
            "postgres://u:p@host:5432/maindb?sslmode=require",
            "nomi_user_abc",
        )
        .unwrap();
        assert_eq!(
            out,
            "postgres://u:p@host:5432/nomi_user_abc?sslmode=require"
        );
    }

    #[test]
    fn swap_database_handles_dsn_without_path_but_with_query() {
        let out = swap_database("postgres://user@host?sslmode=require", "nomi_user_abc").unwrap();
        assert_eq!(out, "postgres://user@host/nomi_user_abc?sslmode=require");
    }

    #[test]
    fn swap_database_handles_dsn_without_path_or_query() {
        let out = swap_database("postgres://user@host:5432", "nomi_user_abc").unwrap();
        assert_eq!(out, "postgres://user@host:5432/nomi_user_abc");
    }
}