mykola/rs-tftpd
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source
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
Show stats Download patch file Download diff file Expand all files Collapse all files

4
src/client.rs
View file

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

9
src/client_config.rs
View file

@ -39,6 +39,8 @@ pub struct ClientConfig {
pub receive_directory: PathBuf,
/// File to Upload or Download.
pub filename: PathBuf,
/// Should clean (delete) files after receiving errors. (default: true)
pub clean_on_error: bool
}
impl Default for ClientConfig {
@ -52,6 +54,7 @@ impl Default for ClientConfig {
mode: Mode::Download,
receive_directory: Default::default(),
filename: Default::default(),
clean_on_error: true
}
}
}
@ -116,6 +119,9 @@ impl ClientConfig {
"-d" | "--download" => {
config.mode = Mode::Download;
}
"--dont-clean" => {
config.clean_on_error = false;
}
"-h" | "--help" => {
println!("TFTP Client\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!(" -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!(" --dont-clean\t\t\t\tWill prevent client from deleting files after receiving errors.");
println!(" -h, --help\t\t\t\tPrint help information");
process::exit(0);
}
@ -169,6 +176,7 @@ mod tests {
"2",
"-t",
"4",
"--dont-clean"
]
.iter()
.map(|s| s.to_string()),
@ -183,6 +191,7 @@ mod tests {
assert_eq!(config.blocksize, 1024);
assert_eq!(config.mode, Mode::Upload);
assert_eq!(config.timeout, Duration::from_secs(4));
assert_eq!(config.clean_on_error, false);
}
#[test]

10
src/config.rs
View file

@ -37,6 +37,8 @@ pub struct Config {
pub duplicate_packets: u8,
/// Overwrite existing files. (default: false)
pub overwrite: bool,
/// Should clean (delete) files after receiving errors. (default: true)
pub clean_on_error: bool
}
impl Default for Config {
@ -51,6 +53,7 @@ impl Default for Config {
read_only: Default::default(),
duplicate_packets: 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!(" --duplicate-packets <NUM>\t\tDuplicate all packets sent from the server (default: 0)");
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");
process::exit(0);
}
@ -153,6 +157,9 @@ impl Config {
"--overwrite" => {
config.overwrite = true;
}
"--dont-clean" => {
config.clean_on_error = false;
}
invalid => return Err(format!("Invalid flag: {invalid}").into()),
}
@ -179,7 +186,7 @@ mod tests {
fn parses_full_config() {
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()
.map(|s| s.to_string()),
@ -191,6 +198,7 @@ mod tests {
assert_eq!(config.directory, PathBuf::from("/"));
assert_eq!(config.receive_directory, PathBuf::from("/"));
assert_eq!(config.send_directory, PathBuf::from("/"));
assert_eq!(config.clean_on_error, false);
assert!(config.single_port);
assert!(config.read_only);
}

4
src/server.rs
View file

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

7
src/worker.rs
View file

@ -32,6 +32,7 @@ const DEFAULT_DUPLICATE_DELAY: Duration = Duration::from_millis(1);
/// let worker = Worker::new(
/// Box::new(socket),
/// PathBuf::from_str("Cargo.toml").unwrap(),
/// true,
/// 512,
/// Duration::from_secs(1),
/// 1,
@ -43,6 +44,7 @@ const DEFAULT_DUPLICATE_DELAY: Duration = Duration::from_millis(1);
pub struct Worker<T: Socket + ?Sized> {
socket: Box<T>,
file_name: PathBuf,
clean_on_error: bool,
blk_size: usize,
timeout: Duration,
windowsize: u16,
@ -54,6 +56,7 @@ impl<T: Socket + ?Sized> Worker<T> {
pub fn new(
socket: Box<T>,
file_name: PathBuf,
clean_on_error: bool,
blk_size: usize,
timeout: Duration,
windowsize: u16,
@ -62,6 +65,7 @@ impl<T: Socket + ?Sized> Worker<T> {
Worker {
socket,
file_name,
clean_on_error,
blk_size,
timeout,
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
/// the supplied socket, asynchronously.
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 remote_addr = self.socket.remote_addr().unwrap();
@ -122,11 +127,13 @@ impl<T: Socket + ?Sized> Worker<T> {
}
Err(err) => {
eprintln!("{err}");
if clean_on_error {
if fs::remove_file(&file_name).is_err() {
eprintln!("Error while cleaning {}", &file_name.to_str().unwrap());
}
}
}
}
});
Ok(handle)