yamaf/src/main.rs

use axum::{
    Router,
    body::Body,
    extract::{DefaultBodyLimit, Multipart, Path},
    http::{HeaderMap, StatusCode, header},
    response::{Html, IntoResponse, Response},
    routing::get,
};
use rand::{Rng, distr::Alphanumeric};
use std::{env, net::SocketAddr, sync::LazyLock};
use std::{env::VarError, path::PathBuf};
use tokio::fs;
use tokio_util::io::ReaderStream;

struct Config {
    root_dir: String,
    key: Result<String, VarError>,
    title: String,
    internal_host: String,
    internal_port: u16,
    external_host: String,
    external_protocol: &'static str,
    max_filesize: usize,
    max_bodysize: usize,
    min_filedays: usize,
    max_filedays: usize,
}

static CONFIG: LazyLock<Config> = LazyLock::new(|| {
    let root_dir = env::var("ROOT_DIR").unwrap_or_else(|_| "/var/files".to_string());
    let key = env::var("KEY");
    let title = env::var("TITLE").unwrap_or_else(|_| "Yet Another Mid Ahh Filehost".to_string());

    let internal_host = env::var("INTERNAL_HOST").unwrap_or_else(|_| "127.0.0.1".to_string());
    let internal_port = env::var("INTERNAL_PORT")
        .ok()
        .and_then(|x| x.parse().ok())
        .unwrap_or(8000);

    let external_host = env::var("EXTERNAL_HOST").unwrap_or_else(|_| internal_host.clone());
    let external_protocol = if env::var("EXTERNAL_HAS_TLS").is_ok() {
        "https"
    } else {
        "http"
    };

    let max_files = env::var("MAX_FILES")
        .ok()
        .and_then(|x| x.parse().ok())
        .unwrap_or(10);
    let max_filesize = env::var("MAX_FILESIZE_MB")
        .ok()
        .and_then(|x| x.parse().ok())
        .unwrap_or(100)
        << 20;
    let max_bodysize = max_files * max_filesize * 2;

    let min_filedays = env::var("MIN_FILEDAYS")
        .ok()
        .and_then(|x| x.parse().ok())
        .unwrap_or(30);

    let max_filedays = env::var("MAX_FILEDAYS")
        .ok()
        .and_then(|x| x.parse().ok())
        .unwrap_or(365);

    assert!(max_filedays >= min_filedays);

    Config {
        root_dir,
        key,
        title,
        internal_host,
        internal_port,
        external_host,
        external_protocol,
        max_filesize,
        max_bodysize,
        min_filedays,
        max_filedays,
    }
});

#[tokio::main]
async fn main() {
    let addr = SocketAddr::new(
        CONFIG.internal_host.parse().expect("Invalid Host Bind"),
        CONFIG.internal_port,
    );

    let app = Router::new()
        .route("/", get(index).post(upload))
        .route("/{filename}", get(serve_file))
        .layer(DefaultBodyLimit::max(CONFIG.max_bodysize));

    let listener = tokio::net::TcpListener::bind(addr).await.unwrap();

    println!(
        "Starting server on {} for directory {}",
        addr, CONFIG.root_dir
    );

    /* cleanup cronjob */
    #[cfg(feature = "cleanup")]
    tokio::spawn(async move {
        use chrono::{Local, NaiveDate};

        let mut last_cleanup_day: NaiveDate = Local::now().date_naive();

        loop {
            use tokio::time::{Duration, sleep};

            /* every 30 minutes */
            sleep(Duration::from_secs(30 * 60)).await;

            let now = Local::now();
            let today = now.date_naive();

            /* if the date changed, midnight has passed */
            if today > last_cleanup_day {
                println!("attempting cleanup at {}", now.to_string());
                if let Err(e) = cleanup_directory().await {
                    eprintln!("cleanup failed: {}", e);
                } else {
                    println!("cleanup succeeded!");
                    last_cleanup_day = today;
                }
            }
        }
    });

    axum::serve(listener, app).await.unwrap();
}

static INDEX_HTML: LazyLock<String> = LazyLock::new(|| {
    let mut html = include_str!("./index.html").to_string();

    html = html.replace("{{TITLE}}", &CONFIG.title);

    html = html.replace(
        "{{USER_URL}}",
        &format!("{}://{}", CONFIG.external_protocol, CONFIG.external_host),
    );

    html = html.replace(
        "{{KEY_FIELD}}",
        if CONFIG.key.is_ok() {
            r#"<input type="password" name="key" placeholder="Upload Key" required><br><br>"#
        } else {
            ""
        },
    );

    if cfg!(feature = "cleanup") {
        html = html
            .replace("{{MAX_SIZE}}", &(CONFIG.max_filesize >> 20).to_string())
            .replace("{{MAX_DAYS}}", &CONFIG.max_filedays.to_string())
            .replace("{{MIN_DAYS}}", &CONFIG.min_filedays.to_string());
    } else {
        if let Some(start) = html.find("<!-- CLEANUP_START -->") {
            if let Some(end) = html.find("<!-- CLEANUP_END -->") {
                html.replace_range(start..end + "<!-- CLEANUP_END -->".len(), "");
            }
        }
    }

    html
});

async fn index() -> Html<&'static str> {
    Html(&INDEX_HTML)
}

#[derive(Debug)]
enum YamafError {
    BadRequest(String),
    InternalError(String),
    FileTooBig(String),
    FileNotFound,
}

impl IntoResponse for YamafError {
    fn into_response(self) -> Response {
        match self {
            YamafError::BadRequest(msg) => (StatusCode::BAD_REQUEST, msg).into_response(),
            YamafError::InternalError(msg) => {
                (StatusCode::INTERNAL_SERVER_ERROR, msg).into_response()
            }
            YamafError::FileTooBig(filename) => (
                StatusCode::PAYLOAD_TOO_LARGE,
                format!("File {} is too big!", filename),
            )
                .into_response(),
            YamafError::FileNotFound => (StatusCode::NOT_FOUND, "File Not Found!").into_response(),
        }
    }
}

fn random(len: usize) -> String {
    rand::rng()
        .sample_iter(&Alphanumeric)
        .take(len)
        .map(char::from)
        .collect()
}

fn clean_filename(filename: &str) -> String {
    let mut slug = String::new();
    let mut prev_dash = false;

    for c in filename.to_lowercase().chars() {
        if c.is_ascii_alphanumeric() || c == '.' {
            slug.push(c);
            prev_dash = false;
        } else if !prev_dash {
            slug.push('-');
            prev_dash = true;
        }
    }

    slug.trim_matches('-').to_string()
}

async fn upload(
    headers: HeaderMap,
    mut payload: Multipart,
) -> Result<impl IntoResponse, YamafError> {
    let mut responses = Vec::new();

    let wants_html = headers
        .get("accept")
        .and_then(|h| h.to_str().ok())
        .map(|a| a.contains("text/html"))
        .unwrap_or(false);

    if let Ok(ref key) = CONFIG.key {
        if let Some(field) = payload.next_field().await.unwrap() {
            if field.name() == Some("key") {
                let bytes = field
                    .bytes()
                    .await
                    .map_err(|e| YamafError::BadRequest(format!("Error reading key: {e}")))?;

                let s = String::from_utf8(bytes.to_vec())
                    .map_err(|_| YamafError::InternalError("Invalid key format".into()))?;

                if s != *key {
                    return Err(YamafError::BadRequest("Wrong key".into()));
                }
            } else {
                return Err(YamafError::BadRequest("Missing key".into()));
            }
        } else {
            return Err(YamafError::BadRequest("Missing key".into()));
        }
    }

    while let Some(mut field) = payload.next_field().await.unwrap() {
        if field.name() == Some("file") {
            let filename = field
                .file_name()
                .map_or(format!("{}-upload", random(10)), |filename| {
                    format!("{}-{}", random(4), clean_filename(filename))
                });

            let save_path = std::path::Path::new(&CONFIG.root_dir).join(&filename);

            let mut file = fs::File::create(&save_path)
                .await
                .map_err(|_| YamafError::InternalError("Internal i/o error".into()))?;

            let mut written: usize = 0;

            while let Some(chunk) = field
                .chunk()
                .await
                .map_err(|err| YamafError::InternalError(err.to_string()))?
            {
                use tokio::io::AsyncWriteExt;

                written = written
                    .checked_add(chunk.len())
                    .ok_or_else(|| YamafError::BadRequest("File too large".into()))?;

                if written > CONFIG.max_filesize {
                    _ = fs::remove_file(&save_path).await;

                    return Err(YamafError::FileTooBig(filename));
                }

                file.write_all(&chunk)
                    .await
                    .map_err(|_| YamafError::InternalError("Internal i/o error".into()))?;
            }

            let url = format!(
                "{proto}://{host}/{file}",
                proto = CONFIG.external_protocol,
                host = CONFIG.external_host,
                file = filename,
            );

            if wants_html {
                responses.push(format!(
                    r#"<a href="{url}">{url}</a> (size ~ {size:.2}k)"#,
                    url = url,
                    size = written as f64 / 1024.0,
                ));
            } else {
                responses.push(url);
            }
        }
    }

    if responses.is_empty() {
        return Err(YamafError::BadRequest("No files uploaded".into()));
    }

    if wants_html {
        Ok(Html(format!(
            "Here are your file(s):<br>{}",
            responses.join("<br>")
        ))
        .into_response())
    } else {
        Ok(responses.join("\n").into_response())
    }
}

async fn serve_file(Path(filename): Path<String>) -> Result<impl IntoResponse, YamafError> {
    let path = PathBuf::from(&CONFIG.root_dir).join(&filename);

    let metadata = fs::metadata(&path)
        .await
        .map_err(|_| YamafError::FileNotFound)?;
    let file = fs::File::open(&path)
        .await
        .map_err(|_| YamafError::FileNotFound)?;
    let mime = mime_guess::from_path(&path).first_or_octet_stream();

    let content_type = mime
        .as_ref()
        .parse()
        .map_err(|_| YamafError::InternalError("Something went wrong".into()))?;
    let content_length = metadata
        .len()
        .to_string()
        .parse()
        .map_err(|_| YamafError::InternalError("Something went wrong".into()))?;

    let headers = HeaderMap::from_iter([
        (header::CONTENT_TYPE, content_type),
        (header::CONTENT_LENGTH, content_length),
        (header::ACCEPT_RANGES, "bytes".parse().unwrap()),
    ]);

    let stream = ReaderStream::new(file);
    let body = Body::from_stream(stream);

    Ok((StatusCode::OK, headers, body).into_response())
}

#[cfg(feature = "cleanup")]
async fn cleanup_directory() -> Result<(), std::io::Error> {
    let dir = std::path::Path::new(&CONFIG.root_dir);
    let mut entries = fs::read_dir(dir).await?;

    while let Some(entry) = entries.next_entry().await? {
        use std::os::unix::fs::MetadataExt;
        use std::time::SystemTime;

        let path = entry.path();
        let meta = entry.metadata().await?;
        let modified = meta.modified().unwrap_or(SystemTime::UNIX_EPOCH);
        let age = SystemTime::now()
            .duration_since(modified)
            .unwrap_or_default();

        /* min_days + (max_days - min_days) * (1 - size/max_size)^e */
        let retention = (CONFIG.min_filedays as f64
            + (CONFIG.max_filedays - CONFIG.min_filedays) as f64
                * (1.0 - (meta.size() as f64 / CONFIG.max_filesize as f64))
                    .powf(std::f64::consts::E))
            * 24.0
            * 60.0
            * 60.0;

        if meta.is_file() && age.as_secs_f64() > retention {
            fs::remove_file(&path).await?;
            println!("Deleted {:?}", path);
        }
    }

    Ok(())
}