diff --git a/src/client.rs b/src/client.rs index e4506e3..9592600 100644 --- a/src/client.rs +++ b/src/client.rs @@ -130,7 +130,7 @@ impl Client { debug!("File '{}': {} bytes remaining", file.hash, remaining); } - let check_hash = crypto::try_hash(&path)?; + let check_hash = crypto::try_hash(&path).unwrap(); let msg = check_hash.as_bytes().to_vec(); handler.send(&msg).await?; diff --git a/src/crypto.rs b/src/crypto.rs index 133dc78..e9596ea 100644 --- a/src/crypto.rs +++ b/src/crypto.rs @@ -102,7 +102,7 @@ impl Crypto { } } -pub fn try_hash(path: &Path) -> Result> { +pub fn try_hash(path: &Path) -> Result> { debug!("Calculating SHA hash"); let hash = sha256::try_digest(path)?; diff --git a/src/main.rs b/src/main.rs index 691e668..f654ec4 100755 --- a/src/main.rs +++ b/src/main.rs @@ -6,9 +6,10 @@ use contego::{ client::Client, parser::{addr_parser, dirpath_parser, filepath_parser}, server::Server, - util::{filepaths, metadata, Ip}, + util::{filepaths, handle_exit, metadata, Ip}, }; use env_logger::Env; +use log::error; use tokio::sync::mpsc; #[derive(Debug, Parser)] @@ -20,16 +21,16 @@ struct Cli { #[derive(Debug, Subcommand)] enum Commands { - #[clap(group(ArgGroup::new("input").required(true).args(&["infile", "files"])))] + #[clap(group(ArgGroup::new("input").required(true).args(&["source", "files"])))] Host { /// Access key #[clap(short = 'k', long)] key: String, /// Path to a source file (alternative to --files) - #[clap(short = 'i', long, value_parser = filepath_parser, conflicts_with = "files", group = "input")] + #[clap(short = 's', long, value_parser = filepath_parser, conflicts_with = "files", group = "input")] source: Option, /// Paths to shareable files (alternative to --source) - #[clap(short = 'f', long, num_args = 1.., value_parser = filepath_parser, conflicts_with = "infile", group = "input")] + #[clap(short = 'f', long, num_args = 1.., value_parser = filepath_parser, conflicts_with = "source", group = "input")] files: Option>, /// Host port #[clap(short = 'p', long, default_value_t = 8080)] @@ -58,9 +59,8 @@ enum Commands { } #[tokio::main] -async fn main() -> Result<(), Box> { +async fn main() -> Result<(), Box> { env_logger::Builder::from_env(Env::default().default_filter_or("info")).init(); - let cli = Cli::parse(); match cli.command { @@ -73,23 +73,33 @@ async fn main() -> Result<(), Box> { local, key, } => { - let files = filepaths(source, files)?; - let (metadata, index) = metadata(&files).await?; + let (tx, rx) = mpsc::channel::<()>(1); + + let paths = filepaths(source, files)?; + let (metadata, index) = metadata(&paths).await?; let addr = match (local, ipv6) { (true, _) => Ip::Local.fetch(port)?, (false, true) => Ip::V6.fetch(port)?, (false, false) => Ip::V4.fetch(port)?, }; - // TODO: handle shutdown signal - let (_tx, rx) = mpsc::channel::<()>(1); - let server = Server::new(addr, key, chunksize, metadata, index); - server.start(rx).await?; + + tokio::spawn(async move { + match server.start(rx).await { + Ok(_) => {} + Err(e) => error!("Error during server execution: {}", e), + }; + }); + + handle_exit(tx).await?; } Commands::Connect { addr, out, key } => { let client = Client::new(addr, key, out); - client.connection().await?; + match client.connection().await { + Ok(_) => {} + Err(e) => error!("Error during client execution: {}", e), + }; } }; diff --git a/src/parser.rs b/src/parser.rs index 998a69e..6c2c6b1 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -1,9 +1,12 @@ use std::{ + env, io::{Error, ErrorKind::NotFound}, net::{AddrParseError, SocketAddr}, path::PathBuf, }; +use log::debug; + pub fn addr_parser(addr: &str) -> Result { let addr = addr .parse::() @@ -13,7 +16,13 @@ pub fn addr_parser(addr: &str) -> Result { } pub fn filepath_parser(path: &str) -> Result { - let path = path.parse::().expect("Failed to parse path"); + debug!("Validating filepath '{}'", path); + + let home = env::var("HOME").unwrap(); + let path = path + .replace('~', &home) + .parse::() + .expect("Failed to parse path"); if path.exists() && path.is_file() { Ok(path) @@ -23,6 +32,8 @@ pub fn filepath_parser(path: &str) -> Result { } pub fn dirpath_parser(path: &str) -> Result { + debug!("Validating dirpath '{}'", path); + let path = path.parse::().expect("Failed to parse path"); if path.exists() && path.is_dir() { diff --git a/src/util.rs b/src/util.rs index 3e43330..d847247 100644 --- a/src/util.rs +++ b/src/util.rs @@ -1,7 +1,15 @@ -use std::{collections::HashMap, error::Error, fs, net::SocketAddr, path::PathBuf}; +use std::{ + collections::HashMap, + env, + error::Error, + fs, + io::{stdin, Read}, + net::SocketAddr, + path::PathBuf, +}; use log::{debug, info}; -use tokio::{fs::File, io::BufWriter}; +use tokio::{fs::File, io::BufWriter, sync::mpsc::Sender}; use crate::crypto; @@ -16,9 +24,7 @@ pub enum Ip { } impl Ip { - pub fn fetch(self, port: u16) -> Result> { - debug!("Fetching IP information"); - + pub fn fetch(self, port: u16) -> Result> { let addr = match self { Ip::V4 => PUBLIC_IPV4, Ip::V6 => PUBLIC_IPV6, @@ -28,6 +34,8 @@ impl Ip { } }; + info!("Fetching IP information from {}", addr); + let res = format!("{}:{}", ureq::get(addr).call()?.into_string()?.trim(), port); let addr = res.parse::()?; @@ -51,38 +59,41 @@ impl FileInfo { } pub fn filepaths( - infile: Option, + source: Option, files: Option>, -) -> Result, Box> { +) -> Result, Box> { info!("Collecting filepaths"); - let mut filepaths = Vec::new(); + let mut paths = Vec::new(); - if let Some(infile) = infile { - let paths = fs::read_to_string(infile)?; - for path in paths.lines() { - filepaths.push(PathBuf::from(path)); - } + if let Some(source) = source { + let home = env::var("HOME")?; + let content = fs::read_to_string(source)?; + paths = content + .lines() + .into_iter() + .map(|p| PathBuf::from(p.replace('~', &home))) + .collect(); } else if let Some(files) = files { - for file in files { - filepaths.push(file); - } + paths = files; } - debug!("Filepaths collection finished (total: {})", filepaths.len()); + debug!("Filepaths collection finished (total: {})", paths.len()); - Ok(filepaths) + Ok(paths) } pub async fn metadata( files: &Vec, -) -> Result<(Vec, HashMap), Box> { +) -> Result<(Vec, HashMap), Box> { info!("Collecting metadata"); let mut metadata = Vec::new(); let mut index = HashMap::new(); for path in files { + debug!("Collecting '{}' metadata", path.to_str().unwrap()); + let split = path.to_str().unwrap().split('/').collect::>(); let name = split[split.len() - 1].to_string(); let handle = File::open(path).await?; @@ -90,8 +101,6 @@ pub async fn metadata( let hash = crypto::try_hash(path)?; if size > 0 { - debug!("Collecting '{}' metadata", name); - let info = FileInfo::new(name, size, hash.clone()); metadata.push(info); index.insert(hash, path.clone()); @@ -117,3 +126,21 @@ pub async fn new_file( Ok((BufWriter::new(handle), path)) } + +pub async fn handle_exit(tx: Sender<()>) -> Result<(), Box> { + let mut stdin = stdin().lock().bytes(); + + loop { + let k = match stdin.next() { + None => continue, + Some(k) => k?, + }; + + if k == b'q' { + tx.send(()).await?; + break; + } + } + + Ok(()) +}