frangipane-backend/src/lib.rs

pub mod auth;
pub mod db;
pub mod errors;
pub mod realtime;
pub mod routes;

use crate::errors::APIError;
pub use crate::routes::users; // etc
use axum::{
    Extension, Json, Router,
    extract::DefaultBodyLimit,
    http::{Method, header},
    routing::get,
};
use serde_json::json;
use std::sync::Arc;
use std::{path::PathBuf, time::Duration};
use tower_governor::{GovernorLayer, governor::GovernorConfigBuilder};
use tower_http::{
    cors::{Any, CorsLayer},
    trace::{DefaultMakeSpan, DefaultOnResponse, TraceLayer},
};
use tracing::Level;
use utoipa::OpenApi;
use utoipa_swagger_ui::SwaggerUi;

const MAX_USERNAME_LENGTH: usize = 35;
const MAX_ROOM_NAME_LENGTH: usize = 35;
const MAX_UPLOAD_SIZE: usize = 5 * 1024 * 1024;
pub const MAX_ATTACHMENTS_PER_MESSAGE: usize = 8;

pub struct AppConfig {
    pub avatar_dir: PathBuf,
    pub prohibit_registration: bool,
    pub max_file_size: usize,
    pub uploads_dir: PathBuf,
}

#[derive(clap::Parser, Debug)]
#[command(author, version, about, long_about = None)]
pub struct Cli {
    /// Server port
    #[arg(short, long, default_value = "8080")]
    pub port: String,

    /// Database URL
    #[arg(short, long, default_value = "0.0.0.0:5432")]
    pub database: String,

    /// Data directory path
    #[arg(short = 'D', long, default_value = "/var/lib/frangipane")]
    pub data_dir: String,

    /// Whether to disable user registration
    #[arg(short, long)]
    pub no_registration: bool,

    /// Max file upload size in bytes (default: 10MB)
    #[arg(long, default_value = "10485760")]
    pub max_file_size: usize,

    /// Verbose mode
    #[arg(short, long)]
    pub verbose: bool,
}

/// Define the OpenAPI specification document structure
#[derive(OpenApi)]
#[openapi(
    paths(
        auth::validate_token,
        routes::friends::list_friends,
        routes::friends::list_requests,
        routes::friends::send_request,
        routes::friends::accept_request,
        routes::friends::decline_request,
        routes::friends::is_friend,
        routes::friends::remove_friend,
        routes::messages::list_messages,
        routes::messages::create_message,
        routes::messages::message_ws_handler,
        routes::messages::create_message_with_attachments,
        routes::messages::get_upload,
        routes::rooms::list_rooms,
        routes::rooms::create_room,
        routes::rooms::get_room,
        routes::rooms::list_invites,
        routes::rooms::send_invite,
        routes::rooms::accept_request,
        routes::rooms::decline_request,
        routes::rooms::leave_room,
        routes::rooms::transfer_ownership,
        routes::rooms::list_members,
        routes::rooms::delete_room,
        routes::users::login,
        routes::users::register_user,
        routes::users::update_user,
        routes::users::upload_avatar,
        routes::users::get_avatar,
        routes::voice::voice_ws_handler,
        routes::ws::issue_ws_token,
    ),
    components(
        schemas(
            errors::ErrorResponse,
            routes::friends::Friend,
            routes::friends::FriendRequest,
            routes::friends::SendFriendRequestPayload,
            routes::friends::AcceptFriendRequestPayload,
            routes::friends::DeclineFriendRequestPayload,
            routes::friends::RemoveFriendPayload,
            routes::messages::Message,
            routes::messages::Attachment,
            routes::messages::NewMessagePayload,
            routes::rooms::Room,
            routes::rooms::NewRoomPayload,
            routes::rooms::RoomInvite,
            routes::rooms::SendRoomInvitePayload,
            routes::rooms::AcceptRoomInvitePayload,
            routes::rooms::TransferOwnershipPayload,
            routes::users::User,
            routes::users::UserProfile,
            routes::users::LoginPayload,
            routes::users::LoginResponse,
            routes::users::NewUserPayload,
            routes::users::UpdateUserPayoad,
            routes::users::UpdateUserResponse,
            routes::ws::WsAuthQuery,
        )
    ),
    modifiers(&SecurityAddon)
)]
pub struct ApiDoc;

struct SecurityAddon;

impl utoipa::Modify for SecurityAddon {
    fn modify(&self, openapi: &mut utoipa::openapi::OpenApi) {
        if let Some(components) = openapi.components.as_mut() {
            components.add_security_scheme(
                "bearer_auth",
                utoipa::openapi::security::SecurityScheme::Http(
                    utoipa::openapi::security::HttpBuilder::new()
                        .scheme(utoipa::openapi::security::HttpAuthScheme::Bearer)
                        .bearer_format("JWT")
                        .build(),
                ),
            );
        }
    }
}

/// Create the app with all routes, middleware and extensions
pub fn create_app(
    db_pool: sqlx::PgPool,
    config: Arc<AppConfig>,
    messages: realtime::RealtimeMessages,
    voice: realtime::RealTimeVoices,
    use_rate_limiter: bool,
) -> Router {
    let governor_conf = GovernorConfigBuilder::default()
        .burst_size(20)
        .per_millisecond(250)
        .finish()
        .unwrap();

    let governor_limiter = governor_conf.limiter().clone();

    let voice_clone = voice.clone();

    // A separate background task to clean up
    let interval = Duration::from_secs(60);
    std::thread::spawn(move || {
        loop {
            std::thread::sleep(interval);

            // tracing::info!("rate limiting storage size: {}", governor_limiter.len());
            governor_limiter.retain_recent();

            voice_clone.retain_active();
        }
    });

    let cors = CorsLayer::new()
        .allow_origin(Any)
        .allow_methods([Method::GET, Method::POST])
        .allow_headers([header::AUTHORIZATION, header::CONTENT_TYPE]);

    let mut router = Router::new()
        .merge(SwaggerUi::new("/swagger-ui").url("/api-docs/openapi.json", ApiDoc::openapi()))
        .route("/version", get(get_version))
        .merge(routes::users::routes())
        .merge(routes::rooms::routes())
        .merge(routes::messages::routes())
        .merge(routes::voice::routes())
        .merge(routes::friends::routes())
        .merge(routes::ws::routes())
        .layer(Extension(db_pool))
        .layer(Extension(messages))
        .layer(Extension(voice))
        .layer(Extension(config))
        .layer(cors)
        .layer(DefaultBodyLimit::max(1024 * 5 * 100))
        .layer(
            TraceLayer::new_for_http()
                .make_span_with(DefaultMakeSpan::new().level(Level::DEBUG))
                .on_request(())
                .on_response(DefaultOnResponse::new().level(Level::DEBUG)),
        );

    if use_rate_limiter {
        router = router.layer(GovernorLayer::new(governor_conf));
    }

    router
}

/// Public route to get current version
async fn get_version() -> Result<Json<serde_json::Value>, APIError> {
    const VERSION: &str = env!("CARGO_PKG_VERSION");
    Ok(Json(json!({ "version": VERSION })))
}