kvitt/server/crates/kvitt-server/src/cli/mod.rs

//! `kvitt-server` subcommands (SPEC §10.4, exhaustive): `serve` and `migrate` here, the others in their own files.

pub mod bootstrap;
pub mod healthcheck;
pub mod invite_partner;
pub mod pair_device;
pub mod rotate_epoch;

use std::net::SocketAddr;
use std::process::ExitCode;
use std::time::Duration;

use axum::Router;
use axum::extract::connect_info::IntoMakeServiceWithConnectInfo;

use clap::{Parser, Subcommand};
use tokio::net::TcpListener;
use tokio::sync::Notify;

use crate::api::invites::{self, NewInvite};
use crate::config::Config;
use crate::db::Db;
use crate::error::AppError;
use crate::{AppState, db, router, tasks};

/// How long `serve` lets in-flight requests finish after SIGTERM (SPEC §10.4).
pub const SHUTDOWN_DRAIN: Duration = Duration::from_secs(10);

#[derive(Debug, Parser)]
#[command(
    name = "kvitt-server",
    version,
    about = "Kvitt sync server (SPEC §10.4)"
)]
pub struct Cli {
    #[command(subcommand)]
    pub command: Command,
}

#[derive(Debug, Subcommand)]
pub enum Command {
    /// Run migrations, start the background tasks and the HTTP server.
    Serve,
    /// Create the admin_setup invite (FR-1). Exit 0 ok, 2 already initialized.
    Bootstrap,
    /// Create a device invite for an existing user (AC-1.6). Exit 0 ok, 3 unknown user.
    PairDevice {
        /// The user's id.
        #[arg(long = "user", value_name = "USER_ID")]
        user: String,
    },
    /// Create a partner invite by the admin, as E-12 does (D-113). Exit 0 ok, 2 household full, 3 no admin.
    InvitePartner,
    /// Set a new random sync_epoch (SPEC §10.9).
    RotateEpoch,
    /// GET /health on this host; exit 0 if 200, 1 otherwise.
    Healthcheck,
    /// Run migrations only.
    Migrate,
}

/// Runs a subcommand to completion and returns the process exit code.
pub async fn run(command: Command, config: Config) -> ExitCode {
    let result = match command {
        Command::Serve => serve(config).await,
        Command::Migrate => migrate(config).await,
        Command::Healthcheck => return healthcheck::run(&config).await,
        Command::Bootstrap => return bootstrap::run(&config).await,
        Command::PairDevice { user } => return pair_device::run(&config, &user).await,
        Command::InvitePartner => return invite_partner::run(&config).await,
        Command::RotateEpoch => return rotate_epoch::run(&config).await,
    };
    match result {
        Ok(()) => ExitCode::SUCCESS,
        Err(error) => {
            tracing::error!(error_code = error.code.as_str(), detail = %error.message, "command failed");
            eprintln!("{error}");
            ExitCode::FAILURE
        }
    }
}

/// The text the invite subcommands print (AC-1.1, AC-1.6, D-113): the code in groups of four, the landing link
/// (naming the inviting admin with `from`, D-087) and the expiry. The code goes to stdout only, never to the log
/// (SPEC §8.7).
pub fn render_invite(invite: &NewInvite, public_url: &str, from: Option<&str>) -> String {
    format!(
        "Code:    {}\nLink:    {}\nExpires: {}\n",
        invites::display_code(&invite.code),
        invites::invite_link(public_url, &invite.code, from),
        invite.expires_at
    )
}

/// Opens the database (see [`open_database`]), runs `job` on it and closes the pool again.
pub async fn with_database<T>(
    config: &Config,
    job: impl AsyncFnOnce(&Db) -> Result<T, AppError>,
) -> Result<T, AppError> {
    let pool = open_database(config).await?;
    let result = job(&pool).await;
    pool.close().await;
    result
}

/// Opens (creating if needed) the data directory and the database, and runs the migrations.
pub async fn open_database(config: &Config) -> Result<db::Db, AppError> {
    std::fs::create_dir_all(&config.data_dir).map_err(|e| {
        AppError::internal(format!("cannot create {}: {e}", config.data_dir.display()))
    })?;
    let pool = db::connect(&config.database_path()).await?;
    db::migrate(&pool).await?;
    Ok(pool)
}

async fn migrate(config: Config) -> Result<(), AppError> {
    with_database(&config, async |_| Ok(())).await?;
    tracing::info!("migrations applied");
    Ok(())
}

async fn serve(config: Config) -> Result<(), AppError> {
    let pool = open_database(&config).await?;
    let bind = config.bind;
    // `tasks::start` installs the push queue (or logs once that push is disabled) and starts the scheduler.
    let state = tasks::start(AppState::new(config, pool.clone()));

    let listener = TcpListener::bind(bind)
        .await
        .map_err(|e| AppError::internal(format!("cannot bind {bind}: {e}")))?;
    tracing::info!(%bind, version = crate::SERVER_VERSION, "listening");

    let shutdown = std::sync::Arc::new(Notify::new());
    let graceful = {
        let shutdown = shutdown.clone();
        async move { shutdown.notified().await }
    };
    let server = axum::serve(listener, make_service(state)).with_graceful_shutdown(graceful);
    let mut server = tokio::spawn(async move { server.await });

    tokio::select! {
        result = &mut server => {
            return finish(result, &pool).await;
        }
        () = shutdown_signal() => {}
    }
    tracing::info!(drain_secs = SHUTDOWN_DRAIN.as_secs(), "shutting down");
    shutdown.notify_one();
    match tokio::time::timeout(SHUTDOWN_DRAIN, &mut server).await {
        Ok(result) => finish(result, &pool).await,
        Err(_) => {
            tracing::warn!("in-flight requests did not finish within the drain period");
            server.abort();
            pool.close().await;
            Ok(())
        }
    }
}

/// The router as served: with the TCP peer address in `ConnectInfo<SocketAddr>`, which the rate limiter keys
/// unauthenticated requests by (SPEC §8.6).
pub fn make_service(state: AppState) -> IntoMakeServiceWithConnectInfo<Router, SocketAddr> {
    router::router(state).into_make_service_with_connect_info::<SocketAddr>()
}

async fn finish(
    result: Result<std::io::Result<()>, tokio::task::JoinError>,
    pool: &db::Db,
) -> Result<(), AppError> {
    pool.close().await;
    match result {
        Ok(Ok(())) => Ok(()),
        Ok(Err(e)) => Err(AppError::internal(format!("server error: {e}"))),
        Err(e) => Err(AppError::internal(format!("server task failed: {e}"))),
    }
}

/// Resolves on SIGTERM or Ctrl-C.
async fn shutdown_signal() {
    let ctrl_c = async {
        if let Err(e) = tokio::signal::ctrl_c().await {
            tracing::warn!(detail = %e, "cannot listen for Ctrl-C");
            std::future::pending::<()>().await;
        }
    };
    #[cfg(unix)]
    let terminate = async {
        match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
            Ok(mut signal) => {
                signal.recv().await;
            }
            Err(e) => {
                tracing::warn!(detail = %e, "cannot listen for SIGTERM");
                std::future::pending::<()>().await;
            }
        }
    };
    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();
    tokio::select! {
        () = ctrl_c => {}
        () = terminate => {}
    }
}