Mercury/src/main.rs

mod auth;
mod cache;
mod config;
mod db;
mod models;
mod routes;
mod state;

use std::io::{self, Write};
use std::sync::Arc;

use axum::{
    extract::DefaultBodyLimit,
    http::{header, HeaderValue, Method},
    middleware,
    routing::{get, post},
    Router,
};
use tower_http::{
    cors::CorsLayer,
    services::{ServeDir, ServeFile},
};
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};

use crate::{
    auth::middleware::{blacklist_layer, require_auth},
    cache::spawn_sweep_task,
    config::Config,
    db::create_pool,
    models::blacklist::BlacklistEntry,
    routes::{admin::admin_router, auth::login, crud::handle_crud},
    state::{AppState, BlacklistCache, QueryCache},
};

fn build_cors(origins: &[String]) -> CorsLayer {
    if origins.is_empty() {
        return CorsLayer::new();
    }
    if origins.iter().any(|o| o == "*") {
        return CorsLayer::permissive();
    }
    let parsed: Vec<HeaderValue> = origins.iter().filter_map(|o| o.parse().ok()).collect();
    CorsLayer::new()
        .allow_origin(parsed)
        .allow_methods([
            Method::GET,
            Method::POST,
            Method::PUT,
            Method::DELETE,
            Method::OPTIONS,
        ])
        .allow_headers([header::CONTENT_TYPE, header::AUTHORIZATION])
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    dotenvy::dotenv().ok();

    tracing_subscriber::registry()
        .with(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| "mercury=info".into()),
        )
        .with(tracing_subscriber::fmt::layer())
        .init();

    let config = Arc::new(Config::from_env()?);
    let pool = create_pool(&config.database_url).await?;

    // If no users exist, prompt to create the first admin interactively.
    let user_count: i64 = sqlx::query_scalar::<_, Option<i64>>("SELECT COUNT(*) FROM users")
        .fetch_one(&pool)
        .await?
        .unwrap_or(0);
    if user_count == 0 {
        println!("\nNo users found. Create the first admin account.");
        print!("Username: ");
        io::stdout().flush()?;
        let mut username = String::new();
        io::stdin().read_line(&mut username)?;
        let username = username.trim().to_string();
        if username.is_empty() {
            anyhow::bail!("username cannot be empty");
        }

        print!("Password: ");
        io::stdout().flush()?;
        let mut password = String::new();
        io::stdin().read_line(&mut password)?;
        let password = password.trim().to_string();
        if password.is_empty() {
            anyhow::bail!("password cannot be empty");
        }

        let hash = bcrypt::hash(&password, bcrypt::DEFAULT_COST)?;
        sqlx::query(
            "INSERT INTO users (username, password_hash, permissions_mask) VALUES ($1, $2, '63')",
        )
        .bind(&username)
        .bind(hash)
        .execute(&pool)
        .await?;
        tracing::info!("created first admin user: {}", username);
        println!("Admin user '{}' created. Starting server...\n", username);
    }

    let query_cache = QueryCache::new();
    let blacklist_cache = BlacklistCache::new();

    // Load blacklist from DB into memory
    let entries = sqlx::query_as::<_, BlacklistEntry>(
        "SELECT id, pattern, method, reason, active, bypass_mask, created_at FROM blacklist ORDER BY id"
    )
    .fetch_all(&pool)
    .await?;
    blacklist_cache.load(entries).await;

    // Start cache sweep
    spawn_sweep_task(
        query_cache.clone(),
        config.cache_idle_timeout_secs,
        config.cache_sweep_interval_secs,
    );

    let state = AppState {
        pool,
        query_cache,
        blacklist_cache,
        config: config.clone(),
    };

    let crud_routes = Router::new()
        .route("/api/:table", get(handle_crud).post(handle_crud))
        .route("/api/:table/", get(handle_crud).post(handle_crud))
        .route(
            "/api/:table/:id",
            get(handle_crud).put(handle_crud).delete(handle_crud),
        )
        .route(
            "/api/:table/:id/",
            get(handle_crud).put(handle_crud).delete(handle_crud),
        )
        .layer(DefaultBodyLimit::max(1024 * 1024))
        // require_auth is inner (added first); blacklist_layer is outer (added last, runs first).
        // Order: blacklist check → auth check → handler.
        .route_layer(middleware::from_fn_with_state(state.clone(), require_auth))
        .route_layer(middleware::from_fn_with_state(
            state.clone(),
            blacklist_layer,
        ));

    let cors_layer = build_cors(&config.cors_origins);

    let app = Router::new()
        .route("/auth/login", post(login))
        .merge(crud_routes)
        .nest("/api/admin", admin_router(state.clone()))
        .nest_service(
            "/",
            ServeDir::new("ui/dist").fallback(ServeFile::new("ui/dist/index.html")),
        )
        .layer(cors_layer)
        .with_state(state);

    let addr = std::net::SocketAddr::from(([0, 0, 0, 0], 3000));
    tracing::info!("listening on {}", addr);
    let listener = tokio::net::TcpListener::bind(addr).await?;
    axum::serve(listener, app).await?;
    Ok(())
}