path: root/src/main.rs

mod config;
mod errors;

use std::{net::SocketAddr, str::FromStr, sync::Arc, time::Duration};

use axum::body;
use axum::extract::Path;
use axum::{error_handling::HandleErrorLayer, routing::get, BoxError, Extension, Router};
use axum::response::{Html, IntoResponse, Response};
use errors::{StringResult, HtmlResult};
use hyper::StatusCode;
use sqlx::{PgPool, postgres::PgPoolOptions};
use tower::ServiceBuilder;
use tower_http::trace::TraceLayer;
use tracing::{error, info, debug};
use crate::errors::ServiceError;
use tracing_subscriber::prelude::*;

pub struct State {
    pub config: config::Config,
    pub conn: PgPool
}

#[tokio::main]
async fn main() {
    color_eyre::install().unwrap();
    tracing_subscriber::registry()
        .with(tracing_subscriber::EnvFilter::new("debug"))
        .with(tracing_subscriber::fmt::layer())
        .init();
    let config = match config::Config::init("/etc/nccd/config.toml".to_owned()) {
        Ok(c) => c,
        Err(e) => {
            error!("Config Error: {:?}", e);
            std::process::exit(1);
        }
    };

    let bind_addr = config.server.bind_addr.clone();

    let db_config = config.database.clone();

    let conn = PgPoolOptions::new()
        .max_connections(db_config.max_connections)
        .connect(&db_config.postgres_url)
        .await.unwrap();

    let shared_state = Arc::new(State { config, conn });

    let app = Router::new()
        .route("/health", get(health_check))
        .route("/", get(index))
        .route("/static/:name", get(statics))
        .layer(
            ServiceBuilder::new()
                .layer(HandleErrorLayer::new(|error: BoxError| async move {
                    if error.is::<tower::timeout::error::Elapsed>() {
                        Ok(StatusCode::REQUEST_TIMEOUT)
                    } else {
                        Err((
                            StatusCode::INTERNAL_SERVER_ERROR,
                            format!("Unhandled internal error: {}", error),
                        ))
                    }
                }))
                .timeout(Duration::from_secs(10))
                .layer(TraceLayer::new_for_http())
                .into_inner(),
        )
        .layer(Extension(shared_state));

    let addr = match SocketAddr::from_str(&bind_addr) {
        Ok(a) => a,
        Err(e) => {
            error!("Invalid bind addr: {:?}", e);
            std::process::exit(1);
        }
    };

    info!("Listening on {}", addr);

    axum::Server::bind(&addr)
        .serve(app.into_make_service())
        .await
        .unwrap();
}

async fn health_check() -> &'static str {
    "OK"
}

#[axum_macros::debug_handler]
async fn index() -> HtmlResult {
    let mut buf = Vec::new();
    crate::templates::index_html(&mut buf).unwrap();

    match String::from_utf8(buf) {
        Ok(s) => Ok(Html(s)),
        Err(_) => Err(ServiceError::NotFound),
    }
}

async fn statics(Path(name): Path<String>) -> Result<Response, ServiceError> {
    for s in templates::statics::STATICS {
        debug!("Name: {}\nContents:\n{:?}\n\n", s.name, s.content);
    }

    match templates::statics::StaticFile::get(&name) {
        Some(s) => match String::from_utf8(s.content.to_vec()) {
            Ok(c) => {
                let body = body::boxed(body::Full::from(c));
                
                Ok(Response::builder()
                    .header("Content-Type", "text/css")
                    .status(StatusCode::OK)
                    .body(body).unwrap())
            },
            Err(_) => Err(ServiceError::NotFound),
        },
        None => Err(ServiceError::NotFound),
    }
}

include!(concat!(env!("OUT_DIR"), "/templates.rs"));