zachery.lol/crates/server/bin/main.rs

use axum::{
    Json, Router,
    http::{HeaderValue, Method, StatusCode, header},
    response::Html,
    routing::{get, post},
};
use message_tools::BoxedMessage;
use serde::Serialize;
use std::{
    fs::OpenOptions,
    io::{BufRead, Read},
    path::{Path, PathBuf},
    time::Duration,
};
use tokio::io::AsyncWriteExt;
use tower_http::{
    cors::CorsLayer, limit::RequestBodyLimitLayer, services::ServeDir,
    set_header::SetResponseHeaderLayer,
};

const PUBKEY_PATH: &str = "./pubkeys/public_keys.json";
const MESSAGE_STORAGE: &str = "./messages";

lazy_static::lazy_static! {
    static ref PUBLIC_KEYS: Vec<Vec<u8>> = load_public_keys(PUBKEY_PATH).expect("static path");
    static ref PUBLISH_LIMITER: std::sync::Mutex<(std::time::Instant, u32)> =
        std::sync::Mutex::new((std::time::Instant::now(), 0));
}

const MAX_PUBLISHES_PER_MINUTE: u32 = 10;

type Error = Box<dyn std::error::Error>;
type Result<T> = std::result::Result<T, Error>;

fn load_public_keys(path: impl AsRef<Path>) -> std::io::Result<Vec<Vec<u8>>> {
    // Open the file
    let f = OpenOptions::new().read(true).write(false).open(path)?;

    // Wrap it in a BufReader for efficient line-by-line reading
    let reader = std::io::BufReader::new(f);

    let mut public_keys = Vec::new();

    // Iterate over each line in the file
    for line_result in reader.lines() {
        let line = line_result?; // Handle potential IO errors during reading

        // Skip empty lines or lines with only whitespace
        let trimmed_line = line.trim();
        if trimmed_line.is_empty() {
            continue;
        }

        // Deserialize the line (which is a single JSON array) into a Vec<u8>
        let key: Vec<u8> = serde_json::from_str(trimmed_line).map_err(|e| {
            std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                format!("Failed to parse JSON line: {}", e),
            )
        })?;

        public_keys.push(key);
    }

    Ok(public_keys)
}

#[tokio::main]
async fn main() -> Result<()> {
    let cors = CorsLayer::new()
        .allow_origin([
            "https://zachery.lol"
                .parse::<axum::http::HeaderValue>()
                .unwrap(),
            "https://www.zachery.lol"
                .parse::<axum::http::HeaderValue>()
                .unwrap(),
        ])
        .allow_methods([Method::GET, Method::POST])
        .allow_headers([axum::http::header::CONTENT_TYPE]);

    // Cache static assets so navigating between pages reuses the streamed field
    // instead of re-downloading it. Scoped to /routes only; HTML routes stay fresh.
    let routes_static = Router::new()
        .fallback_service(ServeDir::new("./routes"))
        .layer(SetResponseHeaderLayer::overriding(
            header::CACHE_CONTROL,
            HeaderValue::from_static("public, max-age=86400"),
        ));

    let app = Router::new()
        .route("/", get(serve_path("./routes/root/index.html")?))
        .route("/about", get(serve_path("./routes/about/index.html")?))
        .route(
            "/projects",
            get(serve_path("./routes/projects/index.html")?),
        )
        .route(
            "/projects/swactor",
            get(serve_path("./routes/projects/swactor/index.html")?),
        )
        .route(
            "/projects/airfrans-neural-cfd-surrogate",
            get(serve_path(
                "./routes/projects/airfrans-neural-cfd-surrogate/index.html",
            )?),
        )
        .route(
            "/projects/cstat-yoke",
            get(serve_path("./routes/projects/cstat-yoke/index.html")?),
        )
        .route("/api/pubkey", get(async || select_key(&PUBLIC_KEYS)))
        .route("/contact", get(serve_path("./routes/contact/index.html")?))
        .route(
            "/contact/message",
            get(serve_path("./routes/contact/message/index.html")?),
        )
        .route(
            "/api/publish",
            post(publish_message).layer(RequestBodyLimitLayer::new(64 * 1024)),
        )
        .nest_service(
            "/gossip-dashboard",
            ServeDir::new("./routes/gossip-dashboard"),
        )
        .nest("/routes", routes_static)
        .layer(cors);

    let addr = "0.0.0.0:3000";
    let listener = tokio::net::TcpListener::bind(addr).await.unwrap();
    println!("Server running at http://{addr}");
    axum::serve(listener, app).await.unwrap();

    Ok(())
}

// Define a struct for the response
#[derive(Serialize)]
struct UserCreated {
    tag: String,
}

/// Strip everything except ASCII alphanumerics, dashes, and underscores
/// so a tag can never escape the storage directory.
fn sanitize_tag(s: &str) -> String {
    s.chars()
        .filter(|c| c.is_ascii_alphanumeric() || *c == '-' || *c == '_')
        .collect()
}

async fn publish_message(Json(msg): Json<BoxedMessage>) -> (StatusCode, Json<UserCreated>) {
    // Rate limit: sliding window of MAX_PUBLISHES_PER_MINUTE
    {
        let mut limiter = PUBLISH_LIMITER.lock().unwrap();
        let now = std::time::Instant::now();
        if now.duration_since(limiter.0) > Duration::from_secs(60) {
            *limiter = (now, 0);
        }
        if limiter.1 >= MAX_PUBLISHES_PER_MINUTE {
            return (
                StatusCode::TOO_MANY_REQUESTS,
                Json(UserCreated { tag: String::new() }),
            );
        }
        limiter.1 += 1;
    }

    let tag = sanitize_tag(&msg.tag());

    if tag.is_empty() {
        return (StatusCode::BAD_REQUEST, Json(UserCreated { tag }));
    }

    let file_path = PathBuf::from(MESSAGE_STORAGE).join(format!("{tag}.json"));

    // 3. Create all necessary parent directories recursively (async operation)
    if let Err(e) = tokio::fs::create_dir_all(&MESSAGE_STORAGE).await {
        eprintln!("Failed to create directory {}: {}", MESSAGE_STORAGE, e);
        return (StatusCode::INTERNAL_SERVER_ERROR, Json(UserCreated { tag }));
    }

    // 4. Open the file asynchronously
    let mut f = match tokio::fs::OpenOptions::new()
        .create(true)
        .write(true)
        .append(true)
        .open(&file_path)
        .await
    {
        Ok(file) => file,
        Err(e) => {
            eprintln!("Failed to open file {}: {}", file_path.display(), e);
            return (StatusCode::INTERNAL_SERVER_ERROR, Json(UserCreated { tag }));
        }
    };

    // 5. Serialize and write the data asynchronously
    // We append a newline character for standard text file formatting
    let s = serde_json::to_string(&msg).expect("failed to serialize boxed message");
    if let Err(e) = f.write_all(format!("{}\n", s).as_bytes()).await {
        eprintln!("Failed to write message to file: {}", e);
        return (StatusCode::INTERNAL_SERVER_ERROR, Json(UserCreated { tag }));
    }

    (StatusCode::CREATED, Json(UserCreated { tag }))
}

fn select_key(pks: &[Vec<u8>]) -> (StatusCode, String) {
    if pks.is_empty() {
        return (
            StatusCode::SERVICE_UNAVAILABLE,
            "no keys available".to_string(),
        );
    }
    let idx = getrandom::u64().expect("random integer") as usize % pks.len();
    (
        StatusCode::OK,
        serde_json::to_string(&pks[idx]).expect("serialize vec of bytes"),
    )
}

fn serve_path(path: impl AsRef<std::path::Path>) -> Result<Html<String>> {
    let mut s = String::new();
    OpenOptions::new()
        .read(true)
        .write(false)
        .open(path)?
        .read_to_string(&mut s)?;

    Ok(Html(s))
}