mod auth_routes; mod helpers; mod models; mod routes; use crate::{auth::jwt_middleware_gurt, config::Config, discord_bot}; use colored::Colorize; use macros_rs::fmt::{crashln, string}; use std::{net::IpAddr, str::FromStr, sync::Arc, collections::HashMap}; use gurt::prelude::*; use gurt::{GurtStatusCode, Route}; #[derive(Clone)] pub(crate) struct AppState { trusted: IpAddr, config: Config, db: sqlx::PgPool, jwt_secret: String, } impl AppState { pub fn new(trusted: IpAddr, config: Config, db: sqlx::PgPool, jwt_secret: String) -> Self { Self { trusted, config, db, jwt_secret, } } } #[derive(Clone)] pub(crate) struct RateLimitState { limits: Arc>>>>, } impl RateLimitState { pub fn new() -> Self { Self { limits: Arc::new(tokio::sync::RwLock::new(HashMap::new())), } } pub async fn check_rate_limit(&self, key: &str, window_secs: i64, max_requests: usize) -> bool { let mut limits = self.limits.write().await; let now = chrono::Utc::now(); let window_start = now - chrono::Duration::seconds(window_secs); let entry = limits.entry(key.to_string()).or_insert_with(Vec::new); entry.retain(|×tamp| timestamp > window_start); if entry.len() >= max_requests { false } else { entry.push(now); true } } } struct AppHandler { app_state: AppState, rate_limit_state: Option, handler_type: HandlerType, } // Macro to reduce JWT middleware duplication macro_rules! handle_authenticated { ($ctx:expr, $app_state:expr, $handler:expr) => { match jwt_middleware_gurt(&$ctx, &$app_state.jwt_secret).await { Ok(claims) => $handler(&$ctx, $app_state, claims).await, Err(e) => Ok(GurtResponse::new(GurtStatusCode::Unauthorized) .with_string_body(&format!("Authentication failed: {}", e))), } }; } #[derive(Clone)] enum HandlerType { Index, GetDomain, GetDomains, GetTlds, CheckDomain, Register, Login, GetUserInfo, CreateInvite, RedeemInvite, CreateDomainInvite, RedeemDomainInvite, CreateDomain, UpdateDomain, DeleteDomain, } impl GurtHandler for AppHandler { fn handle(&self, ctx: &ServerContext) -> std::pin::Pin> + Send + '_>> { let app_state = self.app_state.clone(); let rate_limit_state = self.rate_limit_state.clone(); let handler_type = self.handler_type.clone(); let ctx_data = ( ctx.remote_addr, ctx.request.clone(), ); Box::pin(async move { let start_time = std::time::Instant::now(); let ctx = ServerContext { remote_addr: ctx_data.0, request: ctx_data.1, }; log::info!("Handler started for {} {} from {}", ctx.method(), ctx.path(), ctx.remote_addr); let result = match handler_type { HandlerType::Index => routes::index(app_state).await, HandlerType::GetDomain => routes::get_domain(&ctx, app_state).await, HandlerType::GetDomains => routes::get_domains(&ctx, app_state).await, HandlerType::GetTlds => routes::get_tlds(app_state).await, HandlerType::CheckDomain => routes::check_domain(&ctx, app_state).await, HandlerType::Register => auth_routes::register(&ctx, app_state).await, HandlerType::Login => auth_routes::login(&ctx, app_state).await, HandlerType::GetUserInfo => handle_authenticated!(ctx, app_state, auth_routes::get_user_info), HandlerType::CreateInvite => handle_authenticated!(ctx, app_state, auth_routes::create_invite), HandlerType::RedeemInvite => handle_authenticated!(ctx, app_state, auth_routes::redeem_invite), HandlerType::CreateDomainInvite => handle_authenticated!(ctx, app_state, auth_routes::create_domain_invite), HandlerType::RedeemDomainInvite => handle_authenticated!(ctx, app_state, auth_routes::redeem_domain_invite), HandlerType::CreateDomain => { // Check rate limit first if let Some(ref rate_limit_state) = rate_limit_state { let client_ip = ctx.client_ip().to_string(); if !rate_limit_state.check_rate_limit(&client_ip, 600, 5).await { return Ok(GurtResponse::new(GurtStatusCode::TooLarge).with_string_body("Rate limit exceeded: 5 requests per 10 minutes")); } } handle_authenticated!(ctx, app_state, routes::create_domain) }, HandlerType::UpdateDomain => handle_authenticated!(ctx, app_state, routes::update_domain), HandlerType::DeleteDomain => handle_authenticated!(ctx, app_state, routes::delete_domain), }; let duration = start_time.elapsed(); match &result { Ok(response) => { log::info!("Handler completed for {} {} in {:?} - Status: {}", ctx.method(), ctx.path(), duration, response.status_code); }, Err(e) => { log::error!("Handler failed for {} {} in {:?} - Error: {}", ctx.method(), ctx.path(), duration, e); } } result }) } } pub async fn start(cli: crate::Cli) -> std::io::Result<()> { let config = Config::new().set_path(&cli.config).read(); let trusted_ip = match IpAddr::from_str(&config.server.address) { Ok(addr) => addr, Err(err) => crashln!("Cannot parse address.\n{}", string!(err).white()), }; let db = match config.connect_to_db().await { Ok(pool) => pool, Err(err) => crashln!("Failed to connect to PostgreSQL database.\n{}", string!(err).white()), }; // Start Discord bot if !config.discord.bot_token.is_empty() { if let Err(e) = discord_bot::start_discord_bot(config.discord.bot_token.clone(), db.clone()).await { log::error!("Failed to start Discord bot: {}", e); } } let jwt_secret = config.auth.jwt_secret.clone(); let app_state = AppState::new(trusted_ip, config.clone(), db, jwt_secret); let rate_limit_state = RateLimitState::new(); // Create GURT server let mut server = GurtServer::new(); // Load TLS certificates if let Err(e) = server.load_tls_certificates(&config.server.cert_path, &config.server.key_path) { crashln!("Failed to load TLS certificates: {}", e); } server = server .route(Route::get("/"), AppHandler { app_state: app_state.clone(), rate_limit_state: None, handler_type: HandlerType::Index }) .route(Route::get("/domain/*"), AppHandler { app_state: app_state.clone(), rate_limit_state: None, handler_type: HandlerType::GetDomain }) .route(Route::get("/domains"), AppHandler { app_state: app_state.clone(), rate_limit_state: None, handler_type: HandlerType::GetDomains }) .route(Route::get("/tlds"), AppHandler { app_state: app_state.clone(), rate_limit_state: None, handler_type: HandlerType::GetTlds }) .route(Route::get("/check"), AppHandler { app_state: app_state.clone(), rate_limit_state: None, handler_type: HandlerType::CheckDomain }) .route(Route::post("/auth/register"), AppHandler { app_state: app_state.clone(), rate_limit_state: None, handler_type: HandlerType::Register }) .route(Route::post("/auth/login"), AppHandler { app_state: app_state.clone(), rate_limit_state: None, handler_type: HandlerType::Login }) .route(Route::get("/auth/me"), AppHandler { app_state: app_state.clone(), rate_limit_state: None, handler_type: HandlerType::GetUserInfo }) .route(Route::post("/auth/invite"), AppHandler { app_state: app_state.clone(), rate_limit_state: None, handler_type: HandlerType::CreateInvite }) .route(Route::post("/auth/redeem-invite"), AppHandler { app_state: app_state.clone(), rate_limit_state: None, handler_type: HandlerType::RedeemInvite }) .route(Route::post("/auth/create-domain-invite"), AppHandler { app_state: app_state.clone(), rate_limit_state: None, handler_type: HandlerType::CreateDomainInvite }) .route(Route::post("/auth/redeem-domain-invite"), AppHandler { app_state: app_state.clone(), rate_limit_state: None, handler_type: HandlerType::RedeemDomainInvite }) .route(Route::post("/domain"), AppHandler { app_state: app_state.clone(), rate_limit_state: Some(rate_limit_state), handler_type: HandlerType::CreateDomain }) .route(Route::put("/domain/*"), AppHandler { app_state: app_state.clone(), rate_limit_state: None, handler_type: HandlerType::UpdateDomain }) .route(Route::delete("/domain/*"), AppHandler { app_state: app_state.clone(), rate_limit_state: None, handler_type: HandlerType::DeleteDomain }); log::info!("GURT server listening on {}", config.get_address()); server.listen(&config.get_address()).await.map_err(|e| { std::io::Error::new(std::io::ErrorKind::Other, format!("GURT server error: {}", e)) }) }