feat: make cli option to prevent file removal after receiving errors. integrate option into Config & ClientConfig and populate to Worker applying required behavior

This commit is contained in:
mykola2312 2024-07-17 05:33:30 +03:00
parent b9264c8cb0
commit de0e790962
5 changed files with 35 additions and 3 deletions

View file

@ -30,6 +30,7 @@ pub struct Client {
mode: Mode, mode: Mode,
filename: PathBuf, filename: PathBuf,
save_path: PathBuf, save_path: PathBuf,
clean_on_error: bool
} }
/// Enum used to set the client either in Download Mode or Upload Mode /// Enum used to set the client either in Download Mode or Upload Mode
@ -52,6 +53,7 @@ impl Client {
mode: config.mode, mode: config.mode,
filename: config.filename.clone(), filename: config.filename.clone(),
save_path: config.receive_directory.clone(), save_path: config.receive_directory.clone(),
clean_on_error: config.clean_on_error
}) })
} }
@ -240,6 +242,7 @@ impl Client {
Worker::new( Worker::new(
socket, socket,
file, file,
self.clean_on_error,
self.blocksize, self.blocksize,
DEFAULT_TIMEOUT, DEFAULT_TIMEOUT,
self.windowsize, self.windowsize,
@ -249,6 +252,7 @@ impl Client {
Worker::new( Worker::new(
socket, socket,
PathBuf::from(self.filename.clone()), PathBuf::from(self.filename.clone()),
self.clean_on_error,
self.blocksize, self.blocksize,
DEFAULT_TIMEOUT, DEFAULT_TIMEOUT,
self.windowsize, self.windowsize,

View file

@ -39,6 +39,8 @@ pub struct ClientConfig {
pub receive_directory: PathBuf, pub receive_directory: PathBuf,
/// File to Upload or Download. /// File to Upload or Download.
pub filename: PathBuf, pub filename: PathBuf,
/// Should clean (delete) files after receiving errors. (default: true)
pub clean_on_error: bool
} }
impl Default for ClientConfig { impl Default for ClientConfig {
@ -52,6 +54,7 @@ impl Default for ClientConfig {
mode: Mode::Download, mode: Mode::Download,
receive_directory: Default::default(), receive_directory: Default::default(),
filename: Default::default(), filename: Default::default(),
clean_on_error: true
} }
} }
} }
@ -116,6 +119,9 @@ impl ClientConfig {
"-d" | "--download" => { "-d" | "--download" => {
config.mode = Mode::Download; config.mode = Mode::Download;
} }
"--dont-clean" => {
config.clean_on_error = false;
}
"-h" | "--help" => { "-h" | "--help" => {
println!("TFTP Client\n"); println!("TFTP Client\n");
println!("Usage: tftpd client <File> [OPTIONS]\n"); println!("Usage: tftpd client <File> [OPTIONS]\n");
@ -130,6 +136,7 @@ impl ClientConfig {
println!(" -u, --upload\t\t\t\tSets the client to upload mode, Ignores all previous download flags"); println!(" -u, --upload\t\t\t\tSets the client to upload mode, Ignores all previous download flags");
println!(" -d, --download\t\t\tSet the client to download mode, Invalidates all previous upload flags"); println!(" -d, --download\t\t\tSet the client to download mode, Invalidates all previous upload flags");
println!(" -rd, --receive-directory <DIRECTORY>\tSet the directory to receive files when in Download mode (default: current working directory)"); println!(" -rd, --receive-directory <DIRECTORY>\tSet the directory to receive files when in Download mode (default: current working directory)");
println!(" --dont-clean\t\t\t\tWill prevent client from deleting files after receiving errors.");
println!(" -h, --help\t\t\t\tPrint help information"); println!(" -h, --help\t\t\t\tPrint help information");
process::exit(0); process::exit(0);
} }
@ -169,6 +176,7 @@ mod tests {
"2", "2",
"-t", "-t",
"4", "4",
"--dont-clean"
] ]
.iter() .iter()
.map(|s| s.to_string()), .map(|s| s.to_string()),
@ -183,6 +191,7 @@ mod tests {
assert_eq!(config.blocksize, 1024); assert_eq!(config.blocksize, 1024);
assert_eq!(config.mode, Mode::Upload); assert_eq!(config.mode, Mode::Upload);
assert_eq!(config.timeout, Duration::from_secs(4)); assert_eq!(config.timeout, Duration::from_secs(4));
assert_eq!(config.clean_on_error, false);
} }
#[test] #[test]

View file

@ -37,6 +37,8 @@ pub struct Config {
pub duplicate_packets: u8, pub duplicate_packets: u8,
/// Overwrite existing files. (default: false) /// Overwrite existing files. (default: false)
pub overwrite: bool, pub overwrite: bool,
/// Should clean (delete) files after receiving errors. (default: true)
pub clean_on_error: bool
} }
impl Default for Config { impl Default for Config {
@ -51,6 +53,7 @@ impl Default for Config {
read_only: Default::default(), read_only: Default::default(),
duplicate_packets: Default::default(), duplicate_packets: Default::default(),
overwrite: Default::default(), overwrite: Default::default(),
clean_on_error: true
} }
} }
} }
@ -131,6 +134,7 @@ impl Config {
println!(" -r, --read-only\t\t\tRefuse all write requests, making the server read-only (default: false)"); println!(" -r, --read-only\t\t\tRefuse all write requests, making the server read-only (default: false)");
println!(" --duplicate-packets <NUM>\t\tDuplicate all packets sent from the server (default: 0)"); println!(" --duplicate-packets <NUM>\t\tDuplicate all packets sent from the server (default: 0)");
println!(" --overwrite\t\t\t\tOverwrite existing files (default: false)"); println!(" --overwrite\t\t\t\tOverwrite existing files (default: false)");
println!(" --dont-clean\t\t\t\tWill prevent daemon from deleting files after receiving errors.");
println!(" -h, --help\t\t\t\tPrint help information"); println!(" -h, --help\t\t\t\tPrint help information");
process::exit(0); process::exit(0);
} }
@ -153,6 +157,9 @@ impl Config {
"--overwrite" => { "--overwrite" => {
config.overwrite = true; config.overwrite = true;
} }
"--dont-clean" => {
config.clean_on_error = false;
}
invalid => return Err(format!("Invalid flag: {invalid}").into()), invalid => return Err(format!("Invalid flag: {invalid}").into()),
} }
@ -179,7 +186,7 @@ mod tests {
fn parses_full_config() { fn parses_full_config() {
let config = Config::new( let config = Config::new(
[ [
"/", "-i", "0.0.0.0", "-p", "1234", "-d", "/", "-rd", "/", "-sd", "/", "-s", "-r", "/", "-i", "0.0.0.0", "-p", "1234", "-d", "/", "-rd", "/", "-sd", "/", "-s", "-r", "--dont-clean"
] ]
.iter() .iter()
.map(|s| s.to_string()), .map(|s| s.to_string()),
@ -191,6 +198,7 @@ mod tests {
assert_eq!(config.directory, PathBuf::from("/")); assert_eq!(config.directory, PathBuf::from("/"));
assert_eq!(config.receive_directory, PathBuf::from("/")); assert_eq!(config.receive_directory, PathBuf::from("/"));
assert_eq!(config.send_directory, PathBuf::from("/")); assert_eq!(config.send_directory, PathBuf::from("/"));
assert_eq!(config.clean_on_error, false);
assert!(config.single_port); assert!(config.single_port);
assert!(config.read_only); assert!(config.read_only);
} }

View file

@ -34,6 +34,7 @@ pub struct Server {
single_port: bool, single_port: bool,
read_only: bool, read_only: bool,
overwrite: bool, overwrite: bool,
clean_on_error: bool,
duplicate_packets: u8, duplicate_packets: u8,
largest_block_size: usize, largest_block_size: usize,
clients: HashMap<SocketAddr, Sender<Packet>>, clients: HashMap<SocketAddr, Sender<Packet>>,
@ -50,6 +51,7 @@ impl Server {
single_port: config.single_port, single_port: config.single_port,
read_only: config.read_only, read_only: config.read_only,
overwrite: config.overwrite, overwrite: config.overwrite,
clean_on_error: config.clean_on_error,
duplicate_packets: config.duplicate_packets, duplicate_packets: config.duplicate_packets,
largest_block_size: DEFAULT_BLOCK_SIZE, largest_block_size: DEFAULT_BLOCK_SIZE,
clients: HashMap::new(), clients: HashMap::new(),
@ -179,6 +181,7 @@ impl Server {
let worker = Worker::new( let worker = Worker::new(
socket, socket,
file_path.clone(), file_path.clone(),
self.clean_on_error,
worker_options.block_size, worker_options.block_size,
worker_options.timeout, worker_options.timeout,
worker_options.window_size, worker_options.window_size,
@ -220,6 +223,7 @@ impl Server {
let worker = Worker::new( let worker = Worker::new(
socket, socket,
file_path.clone(), file_path.clone(),
self.clean_on_error,
worker_options.block_size, worker_options.block_size,
worker_options.timeout, worker_options.timeout,
worker_options.window_size, worker_options.window_size,

View file

@ -32,6 +32,7 @@ const DEFAULT_DUPLICATE_DELAY: Duration = Duration::from_millis(1);
/// let worker = Worker::new( /// let worker = Worker::new(
/// Box::new(socket), /// Box::new(socket),
/// PathBuf::from_str("Cargo.toml").unwrap(), /// PathBuf::from_str("Cargo.toml").unwrap(),
/// true,
/// 512, /// 512,
/// Duration::from_secs(1), /// Duration::from_secs(1),
/// 1, /// 1,
@ -43,6 +44,7 @@ const DEFAULT_DUPLICATE_DELAY: Duration = Duration::from_millis(1);
pub struct Worker<T: Socket + ?Sized> { pub struct Worker<T: Socket + ?Sized> {
socket: Box<T>, socket: Box<T>,
file_name: PathBuf, file_name: PathBuf,
clean_on_error: bool,
blk_size: usize, blk_size: usize,
timeout: Duration, timeout: Duration,
windowsize: u16, windowsize: u16,
@ -54,6 +56,7 @@ impl<T: Socket + ?Sized> Worker<T> {
pub fn new( pub fn new(
socket: Box<T>, socket: Box<T>,
file_name: PathBuf, file_name: PathBuf,
clean_on_error: bool,
blk_size: usize, blk_size: usize,
timeout: Duration, timeout: Duration,
windowsize: u16, windowsize: u16,
@ -62,6 +65,7 @@ impl<T: Socket + ?Sized> Worker<T> {
Worker { Worker {
socket, socket,
file_name, file_name,
clean_on_error,
blk_size, blk_size,
timeout, timeout,
windowsize, windowsize,
@ -102,6 +106,7 @@ impl<T: Socket + ?Sized> Worker<T> {
/// Receives a file from the remote [`SocketAddr`] that has sent a write request using /// Receives a file from the remote [`SocketAddr`] that has sent a write request using
/// the supplied socket, asynchronously. /// the supplied socket, asynchronously.
pub fn receive(self) -> Result<JoinHandle<()>, Box<dyn Error>> { pub fn receive(self) -> Result<JoinHandle<()>, Box<dyn Error>> {
let clean_on_error = self.clean_on_error;
let file_name = self.file_name.clone(); let file_name = self.file_name.clone();
let remote_addr = self.socket.remote_addr().unwrap(); let remote_addr = self.socket.remote_addr().unwrap();
@ -122,8 +127,10 @@ impl<T: Socket + ?Sized> Worker<T> {
} }
Err(err) => { Err(err) => {
eprintln!("{err}"); eprintln!("{err}");
if fs::remove_file(&file_name).is_err() { if clean_on_error {
eprintln!("Error while cleaning {}", &file_name.to_str().unwrap()); if fs::remove_file(&file_name).is_err() {
eprintln!("Error while cleaning {}", &file_name.to_str().unwrap());
}
} }
} }
} }