diff --git a/src/main.rs b/src/main.rs index a4c544f..a9e4590 100644 --- a/src/main.rs +++ b/src/main.rs @@ -13,7 +13,6 @@ use std::{ thread, time::{SystemTime, UNIX_EPOCH}, }; -use titanite::response::success::Default; fn main() -> Result<()> { let argument = Arc::new(Argument::parse()); @@ -50,95 +49,134 @@ fn main() -> Result<()> { } fn handle(argument: Arc, peer: SocketAddr, mut stream: TlsStream) { - use titanite::{request::Titan, response::*}; + use titanite::*; println!("[{}] [info] [{peer}] New connection", now()); // 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 { - // read data bytes - let mut input = vec![0; argument.chunk]; - match stream.read(&mut input) { - Ok(0) => { - todo!("Canceled") - } - Ok(l) => { - data.extend(&input[..l]); + Ok(0) => println!("[{}] [warning] [{peer}] Connection closed by peer", now()), + Ok(l) => match request::Titan::from_bytes(&input[..l]) { + Ok(titan) => { + // init memory pool + let mut data: Vec = Vec::with_capacity(titan.size); + loop { + // read data bytes + let mut input = vec![0; argument.chunk]; + match stream.read(&mut input) { + Ok(0) => { + println!("[{}] [warning] [{peer}] Connection closed by peer", now()) + } + 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 titan.size >= total { - if titan.size > total { - println!( - "[{}] [warning] [{peer}] Data size mismatch header declaration", - now() - ); - return; + // calculate once + let total = data.len(); + + // validate server-side limits + if argument.size.is_some_and(|limit| total > limit) { + const MESSAGE: &str = "Allowed max length limit reached"; + return send( + &response::failure::permanent::BadRequest { + message: Some(MESSAGE.to_string()), } - // 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()) - } + .into_bytes(), + stream, + |result| match result { + Ok(()) => { + println!("[{}] [warning] [{peer}] {MESSAGE}", now()) + } + Err(e) => println!("[{}] [error] [{peer}] {e}", now()), + }, + ); + } + + // all expected data received + if titan.size >= total { + // validate client-side limits + if titan.size > total { + const MESSAGE: &str = "Data size mismatch header declaration"; + return send( + &response::failure::permanent::BadRequest { + message: Some(MESSAGE.to_string()), + } + .into_bytes(), + stream, + |result| match result { + Ok(()) => { + println!("[{}] [warning] [{peer}] {MESSAGE}", now()) + } + Err(e) => println!("[{}] [error] [{peer}] {e}", now()), + }, + ); + } + + // @TODO detect/validate/cache mime based on data received + + // success + match storage::Item::create(&argument.directory) { + Ok(mut tmp) => match tmp.file.write(titan.data) { + Ok(_) => match tmp.commit() { + Ok(pmt) => send( + &response::redirect::Permanent { + target: pmt.path.to_str().unwrap().to_owned(), } - }, - Err(e) => { + .into_bytes(), + stream, + |result| match result { + Ok(()) => println!( + "[{}] [info] [{peer}] Data saved to {}", + now(), + pmt.path.to_string_lossy() + ), + Err(e) => { + println!( + "[{}] [warning] [{peer}] {e}", + now() + ) + } + }, + ), + Err((tmp, e)) => { println!("[{}] [error] [{peer}] {e}", now()); if let Err(e) = tmp.delete() { - println!("[{}] [error] [{peer}] {e}", now()) + println!("[{}] [error] [{peer}] {e}", now()); } } }, - Err(e) => println!("[{}] [error] [{peer}] {e}", now()), - } - break; + 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) => todo!("{e}"), } + Err(e) => todo!("{e}"), } } - Err(e) => todo!("{e}"), } - } + Err(e) => todo!("{e}"), + }, Err(e) => todo!("{e}"), } } +fn send(data: &[u8], mut stream: TlsStream, callback: impl FnOnce(Result<()>)) { + callback((|| { + stream.write_all(data)?; + stream.flush()?; + // Close connection gracefully + // https://geminiprotocol.net/docs/protocol-specification.gmi#closing-connections + stream.shutdown()?; + Ok(()) + })()); +} + fn now() -> u128 { SystemTime::now() .duration_since(UNIX_EPOCH) diff --git a/src/storage.rs b/src/storage.rs index c006f30..c89d6a7 100644 --- a/src/storage.rs +++ b/src/storage.rs @@ -1,7 +1,8 @@ -use anyhow::{bail, Result}; +use anyhow::{anyhow, Error, Result}; use std::{ fs::{create_dir_all, remove_file, rename, File}, path::{PathBuf, MAIN_SEPARATOR}, + str::FromStr, thread, time::{Duration, SystemTime, UNIX_EPOCH}, }; @@ -62,21 +63,26 @@ impl Item { // Actions - /// Take object processed, commit its changes - pub fn commit(self) -> Result { + /// Commit changes, return permanent Self + pub fn commit(mut self) -> Result { match self.path.to_str() { Some(old) => match old.strip_suffix(TMP_SUFFIX) { - Some(new) => { - rename(old, new)?; - Ok(new.to_string()) - } - None => bail!("Invalid TMP suffix"), // | panic + Some(new) => match rename(old, new) { + Ok(()) => { + self.path = match PathBuf::from_str(new) { + Ok(path) => path, + Err(e) => return Err((self, anyhow!(e))), + }; + Ok(self) + } + Err(e) => Err((self, anyhow!(e))), + }, + None => Err((self, anyhow!("Unexpected suffix"))), }, - None => bail!("Invalid Item path"), // | panic + None => Err((self, anyhow!("Unexpected file path"))), } } - /// Cleanup object relationships pub fn delete(self) -> Result<()> { Ok(remove_file(self.path)?) }