Hostityourself/server/src/main.rs

use axum::{
    middleware,
    routing::{delete, get, post, put},
    Router,
};
use std::sync::Arc;
use tokio::sync::Mutex;
use std::collections::VecDeque;
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};

mod auth;
mod builder;
mod crypto;
mod db;
mod models;
mod routes;

pub use db::DbPool;

#[derive(Clone)]
pub struct AppState {
    pub db: DbPool,
    /// Admin Postgres pool for schema/user provisioning. None if POSTGRES_URL is unset.
    pub pg: Option<sqlx::PgPool>,
    /// Queue of deploy IDs waiting to be processed.
    pub build_queue: Arc<Mutex<VecDeque<String>>>,
    pub data_dir: String,
    /// token → user_id (in-memory; cleared on restart).
    pub sessions: auth::Sessions,
    /// True when at least one user exists in the database.
    pub auth_enabled: bool,
    /// Bootstrap credentials (used when users table is empty).
    pub admin_user: Option<String>,
    pub admin_pass: Option<String>,
    /// e.g. "yourdomain.com" or "localhost" — used for cookie Domain + redirect URLs.
    pub domain_suffix: String,
    /// Shared secret used by the git-shell and post-receive hook to call internal API routes.
    pub internal_token: String,
    /// Path to the SSH authorized_keys file managed by HIY (e.g. /home/hiy/.ssh/authorized_keys).
    /// Empty string means SSH key management is disabled.
    pub ssh_authorized_keys_file: String,
    /// Path to the hiy-git-shell binary on the host (embedded in authorized_keys command= lines).
    pub git_shell_path: String,
}

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

    tracing_subscriber::registry()
        .with(
            tracing_subscriber::EnvFilter::new(
                std::env::var("RUST_LOG")
                    .unwrap_or_else(|_| "hiy_server=debug,tower_http=debug".into()),
            ),
        )
        .with(tracing_subscriber::fmt::layer())
        .init();

    let data_dir = std::env::var("HIY_DATA_DIR").unwrap_or_else(|_| "./data".into());
    std::fs::create_dir_all(&data_dir)?;

    let db = db::connect(&data_dir).await?;
    db::migrate(&db).await?;

    let admin_user = std::env::var("HIY_ADMIN_USER").ok().filter(|s| !s.is_empty());
    let admin_pass = std::env::var("HIY_ADMIN_PASS").ok().filter(|s| !s.is_empty());
    let domain_suffix =
        std::env::var("DOMAIN_SUFFIX").unwrap_or_else(|_| "localhost".into());

    // Count DB users.
    let user_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM users")
        .fetch_one(&db)
        .await
        .unwrap_or(0);

    // Bootstrap: create initial admin from env vars if no users exist yet.
    if user_count == 0 {
        if let (Some(u), Some(p)) = (&admin_user, &admin_pass) {
            let hash = bcrypt::hash(p, bcrypt::DEFAULT_COST)?;
            let id = uuid::Uuid::new_v4().to_string();
            let now = chrono::Utc::now().to_rfc3339();
            sqlx::query(
                "INSERT INTO users (id, username, password_hash, is_admin, created_at) \
                 VALUES (?,?,?,1,?)",
            )
            .bind(&id)
            .bind(u.as_str())
            .bind(&hash)
            .bind(&now)
            .execute(&db)
            .await?;
            tracing::info!("Bootstrap: created admin user '{}'", u);
        } else {
            tracing::warn!(
                "No users in database and HIY_ADMIN_USER/HIY_ADMIN_PASS not set — \
                 dashboard is unprotected! Set both in .env to create the first admin."
            );
        }
    }

    let auth_enabled: bool = sqlx::query_scalar("SELECT COUNT(*) FROM users")
        .fetch_one(&db)
        .await
        .unwrap_or(0i64)
        > 0;

    let build_queue = Arc::new(Mutex::new(VecDeque::<String>::new()));

    // Internal token: load from file (persists across restarts) or generate fresh.
    let token_file = format!("{}/internal-token", data_dir);
    let internal_token = if let Ok(t) = std::fs::read_to_string(&token_file) {
        t.trim().to_string()
    } else {
        let t = uuid::Uuid::new_v4().to_string();
        let _ = std::fs::write(&token_file, &t);
        t
    };

    let ssh_authorized_keys_file = std::env::var("HIY_SSH_AUTHORIZED_KEYS")
        .unwrap_or_default();
    let git_shell_path = std::env::var("HIY_GIT_SHELL")
        .unwrap_or_else(|_| "/usr/local/bin/hiy-git-shell".into());

    let pg = match std::env::var("POSTGRES_URL") {
        Ok(url) => match sqlx::postgres::PgPoolOptions::new()
            .max_connections(5)
            .connect_lazy(&url)
        {
            Ok(pool) => { tracing::info!("Postgres pool initialised (lazy)"); Some(pool) }
            Err(e)   => { tracing::warn!("Could not create Postgres pool: {}", e); None }
        },
        Err(_) => { tracing::info!("POSTGRES_URL not set — database provisioning disabled"); None }
    };

    let state = AppState {
        db,
        pg,
        build_queue,
        data_dir,
        sessions: auth::new_sessions(),
        auth_enabled,
        admin_user,
        admin_pass,
        domain_suffix,
        internal_token,
        ssh_authorized_keys_file,
        git_shell_path,
    };

    // Single background worker — sequential builds to avoid saturating the Pi.
    let worker_state = state.clone();
    tokio::spawn(async move {
        builder::build_worker(worker_state).await;
    });

    // Re-register all app Caddy routes from the DB on startup.
    // Caddy no longer uses --resume, so routes must be restored each time the
    // stack restarts (ensures Caddyfile changes are always picked up).
    let restore_db = state.db.clone();
    tokio::spawn(async move {
        routes::apps::restore_caddy_routes(&restore_db).await;
    });

    // ── Protected routes (admin login required) ───────────────────────────────
    let protected = Router::new()
        .route("/",           get(routes::ui::index))
        .route("/apps/:id",   get(routes::ui::app_detail))
        .route("/admin/users", get(routes::ui::users_page))
        .route("/api/status", get(routes::ui::status_json))
        // Apps API
        .route("/api/apps",          get(routes::apps::list).post(routes::apps::create))
        .route("/api/apps/:id",      get(routes::apps::get_one)
                                    .put(routes::apps::update)
                                    .patch(routes::apps::update)
                                    .delete(routes::apps::delete))
        .route("/api/apps/:id/stop",    post(routes::apps::stop))
        .route("/api/apps/:id/restart", post(routes::apps::restart))
        // Deploys API
        .route("/api/apps/:id/deploy",   post(routes::deploys::trigger))
        .route("/api/apps/:id/deploys",  get(routes::deploys::list))
        .route("/api/deploys/:id",       get(routes::deploys::get_one))
        .route("/api/deploys/:id/logs",  get(routes::deploys::logs_sse))
        // Env vars API
        .route("/api/apps/:id/env",         get(routes::envvars::list).post(routes::envvars::set))
        .route("/api/apps/:id/env/:key",    delete(routes::envvars::remove))
        // Database provisioning API
        .route("/api/apps/:id/database",    get(routes::databases::get_db)
                                           .post(routes::databases::provision)
                                           .delete(routes::databases::deprovision))
        // Users API (admin only — already behind admin middleware)
        .route("/api/users",                  get(routes::users::list).post(routes::users::create))
        .route("/api/users/:id",              put(routes::users::update).delete(routes::users::delete))
        .route("/api/users/:id/apps/:app_id", post(routes::users::grant_app).delete(routes::users::revoke_app))
        // SSH key management (admin only)
        .route("/api/users/:id/ssh-keys",     get(routes::ssh_keys::list).post(routes::ssh_keys::add))
        .route("/api/ssh-keys/:key_id",       delete(routes::ssh_keys::remove))
        // API key management (admin only)
        .route("/api/users/:id/api-keys",     get(routes::api_keys::list).post(routes::api_keys::create))
        .route("/api/api-keys/:key_id",       delete(routes::api_keys::revoke))
        .route_layer(middleware::from_fn_with_state(state.clone(), auth::auth_middleware));

    // ── Public routes ─────────────────────────────────────────────────────────
    let public = Router::new()
        .route("/login",    get(auth::login_page).post(auth::handle_login))
        .route("/logout",   get(auth::handle_logout))
        .route("/denied",   get(auth::denied_page))
        // Called by Caddy forward_auth for every deployed-app request.
        .route("/auth/verify", get(auth::verify))
        // GitHub webhooks use HMAC-SHA256 — no session needed.
        .route("/webhook/:app_id", post(routes::webhooks::github))
        // HTTP Smart Protocol for git push — authenticated by API key.
        .route("/git/:app/info/refs",         get(routes::git::http_info_refs))
        .route("/git/:app/git-receive-pack",  post(routes::git::http_receive_pack))
        // Internal routes: called by hiy-git-shell and post-receive hooks.
        .route("/internal/git/auth",          get(routes::git::auth_check))
        .route("/internal/git/:app_id/push",  post(routes::git::push_trigger));

    let app = Router::new()
        .merge(protected)
        .merge(public)
        .with_state(state);

    let addr = std::env::var("HIY_ADDR").unwrap_or_else(|_| "0.0.0.0:3000".into());
    let listener = tokio::net::TcpListener::bind(&addr).await?;
    tracing::info!("Listening on http://{}", addr);
    axum::serve(listener, app).await?;

    Ok(())
}