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> = load_public_keys(PUBKEY_PATH).expect("static path"); } type Error = Box; type Result = std::result::Result; type UrlCache = Arc>>; fn load_public_keys(path: impl AsRef) -> std::io::Result>> { // 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 let key: Vec = 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, #[serde(skip_serializing_if = "Option::is_none")] title: Option, #[serde(skip_serializing_if = "Option::is_none")] thumbnail: Option, #[serde(skip_serializing_if = "Option::is_none")] duration: Option, #[serde(skip_serializing_if = "Option::is_none")] stream: Option, #[serde(skip_serializing_if = "Option::is_none")] error: Option, } async fn resolve_handler( State(cache): State, Query(params): Query, ) -> Json { 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", ¶ms.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, AxumPath(id): AxumPath, headers: HeaderMap, ) -> axum::response::Response { 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) -> (StatusCode, Json) { 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>) -> 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) -> Result> { let mut s = String::new(); OpenOptions::new() .read(true) .write(false) .open(path)? .read_to_string(&mut s)?; Ok(Html(s)) }