bot/src/server/mod.rs

pub mod auth;
pub mod chat;
pub mod notifications;
pub mod tasks;

use axum::{
    Router,
    http::HeaderValue,
    routing::{get, post},
};
use migration::{Migrator, MigratorTrait};
use sea_orm::{Database, DatabaseConnection, EntityTrait};
use std::sync::Arc;
use tower_http::cors::{AllowOrigin, CorsLayer};

use crate::entities::task::Entity as Task;
use crate::scheduler::Scheduler;

use crate::error::AppResult;

#[derive(Clone)]
pub struct AppState {
    pub db: DatabaseConnection,
    pub scheduler: Arc<Scheduler>,
    pub config: Arc<crate::config::Config>,
    pub verifier: Arc<crate::domain::auth::JwksVerifier>,
    pub authenticator: Arc<crate::domain::auth::Authenticator>,
    pub tx: tokio::sync::broadcast::Sender<crate::server::notifications::WsEvent>,
}

pub async fn start(config: crate::config::Config) -> AppResult<()> {
    let db = setup_database(&config.database_url).await?;

    let config = Arc::new(config);

    let (tx, _) = tokio::sync::broadcast::channel(100);

    let scheduler = Arc::new(
        Scheduler::new(db.clone(), config.clone(), tx.clone())
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))?,
    );

    // Load existing scheduled tasks
    let existing_tasks = Task::find()
        .all(&db)
        .await
        .map_err(crate::error::AppError::Database)?;
    for task in existing_tasks {
        if let Some(cron) = task.cron {
            let _ = scheduler.add_task_job(task.id, &cron).await;
        }
    }

    let (verifier, authenticator) = setup_auth(&config).await?;

    let state = Arc::new(AppState {
        db,
        scheduler,
        config: config.clone(),
        verifier,
        authenticator,
        tx,
    });

    let app = build_app(state, &config);

    let addr = format!("0.0.0.0:{}", config.port);
    let listener = tokio::net::TcpListener::bind(&addr)
        .await
        .map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
    tracing::info!("Server running on http://localhost:{}", config.port);
    axum::serve(listener, app)
        .await
        .map_err(|e| crate::error::AppError::Internal(e.to_string()))?;

    Ok(())
}

async fn setup_database(database_url: &str) -> AppResult<DatabaseConnection> {
    let db = Database::connect(database_url)
        .await
        .map_err(crate::error::AppError::Database)?;
    Migrator::up(&db, None)
        .await
        .map_err(crate::error::AppError::Database)?;
    Ok(db)
}

async fn setup_auth(
    config: &crate::config::Config,
) -> AppResult<(
    Arc<crate::domain::auth::JwksVerifier>,
    Arc<crate::domain::auth::Authenticator>,
)> {
    let verifier = Arc::new(
        crate::domain::auth::JwksVerifier::new(
            config.authentik_issuer.clone(),
            config.authentik_client_id.clone(),
        )
        .await
        .map_err(|e| crate::error::AppError::Internal(e.to_string()))?,
    );
    let authenticator = Arc::new(
        crate::domain::auth::Authenticator::new(
            config.authentik_issuer.clone(),
            config.authentik_client_id.clone(),
            config.authentik_client_secret.clone(),
        )
        .await
        .map_err(|e| crate::error::AppError::Internal(e.to_string()))?,
    );

    Ok((verifier, authenticator))
}

fn build_app(state: Arc<AppState>, config: &crate::config::Config) -> Router {
    let cors = build_cors_layer(config);

    Router::new()
        .route("/api/tasks", post(tasks::create_task).get(tasks::list_tasks))
        .route("/api/tasks/:id", get(tasks::get_task).put(tasks::update_task))
        .route("/api/tasks/:id/runs", post(tasks::rerun_task))
        .route("/api/runs/recent", get(tasks::get_recent_runs))
        .route("/api/auth/session", get(auth::auth_session))
        .route("/api/auth/callback", get(auth::auth_callback))
        .route("/api/auth/refresh", post(auth::auth_refresh))
        .route("/api/auth/logout", post(auth::auth_logout))
        .route("/api/chat", post(chat::chat_handler))
        .route("/api/ws", get(notifications::ws_handler))
        .layer(cors)
        .layer(tower_http::set_header::SetResponseHeaderLayer::overriding(
            axum::http::header::CONTENT_SECURITY_POLICY,
            HeaderValue::from_static("default-src 'self'; script-src 'self' 'unsafe-inline'; style-src 'self' 'unsafe-inline'; img-src 'self' data:;"),
        ))
        .layer(tower_http::set_header::SetResponseHeaderLayer::overriding(
            axum::http::header::X_CONTENT_TYPE_OPTIONS,
            HeaderValue::from_static("nosniff"),
        ))
        .layer(tower_http::set_header::SetResponseHeaderLayer::overriding(
            axum::http::header::REFERRER_POLICY,
            HeaderValue::from_static("strict-origin-when-cross-origin"),
        ))
        .layer(tower_http::limit::RequestBodyLimitLayer::new(1024 * 1024)) // 1MB limit
        .with_state(state)
}

fn build_cors_layer(config: &crate::config::Config) -> CorsLayer {
    let allow_origin = if let Some(origins) = &config.cors_allowed_origins {
        let values: Vec<HeaderValue> = origins
            .split(',')
            .map(|origin| origin.trim())
            .filter(|origin| !origin.is_empty())
            .filter_map(|origin| HeaderValue::from_str(origin).ok())
            .collect();

        if values.is_empty() {
            AllowOrigin::mirror_request()
        } else {
            AllowOrigin::list(values)
        }
    } else {
        AllowOrigin::mirror_request()
    };

    CorsLayer::new()
        .allow_origin(allow_origin)
        .allow_methods([
            axum::http::Method::GET,
            axum::http::Method::POST,
            axum::http::Method::PUT,
            axum::http::Method::PATCH,
            axum::http::Method::DELETE,
            axum::http::Method::OPTIONS,
        ])
        .allow_headers([
            axum::http::header::CONTENT_TYPE,
            axum::http::header::AUTHORIZATION,
            axum::http::header::ACCEPT,
        ])
        .allow_credentials(true)
}