nyazoom/src/main.rs

use tokio_util::io::ReaderStream;

use axum::{
    body::Body,
    extract::{ConnectInfo, Request, State},
    http::{HeaderMap, StatusCode},
    middleware::{self, Next},
    response::{Html, IntoResponse, Redirect},
    routing::get,
    Router,
};
use axum_extra::TypedHeader;
use tower_http::{services::ServeDir, trace::TraceLayer};

use std::{io, net::SocketAddr};

mod router {
    pub mod link;
    pub mod records;
    pub mod upload;
}
mod cache;
mod state;
mod util;
mod views;

use util::{headers::ForwardedFor, logging, ssr, sweeper};

use state::*;
use views::*;

use router::link::get_link_router;
use router::records::get_records_router;
use router::upload::get_upload_router;

#[tokio::main]
async fn main() -> io::Result<()> {
    logging::init_tracing();

    // uses create_dir_all to create both .cache and serve inside it in one go
    util::make_dir(".cache/serve").await?;

    let state = cache::fetch_cache().await;

    sweeper::spawn(state.clone());

    // Router Setup
    let app = Router::new()
        .route("/", get(welcome))
        .route("/download/:id", get(download))
        .nest("/upload", get_upload_router())
        .nest("/records", get_records_router())
        .nest("/link", get_link_router())
        .with_state(state)
        .fallback_service(ServeDir::new("dist"))
        .layer(TraceLayer::new_for_http())
        .layer(middleware::from_fn(log_source));

    serve(app).await;

    Ok(())
}

async fn welcome() -> impl IntoResponse {
    let cat_fact = views::get_cat_fact().await;
    Html(ssr::render(move || {
        leptos::view! { <WelcomePage fact=cat_fact /> }
    }))
}

async fn log_source(
    ConnectInfo(addr): ConnectInfo<SocketAddr>,
    forwarded_for: Option<TypedHeader<ForwardedFor>>,
    req: Request,
    next: Next,
) -> impl IntoResponse {
    tracing::info!("{} : {:?}", addr, forwarded_for);

    next.run(req).await
}

async fn download(
    axum::extract::Path(id): axum::extract::Path<String>,
    headers: HeaderMap,
    State(state): State<AppState>,
) -> Result<axum::response::Response, (StatusCode, String)> {
    {
        let mut records = state.records.lock().await;
        if headers.get("hx-request").is_some() {
            return Ok(axum::http::Response::builder()
                .header("HX-Redirect", format!("/download/{id}"))
                .status(204)
                .body("".to_owned())
                .unwrap()
                .into_response());
        }

        if let Some(record) = records
            .get_mut(&id)
            .filter(|record| record.can_be_downloaded())
        {
            record.downloads += 1;

            let file = tokio::fs::File::open(&record.file).await.unwrap();

            return Ok(axum::response::Response::builder()
                .header("Content-Type", "application/zip")
                .body(Body::from_stream(ReaderStream::new(file)))
                .unwrap());
        } else {
            records.remove_record(&id).await.unwrap()
        }
    }

    Ok(Redirect::to("/404.html").into_response())
}

async fn serve(app: Router) {
    // // Server creation
    let addr = SocketAddr::from(([0, 0, 0, 0], 3000));
    let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
    axum::serve(
        listener,
        app.into_make_service_with_connect_info::<SocketAddr>(),
    )
    .await
    .unwrap();

    tracing::debug!("listening on http://{}/", addr);
}