diff --git a/Cargo.toml b/Cargo.toml index 3c57022..0786031 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,7 +12,7 @@ edition = "2018" [dependencies] tokio-rustls = "0.22.0" -tokio = { version = "1.0", features = ["fs", "io-util", "net", "rt-multi-thread"] } +tokio = { version = "1.0", features = ["fs", "io-util", "net", "rt-multi-thread", "sync"] } env_logger = { version = "0.8", default-features = false, features = ["atty", "humantime", "termcolor"] } getopts = "0.2.21" log = "0.4" diff --git a/README.md b/README.md index fe050e7..68af2d4 100644 --- a/README.md +++ b/README.md @@ -50,9 +50,8 @@ The lines of the file should have this format: : ``` -Where `` is just a filename (not a path) of a file in the same directory, and `` is the metadata to be stored. -Lines that start with optional whitespace and `#` are ignored, as are lines that do not fit the above basic format. -Both parts are stripped of any leading and/or trailing whitespace. +Where `` is just a filename (not a path) of a file in the same directory, and `` is the MIME type to be stored. If `` starts with a semicolon, agate will use the usual mechanism to determine the mime type of the file and append the specified parameters. +Lines that start with optional whitespace and `#` are ignored, as are lines that do not contain a `:`. Both parts are stripped of any leading and/or trailing whitespace. [Gemini]: https://gemini.circumlunar.space/ [Rust]: https://www.rust-lang.org/ diff --git a/src/main.rs b/src/main.rs index 3f1efe1..cb9a494 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,5 @@ mod metadata; +use metadata::FileOptions; use { once_cell::sync::Lazy, @@ -21,6 +22,7 @@ use { io::{AsyncReadExt, AsyncWriteExt}, net::{TcpListener, TcpStream}, runtime::Runtime, + sync::RwLock, }, tokio_rustls::{server::TlsStream, TlsAcceptor}, url::{Host, Url}, @@ -31,12 +33,19 @@ fn main() -> Result { env_logger::Builder::new().parse_filters("info").init(); } Runtime::new()?.block_on(async { + let mimetypes = Arc::new(RwLock::new(FileOptions::new( + &ARGS + .language + .as_ref() + .map_or(String::new(), |lang| format!(";lang={}", lang)), + ))); let listener = TcpListener::bind(&ARGS.addrs[..]).await?; log::info!("Listening on {:?}...", ARGS.addrs); loop { let (stream, _) = listener.accept().await?; + let arc = mimetypes.clone(); tokio::spawn(async { - if let Err(e) = handle_request(stream).await { + if let Err(e) = handle_request(stream, arc).await { log::error!("{:?}", e); } }); @@ -147,11 +156,11 @@ fn check_path(s: String) -> Result { } /// Handle a single client session (request + response). -async fn handle_request(stream: TcpStream) -> Result { +async fn handle_request(stream: TcpStream, mimetypes: Arc>) -> Result { let stream = &mut TLS.accept(stream).await?; match parse_request(stream).await { - Ok(url) => send_response(url, stream).await?, + Ok(url) => send_response(url, stream, mimetypes).await?, Err((status, msg)) => send_header(stream, status, &[msg]).await?, } stream.shutdown().await?; @@ -223,7 +232,11 @@ async fn parse_request( } /// Send the client the file located at the requested URL. -async fn send_response(url: Url, stream: &mut TlsStream) -> Result { +async fn send_response( + url: Url, + stream: &mut TlsStream, + mimetypes: Arc>, +) -> Result { let mut path = std::path::PathBuf::from(&ARGS.content_dir); if let Some(segments) = url.path_segments() { for segment in segments { @@ -265,12 +278,21 @@ async fn send_response(url: Url, stream: &mut TlsStream) -> Result { }; // Send header. - if path.extension() == Some(OsStr::new("gmi")) { - send_text_gemini_header(stream).await?; + let mut locked = mimetypes.write().await; + let data = locked.get(&path); + if data.is_empty() || data.starts_with(";") { + // guess MIME type + if path.extension() == Some(OsStr::new("gmi")) { + send_header(stream, 20, &["text/gemini", data]).await?; + } else { + let mime = mime_guess::from_path(&path).first_or_octet_stream(); + send_header(stream, 20, &[mime.essence_str(), data]).await?; + }; } else { - let mime = mime_guess::from_path(&path).first_or_octet_stream(); - send_header(stream, 20, &[mime.essence_str()]).await?; + // this must be a full MIME type + send_header(stream, 20, &[data]).await?; } + drop(locked); // Send body. tokio::io::copy(&mut file, stream).await?; @@ -284,7 +306,7 @@ async fn list_directory(stream: &mut TlsStream, path: &Path) -> Resul .add(b'?').add(b'`').add(b'{').add(b'}'); log::info!("Listing directory {:?}", path); - send_text_gemini_header(stream).await?; + send_header(stream, 20, &["text/gemini"]).await?; let mut entries = tokio::fs::read_dir(path).await?; let mut lines = vec![]; while let Some(entry) = entries.next_entry().await? { @@ -321,11 +343,3 @@ async fn send_header(stream: &mut TlsStream, status: u8, meta: &[&str stream.write_all(response.as_bytes()).await?; Ok(()) } - -async fn send_text_gemini_header(stream: &mut TlsStream) -> Result { - if let Some(lang) = ARGS.language.as_deref() { - send_header(stream, 20, &["text/gemini;lang=", lang]).await - } else { - send_header(stream, 20, &["text/gemini"]).await - } -}