Add upload directory and download directory (#14)

Co-authored-by: y1shu <yu.shu@nokia-sbell.com>
Co-authored-by: Altuğ Bakan <mail@alt.ug>

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 <PORT>\t\tSet the listening port of the server (default: 69)"
                     );
-                    println!("  -d, --directory <DIRECTORY>\tSet the serving directory (default: Current Working Directory)");
+                    println!("  -d, --directory <DIRECTORY>\tSet the serving directory (default: current working directory)");
+                    println!("  -rd, --receive-directory <DIRECTORY>\tSet the directory to receive files to (default: the directory setting)");
+                    println!("  -sd, --send-directory <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 <NUM>\tDuplicate all packets sent from the server (default: 0)");
@@ -123,20 +151,27 @@ 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"]
+            [
+                "/", "-i", "0.0.0.0", "-p", "1234", "-d", "/", "-rd", "/", "-sd", "/", "-s", "-r",
+            ]
             .iter()
             .map(|s| s.to_string()),
         )
@@ -144,7 +179,9 @@ mod tests {
 
         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(

src/main.rs

@@ -15,12 +15,22 @@ fn main() {
         process::exit(1)
     });
 
+    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();
 }

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<Server, Box<dyn Error>> {
         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<dyn Error>> {
-        let file_path = &self.directory.join(filename);
+        let file_path = &self.send_directory.join(filename);
-        match check_file_exists(file_path, &self.directory) {
+        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<dyn Error>> {
-        let file_path = &self.directory.join(file_name);
+        let file_path = &self.receive_directory.join(file_name);
         let initialize_write = &mut || -> Result<(), Box<dyn Error>> {
             let worker_options = parse_options(options, RequestType::Write)?;
             let mut socket: Box<dyn Socket>;
@@ -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()

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 =
+    let mut client = CommandRunner::new(
-        CommandRunner::new("time", &["atftp", "-g", "-r", file_name, "127.0.0.1", port]);
+        "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",
+        &[
             "-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");
 }