prvng_platform/control-center/src/main.rs

//! Control Center Main Binary
//!
//! JWT authentication service with user management, role-based access control,
//! and real-time WebSocket events.

use control_center::{AppState, Config, Result};
use axum::{
    middleware,
    routing::{get, post},
    Router,
};
use clap::Parser;
use std::{path::PathBuf, sync::Arc};
use tokio::signal;
use tower::ServiceBuilder;
use tower_http::{compression::CompressionLayer, timeout::TimeoutLayer, trace::TraceLayer};
use tracing::{error, info};
use tracing_subscriber::EnvFilter;
use std::time::Duration;

use control_center::handlers::{
    auth::*,
    permission::list_permissions,
    websocket::websocket_handler,
};
use control_center::middleware::{
    auth::auth_middleware,
    cors::create_cors_from_env,
    rate_limit::{RateLimitConfig, RateLimitLayer},
};

#[derive(Parser)]
#[command(name = "control-center")]
#[command(about = "Control Center - JWT Authentication & User Management Service")]
#[command(version = env!("CARGO_PKG_VERSION"))]
struct Cli {
    /// Configuration file path
    #[arg(short, long, default_value = "config.toml")]
    config: Option<PathBuf>,

    /// Server port (overrides config file)
    #[arg(short, long)]
    port: Option<u16>,

    /// Server host (overrides config file)
    #[arg(long)]
    host: Option<String>,

    /// Enable debug logging
    #[arg(short, long)]
    debug: bool,

    /// Generate default configuration file
    #[arg(long)]
    generate_config: bool,
}

#[tokio::main]
async fn main() -> Result<()> {
    let cli = Cli::parse();

    // Generate default config if requested
    if cli.generate_config {
        let config_path = cli.config.unwrap_or_else(|| PathBuf::from("config.toml"));
        control_center::config::create_default_config_file(&config_path)?;
        info!("Default configuration file created at: {:?}", config_path);
        return Ok(());
    }

    // Initialize logging
    let log_level = if cli.debug { "debug" } else { "info" };
    let filter = EnvFilter::new(format!("control_center={},tower_http=info", log_level));

    tracing_subscriber::fmt()
        .with_env_filter(filter)
        .with_target(false)
        .init();

    // Load configuration
    let mut config = if let Some(config_path) = cli.config {
        Config::load_from_file(config_path)?
    } else {
        Config::load()?
    };

    // Apply CLI overrides
    if let Some(port) = cli.port {
        config.server.port = port;
    }
    if let Some(host) = cli.host {
        config.server.host = host;
    }

    info!(
        "Starting Control Center v{} on {}",
        env!("CARGO_PKG_VERSION"),
        config.bind_address()
    );

    // Initialize application state
    let app_state = Arc::new(AppState::new(config.clone()).await?);

    // Health check
    if let Err(e) = app_state.health_check().await {
        error!("Health check failed: {}", e);
        return Err(e);
    }

    info!("All services initialized successfully");

    // Create router
    let app = create_router(app_state.clone()).await?;

    // Start session cleanup task
    start_background_tasks(app_state.clone()).await;

    // Start server
    let bind_address = config.bind_address();
    let listener = tokio::net::TcpListener::bind(&bind_address).await?;

    info!("Control Center listening on http://{}", bind_address);
    info!("WebSocket endpoint available at: ws://{}/ws", bind_address);

    // Graceful shutdown
    axum::serve(listener, app)
        .with_graceful_shutdown(shutdown_signal())
        .await?;

    info!("Control Center shutdown complete");
    Ok(())
}

/// Create the main application router
async fn create_router(app_state: Arc<AppState>) -> Result<Router> {
    // Create rate limiting layers
    let auth_rate_limit = RateLimitLayer::new(RateLimitConfig::auth());
    let general_rate_limit = RateLimitLayer::new(RateLimitConfig::default());

    // Create CORS layer
    let cors_layer = create_cors_from_env();

    // Public routes (no authentication required)
    let public_routes = Router::new()
        .route("/health", get(health_check))
        .route("/auth/login", post(login))
        .route("/auth/refresh", post(refresh_token))
        .layer(auth_rate_limit);

    // Protected routes (authentication required)
    let protected_routes = Router::new()
        // Authentication routes
        .route("/auth/logout", post(logout))
        // .route("/auth/verify", get(verify_token))
        // .route("/auth/sessions", get(get_sessions))
        // .route("/auth/sessions/invalidate", post(invalidate_all_sessions))
        // .route("/auth/change-password", post(change_password))
        // User routes
        // .route("/users", get(list_users).post(create_user))
        // .route("/users/me", get(get_current_user).put(update_current_user))
        // .route(
        //     "/users/:id",
        //     get(get_user_by_id)
        //         .put(update_user_by_id)
        //         .delete(delete_user),
        // )
        // .route("/users/:id/activate", post(activate_user))
        // .route("/users/:id/deactivate", post(deactivate_user))
        // .route("/users/:id/verify", post(verify_user))
        // .route("/users/:id/roles/:role", post(add_user_role).delete(remove_user_role))
        // Role routes
        // .route("/roles", get(list_roles).post(create_role))
        // .route("/roles/:id", get(get_role_by_id).put(update_role).delete(delete_role))
        // Permission routes
        .route("/permissions", get(list_permissions))
        // .route("/permissions/resources", get(get_resources))
        // .route("/permissions/actions", get(get_actions))
        // WebSocket route
        .route("/ws", get(websocket_handler))
        // Apply authentication middleware to all protected routes
        .route_layer(middleware::from_fn_with_state(
            app_state.jwt_service.clone(),
            auth_middleware,
        ));

    // Combine all routes
    let app = Router::new()
        .merge(public_routes)
        .merge(protected_routes)
        .layer(
            ServiceBuilder::new()
                .layer(TraceLayer::new_for_http())
                .layer(CompressionLayer::new())
                .layer(TimeoutLayer::new(Duration::from_secs(30)))
                .layer(cors_layer)
                .layer(general_rate_limit),
        )
        .with_state(app_state);

    Ok(app)
}

/// Start background tasks
async fn start_background_tasks(app_state: Arc<AppState>) {
    let cleanup_interval = Duration::from_secs(
        app_state.config.security.session_cleanup_interval_minutes * 60,
    );

    // Session cleanup task
    let app_state_cleanup = app_state.clone();
    tokio::spawn(async move {
        let mut interval = tokio::time::interval(cleanup_interval);
        loop {
            interval.tick().await;

            if let Err(e) = app_state_cleanup.auth_service.cleanup_expired_sessions().await {
                error!("Failed to cleanup expired sessions: {}", e);
            } else {
                info!("Cleaned up expired sessions");
            }
        }
    });

    info!("Background tasks started");
}

/// Graceful shutdown signal handler
async fn shutdown_signal() {
    let ctrl_c = async {
        signal::ctrl_c()
            .await
            .expect("Failed to install Ctrl+C handler");
    };

    #[cfg(unix)]
    let terminate = async {
        signal::unix::signal(signal::unix::SignalKind::terminate())
            .expect("Failed to install signal handler")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c => {},
        _ = terminate => {},
    }

    info!("Shutdown signal received, starting graceful shutdown");
}