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

use axum::{
    Json, Router,
    body::Body,
    extract::{Path as AxumPath, Query, State},
    http::{HeaderMap, StatusCode},
    response::Html,
    routing::{get, post},
};
use futures::TryStreamExt;
use message_tools::BoxedMessage;
use serde::{Deserialize, Serialize};
use std::{
    collections::HashMap,
    fs::OpenOptions,
    io::{BufRead, Read},
    path::{Path, PathBuf},
    sync::Arc,
};
use tokio::io::AsyncWriteExt;
use tokio::sync::RwLock;
use tower_http::services::ServeDir;

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");
}

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

type UrlCache = Arc<RwLock<HashMap<String, String>>>;

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)
}

#[derive(Deserialize)]
struct ResolveQuery {
    url: String,
}

#[derive(Serialize)]
struct ResolveResponse {
    #[serde(skip_serializing_if = "Option::is_none")]
    id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    title: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    thumbnail: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    duration: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    stream: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    error: Option<String>,
}

async fn resolve_handler(
    State(cache): State<UrlCache>,
    Query(params): Query<ResolveQuery>,
) -> Json<ResolveResponse> {
    let output = tokio::process::Command::new("yt-dlp")
        .args([
            "-j",
            "-f", "best[ext=mp4][protocol=https]/best[ext=mp4][protocol=http]/best[protocol=https]/best[protocol=http]/best",
            "--no-playlist",
            &params.url,
        ])
        .output()
        .await;

    let output = match output {
        Ok(o) => o,
        Err(e) => {
            return Json(ResolveResponse {
                id: None,
                title: None,
                thumbnail: None,
                duration: None,
                stream: None,
                error: Some(format!("Failed to run yt-dlp: {e}")),
            });
        }
    };

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        return Json(ResolveResponse {
            id: None,
            title: None,
            thumbnail: None,
            duration: None,
            stream: None,
            error: Some(format!("yt-dlp failed: {stderr}")),
        });
    }

    let info: serde_json::Value = match serde_json::from_slice(&output.stdout) {
        Ok(v) => v,
        Err(e) => {
            return Json(ResolveResponse {
                id: None,
                title: None,
                thumbnail: None,
                duration: None,
                stream: None,
                error: Some(format!("Failed to parse yt-dlp output: {e}")),
            });
        }
    };

    let video_id = info["id"].as_str().unwrap_or("unknown").to_string();
    let title = info["title"].as_str().map(|s| s.to_string());
    let thumbnail = info["thumbnail"].as_str().map(|s| s.to_string());
    let duration = info["duration"].as_f64();
    let stream_url = info["url"].as_str().unwrap_or("").to_string();

    if stream_url.is_empty() {
        return Json(ResolveResponse {
            id: None,
            title: None,
            thumbnail: None,
            duration: None,
            stream: None,
            error: Some("No stream URL found".to_string()),
        });
    }

    cache.write().await.insert(video_id.clone(), stream_url);

    Json(ResolveResponse {
        id: Some(video_id.clone()),
        title,
        thumbnail,
        duration,
        stream: Some(format!("/music/api/stream/{video_id}")),
        error: None,
    })
}

async fn stream_handler(
    State(cache): State<UrlCache>,
    AxumPath(id): AxumPath<String>,
    headers: HeaderMap,
) -> axum::response::Response<Body> {
    let stream_url = {
        let map = cache.read().await;
        map.get(&id).cloned()
    };

    let stream_url = match stream_url {
        Some(u) => u,
        None => {
            return axum::response::Response::builder()
                .status(StatusCode::NOT_FOUND)
                .body(Body::from("Stream not found"))
                .unwrap();
        }
    };

    let client = reqwest::Client::new();
    let mut req = client.get(&stream_url);

    if let Some(range) = headers.get("range") {
        req = req.header("Range", range);
    }

    let upstream = match req.send().await {
        Ok(r) => r,
        Err(e) => {
            return axum::response::Response::builder()
                .status(StatusCode::BAD_GATEWAY)
                .body(Body::from(format!("Upstream error: {e}")))
                .unwrap();
        }
    };

    let status = upstream.status();
    let mut response = axum::response::Response::builder().status(status.as_u16());

    for key in ["content-type", "content-length", "content-range", "accept-ranges"] {
        if let Some(val) = upstream.headers().get(key) {
            response = response.header(key, val);
        }
    }

    let stream = upstream
        .bytes_stream()
        .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e));
    let body = Body::from_stream(stream);

    response.body(body).unwrap()
}

#[tokio::main]
async fn main() -> Result<()> {
    let url_cache: UrlCache = Arc::new(RwLock::new(HashMap::new()));

    let music_routes = Router::new()
        .route("/", get(serve_path("./routes/music/index.html")?))
        .route("/api/resolve", get(resolve_handler))
        .route("/api/stream/{id}", get(stream_handler))
        .with_state(url_cache);

    // build our application with a single route
    let app = Router::new()
        .route("/", get(serve_path("./routes/root/index.html")?))
        .route("/who", get(serve_path("./routes/who/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))
        .nest("/music", music_routes)
        .nest_service("/routes", ServeDir::new("./routes"));

    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,
}

// FIXME: sanitation of boxed messages
async fn publish_message(Json(msg): Json<BoxedMessage>) -> (StatusCode, Json<UserCreated>) {
    let tag = msg.tag(); // Capture tag early

    // 1. Define the full *file* path: e.g., "MESSAGE_STORAGE/some_tag/.json"
    let file_path = PathBuf::from(MESSAGE_STORAGE)
        .join(format!("{tag}.json")); // The file inside that directory

    // 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<Vec<u8>>) -> String {
    let idx = getrandom::u64().expect("random integer") as usize % pks.len();
    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))
}