diff --git a/src/config.rs b/src/config.rs index ab362f8..76b18e5 100644 --- a/src/config.rs +++ b/src/config.rs @@ -25,6 +25,10 @@ pub struct Config { pub port: u16, /// Default directory of the TFTP Server. (default: current working directory) pub directory: PathBuf, + /// Upload directory of the TFTP Server. (default: directory) + pub receive_directory: PathBuf, + /// Download directory of the TFTP Server. (default: directory) + pub send_directory: PathBuf, /// Use a single port for both sending and receiving. (default: false) pub single_port: bool, /// Refuse all write requests, making the server read-only. (default: false) @@ -43,6 +47,8 @@ impl Config { ip_address: Ipv4Addr::new(127, 0, 0, 1), port: 69, directory: env::current_dir().unwrap_or_else(|_| env::temp_dir()), + receive_directory: PathBuf::new(), + send_directory: PathBuf::new(), single_port: false, read_only: false, duplicate_packets: 0, @@ -72,11 +78,31 @@ impl Config { if !Path::new(&dir_str).exists() { return Err(format!("{dir_str} does not exist").into()); } - config.directory = PathBuf::from(dir_str); + config.directory = dir_str.into(); } else { return Err("Missing directory after flag".into()); } } + "-rd" | "--receive-directory" => { + if let Some(dir_str) = args.next() { + if !Path::new(&dir_str).exists() { + return Err(format!("{dir_str} does not exist").into()); + } + config.receive_directory = dir_str.into(); + } else { + return Err("Missing receive directory after flag".into()); + } + } + "-sd" | "--send-directory" => { + if let Some(dir_str) = args.next() { + if !Path::new(&dir_str).exists() { + return Err(format!("{dir_str} does not exist").into()); + } + config.send_directory = dir_str.into(); + } else { + return Err("Missing send directory after flag".into()); + } + } "-s" | "--single-port" => { config.single_port = true; } @@ -91,7 +117,9 @@ impl Config { println!( " -p, --port \t\tSet the listening port of the server (default: 69)" ); - println!(" -d, --directory \tSet the serving directory (default: Current Working Directory)"); + println!(" -d, --directory \tSet the serving directory (default: current working directory)"); + println!(" -rd, --receive-directory \tSet the directory to receive files to (default: the directory setting)"); + println!(" -sd, --send-directory \tSet the directory to send files from (default: the directory setting)"); println!(" -s, --single-port\t\tUse a single port for both sending and receiving (default: false)"); println!(" -r, --read-only\t\tRefuse all write requests, making the server read-only (default: false)"); println!(" --duplicate-packets \tDuplicate all packets sent from the server (default: 0)"); @@ -123,28 +151,37 @@ impl Config { } } + if config.receive_directory == PathBuf::new() { + config.receive_directory = config.directory.clone(); + } + if config.send_directory == PathBuf::new() { + config.send_directory = config.directory.clone(); + } + Ok(config) } } #[cfg(test)] mod tests { - use std::str::FromStr; - use super::*; #[test] fn parses_full_config() { let config = Config::new( - ["/", "-i", "0.0.0.0", "-p", "1234", "-d", "/", "-s", "-r"] - .iter() - .map(|s| s.to_string()), + [ + "/", "-i", "0.0.0.0", "-p", "1234", "-d", "/", "-rd", "/", "-sd", "/", "-s", "-r", + ] + .iter() + .map(|s| s.to_string()), ) .unwrap(); assert_eq!(config.ip_address, Ipv4Addr::new(0, 0, 0, 0)); assert_eq!(config.port, 1234); - assert_eq!(config.directory, PathBuf::from_str("/").unwrap()); + assert_eq!(config.directory, PathBuf::from("/")); + assert_eq!(config.receive_directory, PathBuf::from("/")); + assert_eq!(config.send_directory, PathBuf::from("/")); assert!(config.single_port); assert!(config.read_only); } @@ -160,7 +197,21 @@ mod tests { assert_eq!(config.ip_address, Ipv4Addr::new(0, 0, 0, 0)); assert_eq!(config.port, 69); - assert_eq!(config.directory, PathBuf::from_str("/").unwrap()); + assert_eq!(config.directory, PathBuf::from("/")); + } + + #[test] + fn sets_receive_directory_to_directory() { + let config = Config::new(["/", "-d", "/"].iter().map(|s| s.to_string())).unwrap(); + + assert_eq!(config.receive_directory, PathBuf::from("/")); + } + + #[test] + fn sets_send_directory_to_directory() { + let config = Config::new(["/", "-d", "/"].iter().map(|s| s.to_string())).unwrap(); + + assert_eq!(config.send_directory, PathBuf::from("/")); } #[test] @@ -188,6 +239,26 @@ mod tests { .is_err()); } + #[test] + fn returns_error_on_invalid_up_directory() { + assert!(Config::new( + ["/", "-ud", "/this/does/not/exist"] + .iter() + .map(|s| s.to_string()), + ) + .is_err()); + } + + #[test] + fn returns_error_on_invalid_down_directory() { + assert!(Config::new( + ["/", "-dd", "/this/does/not/exist"] + .iter() + .map(|s| s.to_string()), + ) + .is_err()); + } + #[test] fn returns_error_on_invalid_duplicate_packets() { assert!(Config::new( diff --git a/src/main.rs b/src/main.rs index 8bc460a..cf2e7b6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -15,12 +15,22 @@ fn main() { process::exit(1) }); - println!( - "Running TFTP Server on {}:{} in {}", - config.ip_address, - config.port, - config.directory.display() - ); + if config.receive_directory == config.send_directory { + println!( + "Running TFTP Server on {}:{} in {}", + config.ip_address, + config.port, + config.directory.display() + ); + } else { + println!( + "Running TFTP Server on {}:{}. Sending from {}, receiving to {}", + config.ip_address, + config.port, + config.send_directory.display(), + config.receive_directory.display(), + ); + } server.listen(); } diff --git a/src/server.rs b/src/server.rs index b607fe3..e241a17 100644 --- a/src/server.rs +++ b/src/server.rs @@ -29,7 +29,8 @@ const DEFAULT_WINDOW_SIZE: u16 = 1; /// ``` pub struct Server { socket: UdpSocket, - directory: PathBuf, + receive_directory: PathBuf, + send_directory: PathBuf, single_port: bool, read_only: bool, overwrite: bool, @@ -42,10 +43,10 @@ impl Server { /// Creates the TFTP Server with the supplied [`Config`]. pub fn new(config: &Config) -> Result> { let socket = UdpSocket::bind(SocketAddr::from((config.ip_address, config.port)))?; - let server = Server { socket, - directory: config.directory.clone(), + receive_directory: config.receive_directory.clone(), + send_directory: config.send_directory.clone(), single_port: config.single_port, read_only: config.read_only, overwrite: config.overwrite, @@ -132,8 +133,8 @@ impl Server { options: &mut [TransferOption], to: &SocketAddr, ) -> Result<(), Box> { - let file_path = &self.directory.join(filename); - match check_file_exists(file_path, &self.directory) { + let file_path = &self.send_directory.join(filename); + match check_file_exists(file_path, &self.send_directory) { ErrorCode::FileNotFound => Socket::send_to( &self.socket, &Packet::Error { @@ -195,7 +196,7 @@ impl Server { options: &mut [TransferOption], to: &SocketAddr, ) -> Result<(), Box> { - let file_path = &self.directory.join(file_name); + let file_path = &self.receive_directory.join(file_name); let initialize_write = &mut || -> Result<(), Box> { let worker_options = parse_options(options, RequestType::Write)?; let mut socket: Box; @@ -226,7 +227,7 @@ impl Server { worker.receive() }; - match check_file_exists(file_path, &self.directory) { + match check_file_exists(file_path, &self.receive_directory) { ErrorCode::FileExists => { if self.overwrite { initialize_write() diff --git a/tests/integration_test.rs b/tests/integration_test.rs index 9159d02..b16ee8c 100644 --- a/tests/integration_test.rs +++ b/tests/integration_test.rs @@ -1,6 +1,6 @@ #![cfg(feature = "integration")] -use std::fs::create_dir_all; +use std::fs::{create_dir_all, remove_dir_all}; use std::process::{Child, Command, ExitStatus}; const SERVER_DIR: &str = "target/integration/server"; @@ -41,8 +41,18 @@ fn test_send() { initialize(format!("{SERVER_DIR}/{file_name}").as_str()); let _server = CommandRunner::new("target/debug/tftpd", &["-p", port, "-d", SERVER_DIR]); - let mut client = - CommandRunner::new("time", &["atftp", "-g", "-r", file_name, "127.0.0.1", port]); + let mut client = CommandRunner::new( + "atftp", + &[ + "-g", + "-r", + file_name, + "-l", + format!("{CLIENT_DIR}/{file_name}").as_str(), + "127.0.0.1", + port, + ], + ); let status = client.wait(); assert!(status.success()); @@ -56,9 +66,56 @@ fn test_receive() { let _server = CommandRunner::new("target/debug/tftpd", &["-p", port, "-d", SERVER_DIR]); let mut client = CommandRunner::new( - "time", + "atftp", + &[ + "-p", + "-r", + file_name, + "-l", + format!("{CLIENT_DIR}/{file_name}").as_str(), + "127.0.0.1", + port, + ], + ); + + let status = client.wait(); + assert!(status.success()); +} + +#[test] +fn test_send_dir() { + let file_name = "send_dir"; + let port = "6971"; + initialize(format!("{SERVER_DIR}/{file_name}").as_str()); + + let _server = CommandRunner::new("target/debug/tftpd", &["-p", port, "-sd", SERVER_DIR]); + let mut client = CommandRunner::new( + "atftp", + &[ + "-g", + "-r", + file_name, + "-l", + format!("{CLIENT_DIR}/{file_name}").as_str(), + "127.0.0.1", + port, + ], + ); + + let status = client.wait(); + assert!(status.success()); +} + +#[test] +fn test_receive_dir() { + let file_name = "receive_dir"; + let port = "6972"; + initialize(format!("{CLIENT_DIR}/{file_name}").as_str()); + + let _server = CommandRunner::new("target/debug/tftpd", &["-p", port, "-rd", SERVER_DIR]); + let mut client = CommandRunner::new( + "atftp", &[ - "atftp", "-p", "-r", file_name, @@ -79,6 +136,8 @@ fn initialize(file_name: &str) { } fn create_folders() { + let _ = remove_dir_all(SERVER_DIR); + let _ = remove_dir_all(CLIENT_DIR); create_dir_all(SERVER_DIR).expect("error creating server directory"); create_dir_all(CLIENT_DIR).expect("error creating client directory"); }