ae/contego
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Better req handling & reading

Browse Source 
* Fixed old typos
* Fixed problem caused by empty data folder
* Included newlines to requests to separate individual messages
This commit is contained in:
einisto 2022-07-16 04:49:55 +03:00
parent 44ef3d8db1
commit 530e33867b
2 changed files with 82 additions and 33 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
src/client.rs
View File

@ -29,18 +29,27 @@ pub async fn connect(addr: String, fileroot: PathBuf) -> Result<(), Box<dyn std:
buf.clear(); buf.clear();
// ACK buffersize // ACK buffersize
writer.write_all(b"ACK").await.unwrap(); writer.write_all(b"ACK\n").await.unwrap();
writer.flush().await?; writer.flush().await?;
// Receive file amount // Receive file amount (or termination request if the server does not have any files available)
let _bytes_read = reader.read_buf(&mut buf).await?; let file_amount: usize;
let file_amount = String::from_utf8(buf.clone())?.parse::<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::<usize>()?;
println!("[+] Total of {} files available", file_amount); println!("[+] Total of {} files available", file_amount);
buf.clear(); buf.clear();
// ACK file amount // ACK file amount
writer.write_all(b"ACK").await?; writer.write_all(b"ACK\n").await?;
writer.flush().await?; writer.flush().await?;
}
// Receive file metadata // Receive file metadata
println!("[+] Receiving file metadata"); println!("[+] Receiving file metadata");
@ -64,7 +73,8 @@ pub async fn connect(addr: String, fileroot: PathBuf) -> Result<(), Box<dyn std:
println!("[+] Requesting files individually"); println!("[+] Requesting files individually");
for file in &metadata { for file in &metadata {
println!("[INFO] Current request: [{:?}]", file); println!("[INFO] Current request: [{:?}]", file);
writer.write_all(file.0.as_bytes()).await?; let msg = file.0.to_string() + "\n";
writer.write_all(msg.as_bytes()).await?;
writer.flush().await?; writer.flush().await?;
// Create file locally // Create file locally
@ -112,7 +122,7 @@ pub async fn connect(addr: String, fileroot: PathBuf) -> Result<(), Box<dyn std:
} }
// ACK file // ACK file
writer.write_all(b"ACK").await?; writer.write_all(b"ACK\n").await?;
writer.flush().await?; writer.flush().await?;
println!( println!(
"[+] Successfully wrote {} bytes to {:#?}\n", "[+] Successfully wrote {} bytes to {:#?}\n",
@ -121,7 +131,7 @@ pub async fn connect(addr: String, fileroot: PathBuf) -> Result<(), Box<dyn std:
} }
println!("[+] All files finished, requesting connection termination"); println!("[+] All files finished, requesting connection termination");
writer.write_all(b"FIN").await?; writer.write_all(b"FIN\n").await?;
writer.flush().await?; writer.flush().await?;
} }

81
src/server.rs
View File

@ -9,7 +9,7 @@ use std::{
use tokio::{ use tokio::{
self, self,
fs::File, fs::File,
io::{AsyncReadExt, AsyncWriteExt, BufReader, BufWriter}, io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader, BufWriter},
net::TcpListener, net::TcpListener,
time::timeout, time::timeout,
}; };
@ -36,13 +36,16 @@ pub async fn listen(
match timeout(Duration::from_secs(timeout_duration), listener.accept()).await { match timeout(Duration::from_secs(timeout_duration), listener.accept()).await {
Ok(n) => n?, Ok(n) => n?,
Err(_) => { Err(_) => {
println!("\nConnection timed out after {} seconds", timeout_duration); println!(
"\n[-] Connection timed out after {} seconds",
timeout_duration
);
break; break;
} }
}; };
println!("\n[+] New client: {}", addr); println!("\n[+] New client: {}", addr);
tokio::spawn(async move { let task = tokio::spawn(async move {
let (reader, writer) = socket.split(); let (reader, writer) = socket.split();
let mut reader = BufReader::new(reader); let mut reader = BufReader::new(reader);
let mut writer = BufWriter::new(writer); let mut writer = BufWriter::new(writer);
@ -54,23 +57,49 @@ pub async fn listen(
writer.flush().await?; writer.flush().await?;
// Read ACK // Read ACK
let _bytes_read = reader.read_buf(&mut vec_buf).await?; //let _bytes_read = reader.read_buf(&mut vec_buf).await?;
if String::from_utf8(vec_buf.clone())? != "ACK" { let _bytes_read = reader.read_until(b'\n', &mut vec_buf).await?;
panic!("ACK not received (buffersize)"); let msg = String::from_utf8(vec_buf.clone())?;
if msg.trim() != "ACK" {
panic!("ACK not received (buffersize), instead got {}", msg);
} else { } else {
vec_buf.clear(); 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 {
// Send file amount println!(
writer.write_all(file_amount.to_string().as_bytes()).await?; "[-] 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?; 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<dyn std::error::Error + Send + Sync> = "".into();
return Err::<(), Box<dyn std::error::Error + Send + Sync>>(e);
}
} else {
// Send file amount
let msg = file_amount.to_string() + "\n";
writer.write_all(msg.as_bytes()).await?;
writer.flush().await?;
}
// Read ACK // Read ACK
let _bytes_read = reader.read_buf(&mut vec_buf).await?; //let _bytes_read = reader.read_buf(&mut vec_buf).await?;
if String::from_utf8(vec_buf.clone())? != "ACK" { let _bytes_read = reader.read_until(b'\n', &mut vec_buf).await?;
panic!("ACK not received (amount)"); let msg = String::from_utf8(vec_buf.clone())?;
if msg.trim() != "ACK" {
panic!("ACK not received (amount), instead got {}", msg);
} else { } else {
vec_buf.clear(); vec_buf.clear();
} }
@ -86,22 +115,24 @@ pub async fn listen(
// Handle file request(s) // Handle file request(s)
println!("[+] Ready to serve files"); println!("[+] Ready to serve files");
loop { 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 { if bytes_read == 0 {
println!("File request never received"); println!("[-] File request never received");
vec_buf.clear();
break; break;
} else { } else {
let msg = String::from_utf8(vec_buf.clone())?; let msg = String::from_utf8(vec_buf.clone())?;
vec_buf.clear(); vec_buf.clear();
if msg == "FIN" { if msg.trim() == "FIN" {
println!("[+] FIN received, terminating individual connection..."); println!("[+] FIN received, terminating individual connection...");
break; break;
} }
let mut input_path = alt_fileroot.clone(); let mut input_path = alt_fileroot.clone();
input_path.push(msg); input_path.push(msg.trim());
println!("\n[+] File requested: {:#?}", input_path); println!("\n[+] File requested: {:#?}", input_path);
let mut file = File::open(input_path.clone()).await?; let mut file = File::open(input_path.clone()).await?;
@ -122,9 +153,11 @@ pub async fn listen(
} }
// Read ACK // Read ACK
let _bytes_read = reader.read_buf(&mut vec_buf).await?; //let _bytes_read = reader.read_buf(&mut vec_buf).await?;
if String::from_utf8(vec_buf.clone())? != "ACK" { let _bytes_read = reader.read_until(b'\n', &mut vec_buf).await?;
panic!("ACK not received (amount)"); let msg = String::from_utf8(vec_buf.clone())?;
if msg.trim() != "ACK" {
panic!("ACK not received (file), instead got {}", msg);
} else { } else {
println!("[+] File transfer successfully done"); println!("[+] File transfer successfully done");
vec_buf.clear(); vec_buf.clear();
@ -133,15 +166,21 @@ pub async fn listen(
Ok::<(), Box<dyn std::error::Error + Send + Sync>>(()) Ok::<(), Box<dyn std::error::Error + Send + Sync>>(())
}); });
match task.await? {
Ok(_) => continue,
Err(_) => break,
}
} }
Ok(()) Ok(())
} }
async fn get_metadata( async fn get_metadata(
fileroot: PathBuf,
) -> Result<(Vec<(String, u64)>, usize), Box<dyn std::error::Error + Send + Sync>> { ) -> Result<(Vec<(String, u64)>, usize), Box<dyn std::error::Error + Send + Sync>> {
let mut metadata = Vec::<(String, u64)>::new(); let mut metadata = Vec::<(String, u64)>::new();
let paths = read_dir("./data")?; let paths = read_dir(fileroot)?;
for filename in paths { for filename in paths {
let filepath = filename?.path().display().to_string(); let filepath = filename?.path().display().to_string();