From 530e33867b1a2917ba9e7ac0caa423c8ecb890b7 Mon Sep 17 00:00:00 2001 From: einisto <79069176+einisto@users.noreply.github.com> Date: Sat, 16 Jul 2022 04:49:55 +0300 Subject: [PATCH] Better req handling & reading * Fixed old typos * Fixed problem caused by empty data folder * Included newlines to requests to separate individual messages --- src/client.rs | 34 +++++++++++++-------- src/server.rs | 81 ++++++++++++++++++++++++++++++++++++++------------- 2 files changed, 82 insertions(+), 33 deletions(-) diff --git a/src/client.rs b/src/client.rs index 050c247..4d667ac 100644 --- a/src/client.rs +++ b/src/client.rs @@ -29,18 +29,27 @@ pub async fn connect(addr: String, fileroot: PathBuf) -> Result<(), Box()?; - println!("[+] Total of {} files available", file_amount); - buf.clear(); + // Receive file amount (or termination request if the server does not have any files available) + let file_amount: usize; + let _bytes_read = reader.read_until(b'\n', &mut buf).await?; + let msg = String::from_utf8(buf.clone())?; + if msg.trim() == "FIN" { + println!("[-] Server does not have any files available, closing connection"); + writer.write_all(b"FIN\n").await?; + writer.flush().await?; + break; + } else { + file_amount = msg.trim().parse::()?; + println!("[+] Total of {} files available", file_amount); + buf.clear(); - // ACK file amount - writer.write_all(b"ACK").await?; - writer.flush().await?; + // ACK file amount + writer.write_all(b"ACK\n").await?; + writer.flush().await?; + } // Receive file metadata println!("[+] Receiving file metadata"); @@ -64,7 +73,8 @@ pub async fn connect(addr: String, fileroot: PathBuf) -> Result<(), Box Result<(), Box Result<(), Box n?, Err(_) => { - println!("\nConnection timed out after {} seconds", timeout_duration); + println!( + "\n[-] Connection timed out after {} seconds", + timeout_duration + ); break; } }; println!("\n[+] New client: {}", addr); - tokio::spawn(async move { + let task = tokio::spawn(async move { let (reader, writer) = socket.split(); let mut reader = BufReader::new(reader); let mut writer = BufWriter::new(writer); @@ -54,23 +57,49 @@ pub async fn listen( writer.flush().await?; // Read ACK - let _bytes_read = reader.read_buf(&mut vec_buf).await?; - if String::from_utf8(vec_buf.clone())? != "ACK" { - panic!("ACK not received (buffersize)"); + //let _bytes_read = reader.read_buf(&mut vec_buf).await?; + let _bytes_read = reader.read_until(b'\n', &mut vec_buf).await?; + let msg = String::from_utf8(vec_buf.clone())?; + if msg.trim() != "ACK" { + panic!("ACK not received (buffersize), instead got {}", msg); } else { vec_buf.clear(); } - let (metadata_list, file_amount) = get_metadata().await?; + let (metadata_list, file_amount) = get_metadata(alt_fileroot.clone()).await?; + if file_amount == 0 { + println!( + "[-] No files available inside fileroot {:#?}, shutting down", + alt_fileroot + ); + // No files available, request connection termination + writer.write_all(b"FIN\n").await?; + writer.flush().await?; - // Send file amount - writer.write_all(file_amount.to_string().as_bytes()).await?; - writer.flush().await?; + // Read FIN + //let _bytes_read = reader.read_buf(&mut vec_buf).await?; + let _bytes_read = reader.read_until(b'\n', &mut vec_buf).await?; + let msg = String::from_utf8(vec_buf.clone())?; + if msg.trim() != "FIN" { + panic!("ACK not received (server-side-termination)"); + } else { + // Empty error as error's reason is already logged with println + let e: Box = "".into(); + return Err::<(), Box>(e); + } + } else { + // Send file amount + let msg = file_amount.to_string() + "\n"; + writer.write_all(msg.as_bytes()).await?; + writer.flush().await?; + } // Read ACK - let _bytes_read = reader.read_buf(&mut vec_buf).await?; - if String::from_utf8(vec_buf.clone())? != "ACK" { - panic!("ACK not received (amount)"); + //let _bytes_read = reader.read_buf(&mut vec_buf).await?; + let _bytes_read = reader.read_until(b'\n', &mut vec_buf).await?; + let msg = String::from_utf8(vec_buf.clone())?; + if msg.trim() != "ACK" { + panic!("ACK not received (amount), instead got {}", msg); } else { vec_buf.clear(); } @@ -86,22 +115,24 @@ pub async fn listen( // Handle file request(s) println!("[+] Ready to serve files"); loop { - let bytes_read = reader.read_buf(&mut vec_buf).await?; + //let bytes_read = reader.read_buf(&mut vec_buf).await?; + let bytes_read = reader.read_until(b'\n', &mut vec_buf).await?; if bytes_read == 0 { - println!("File request never received"); + println!("[-] File request never received"); + vec_buf.clear(); break; } else { let msg = String::from_utf8(vec_buf.clone())?; vec_buf.clear(); - if msg == "FIN" { + if msg.trim() == "FIN" { println!("[+] FIN received, terminating individual connection..."); break; } let mut input_path = alt_fileroot.clone(); - input_path.push(msg); + input_path.push(msg.trim()); println!("\n[+] File requested: {:#?}", input_path); let mut file = File::open(input_path.clone()).await?; @@ -122,9 +153,11 @@ pub async fn listen( } // Read ACK - let _bytes_read = reader.read_buf(&mut vec_buf).await?; - if String::from_utf8(vec_buf.clone())? != "ACK" { - panic!("ACK not received (amount)"); + //let _bytes_read = reader.read_buf(&mut vec_buf).await?; + let _bytes_read = reader.read_until(b'\n', &mut vec_buf).await?; + let msg = String::from_utf8(vec_buf.clone())?; + if msg.trim() != "ACK" { + panic!("ACK not received (file), instead got {}", msg); } else { println!("[+] File transfer successfully done"); vec_buf.clear(); @@ -133,15 +166,21 @@ pub async fn listen( Ok::<(), Box>(()) }); + + match task.await? { + Ok(_) => continue, + Err(_) => break, + } } Ok(()) } async fn get_metadata( + fileroot: PathBuf, ) -> Result<(Vec<(String, u64)>, usize), Box> { let mut metadata = Vec::<(String, u64)>::new(); - let paths = read_dir("./data")?; + let paths = read_dir(fileroot)?; for filename in paths { let filepath = filename?.path().display().to_string();