From 08c437920c175a2a457143cbdf5d7fce33194200 Mon Sep 17 00:00:00 2001 From: yggverse Date: Fri, 21 Feb 2025 22:18:07 +0200 Subject: [PATCH] use `titanite` library for backend --- Cargo.toml | 7 ++- README.md | 1 + src/argument.rs | 4 ++ src/header.rs | 56 ----------------- src/main.rs | 161 ++++++++++++++++++++++-------------------------- src/storage.rs | 7 ++- 6 files changed, 89 insertions(+), 147 deletions(-) delete mode 100644 src/header.rs diff --git a/Cargo.toml b/Cargo.toml index d7eff04..ae8a3e4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "titanit" -version = "0.1.0" +version = "0.1.1" edition = "2021" license = "MIT" readme = "README.md" @@ -14,3 +14,8 @@ anyhow = "1.0.95" clap = { version = "4.5.30", features = ["derive"] } native-tls = "0.2.13" regex = "1.11.1" +titanite = "0.1.1" + +# development +# [patch.crates-io] +# titanite = { path = "../titanite" } diff --git a/README.md b/README.md index 66520fd..9dcd4be 100644 --- a/README.md +++ b/README.md @@ -34,6 +34,7 @@ openssl pkcs12 -export -out server.pfx -inkey server.pem -in server.crt * `--bind`, `-b` required, server `host:port` to listen incoming connections * `--identity`, `-i` required, filepath to server identity in PKCS (PFX) format * `--password`, `-p` optional, unlock encrypted `identity` by passphrase +* `--chunk`, `-c` optional, buffer chunk size (`1024` by default) * `--size`, `-s` optional, max size limit in bytes (unlimited by default) * `--mime`, `-m` optional, uploads MIME type whitelist (comma separated, all by default) * `--directory`, `-d` optional, uploads target directory (`public` by default) diff --git a/src/argument.rs b/src/argument.rs index adca508..1f6fbe4 100644 --- a/src/argument.rs +++ b/src/argument.rs @@ -23,6 +23,10 @@ pub struct Argument { #[arg(short, long)] pub size: Option, + /// Buffer chunk size (`1024` by default) + #[arg(short, long, default_value_t = 1024)] + pub chunk: usize, + /// Uploads MIME type allowed (comma separated, all by default) /// * based on headers #[arg(short, long)] diff --git a/src/header.rs b/src/header.rs deleted file mode 100644 index 71fb30e..0000000 --- a/src/header.rs +++ /dev/null @@ -1,56 +0,0 @@ -use anyhow::{bail, Result}; -use native_tls::TlsStream; -use std::{io::Read, net::TcpStream}; - -pub struct Header { - pub mime: Option, - pub size: usize, -} - -impl Header { - pub fn for_stream(stream: &mut TlsStream) -> Result { - let mut header: Vec = Vec::new(); - let mut buffer = vec![0]; - loop { - match stream.read(&mut buffer) { - Ok(0) => bail!("Invalid protocol"), - Ok(_) => { - if header.len() > 1024 { - bail!("Invalid length") - } - if buffer[0] == b'\r' { - continue; - } - if buffer[0] == b'\n' { - break; - } - header.push(buffer[0]) - } - Err(e) => bail!(e), - } - } - Self::from_bytes(header) - } - - fn from_bytes(buffer: Vec) -> Result { - use anyhow::bail; - use regex::Regex; - let header = String::from_utf8(buffer)?; - Ok(Self { - mime: match Regex::new(r"mime=([^\/]+\/[^\s;]+)")?.captures(&header) { - Some(c) => c.get(1).map(|v| v.as_str().to_string()), - None => None, - }, - size: match Regex::new(r"size=(\d+)")?.captures(&header) { - Some(c) => match c.get(1) { - Some(v) => match v.as_str().parse::() { - Ok(s) => s, - Err(e) => bail!(e), - }, - None => bail!("Size required"), - }, - None => bail!("Size required"), - }, - }) - } -} diff --git a/src/main.rs b/src/main.rs index a104e09..a4c544f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,11 +1,9 @@ mod argument; -mod header; mod storage; use anyhow::Result; use argument::Argument; use clap::Parser; -use header::Header; use native_tls::{Identity, TlsAcceptor, TlsStream}; use std::{ fs::File, @@ -15,6 +13,7 @@ use std::{ thread, time::{SystemTime, UNIX_EPOCH}, }; +use titanite::response::success::Default; fn main() -> Result<()> { let argument = Arc::new(Argument::parse()); @@ -51,106 +50,92 @@ fn main() -> Result<()> { } fn handle(argument: Arc, peer: SocketAddr, mut stream: TlsStream) { + use titanite::{request::Titan, response::*}; println!("[{}] [info] [{peer}] New connection", now()); - match Header::for_stream(&mut stream) { - Ok(header) => { - // do not trust header values, but check it to continue - if argument.size.is_some_and(|s| header.size > s) { - println!("[{}] [error] [{peer}] Max size limit reached", now()); - return; - } - // make sure mime type whitelisted - if argument - .mime - .as_ref() - .is_some_and(|m| header.mime.is_some_and(|h| m.contains(&h))) - { - println!("[{}] [error] [{peer}] MIME type not allowed", now()); - return; - } - // begin data handle - let mut total = 0; - let mut buffer = vec![0]; // @TODO optional chunk size - match storage::Item::create(&argument.directory) { - Ok(mut item) => { + + // read header bytes + let mut input = vec![0; 1024]; + match stream.read(&mut input) { + Ok(0) => todo!("Canceled"), + Ok(l) => { + match Titan::from_bytes(&input[..l]) { + Ok(titan) => { + // init memory pool + let mut data: Vec = Vec::with_capacity(titan.size); loop { - match stream.read(&mut buffer) { + // read data bytes + let mut input = vec![0; argument.chunk]; + match stream.read(&mut input) { Ok(0) => { - item.delete().unwrap(); - println!( - "[{}] [warning] [{peer}] Connection closed with rollback", - now() - ); - break; + todo!("Canceled") } - Ok(n) => { - total += n; // validate real data size - if argument.size.is_some_and(|s| s > total) { - item.delete().unwrap(); - println!("[{}] [error] [{peer}] Max size limit reached", now()); - break; + Ok(l) => { + data.extend(&input[..l]); + + let total = data.len(); + if argument.size.is_some_and(|limit| total > limit) { + todo!("Allowed max length limit reached") } - if total > header.size { - item.delete().unwrap(); - println!( - "[{}] [error] [{peer}] Content size larger than declared in headers", now() - ); - break; - } - if let Err(e) = item.file.write(&buffer) { - item.delete().unwrap(); - println!("[{}] [error] [{peer}] Failed to write file from stream: {e}", now()); - break; - } - // transfer completed - if total == header.size { - match &stream.write_all( - // https://geminiprotocol.net/docs/protocol-specification.gmi#status-31-permanent-redirection - format!( - "31 {}\r\n", - argument - .redirect - .clone() - .unwrap_or(format!("gemini://{}", argument.bind)) - ) - .as_bytes(), - ) { - Ok(_) => { - item.commit().unwrap(); // @TODO detect/cache mime based on content type received - println!("[{}] [info] [{peer}] Success", now()); - // @TODO close connection gracefully - // https://geminiprotocol.net/docs/protocol-specification.gmi#closing-connections - match stream.flush() { - Ok(_) => println!("[{}] [info] [{peer}] Connection closed by server.", now()), - Err(e) => println!("[{}] [error] [{peer}] {e}", now()) - }; - } - Err(e) => { - item.delete().unwrap(); - println!( - "[{}] [error] [{peer}] Failed write to stream: {e}", - now() - ) - } + if titan.size >= total { + if titan.size > total { + println!( + "[{}] [warning] [{peer}] Data size mismatch header declaration", + now() + ); + return; + } + // success + match storage::Item::create(&argument.directory) { + Ok(mut tmp) => match tmp.file.write(titan.data) { + Ok(_) => match &stream.write_all( + &Success::Default(Default { + mime: "text/gemini".to_string(), + }) + .into_bytes(), + ) { + // @TODO detect/validate/cache mime based on data received + Ok(()) => match tmp.commit() { + Ok(path) => match stream.flush() { + // Close connection gracefully + // https://geminiprotocol.net/docs/protocol-specification.gmi#closing-connections + Ok(()) => match stream.shutdown() { + Ok(()) => println!( + "[{}] [info] [{peer}] Data saved to {path}", + now() + ), + Err(e) => println!("[{}] [warning] [{peer}] {e}", now()) + }, + Err(e) => println!("[{}] [warning] [{peer}] {e}", now()) + } + Err(e) => println!("[{}] [warning] [{peer}] {e}", now()) + }, + Err(e) => { + println!("[{}] [error] [{peer}] {e}", now()); + if let Err(e) = tmp.delete() { + println!("[{}] [error] [{peer}] {e}", now()) + } + } + }, + Err(e) => { + println!("[{}] [error] [{peer}] {e}", now()); + if let Err(e) = tmp.delete() { + println!("[{}] [error] [{peer}] {e}", now()) + } + } + }, + Err(e) => println!("[{}] [error] [{peer}] {e}", now()), } break; } } - Err(e) => { - item.delete().unwrap(); - println!( - "[{}] [error] [{peer}] Failed read from stream: {e}", - now() - ); - break; - } + Err(e) => todo!("{e}"), } } } - Err(e) => println!("[{}] [error] [{peer}] Could not create Item: {e}", now()), + Err(e) => todo!("{e}"), } } - Err(e) => println!("[{}] [error] [{peer}] Header error: {e}", now()), + Err(e) => todo!("{e}"), } } diff --git a/src/storage.rs b/src/storage.rs index 2371de5..c006f30 100644 --- a/src/storage.rs +++ b/src/storage.rs @@ -63,10 +63,13 @@ impl Item { // Actions /// Take object processed, commit its changes - pub fn commit(self) -> Result<()> { + pub fn commit(self) -> Result { match self.path.to_str() { Some(old) => match old.strip_suffix(TMP_SUFFIX) { - Some(new) => Ok(rename(old, new)?), + Some(new) => { + rename(old, new)?; + Ok(new.to_string()) + } None => bail!("Invalid TMP suffix"), // | panic }, None => bail!("Invalid Item path"), // | panic