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();
// ACK buffersize
writer.write_all(b"ACK").await.unwrap();
writer.write_all(b"ACK\n").await.unwrap();
writer.flush().await?;
// Receive file amount
let _bytes_read = reader.read_buf(&mut buf).await?;
let file_amount = String::from_utf8(buf.clone())?.parse::<usize>()?;
// 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::<usize>()?;
println!("[+] Total of {} files available", file_amount);
buf.clear();
// ACK file amount
writer.write_all(b"ACK").await?;
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<dyn std:
println!("[+] Requesting files individually");
for file in &metadata {
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?;
// Create file locally
@ -112,7 +122,7 @@ pub async fn connect(addr: String, fileroot: PathBuf) -> Result<(), Box<dyn std:
}
// ACK file
writer.write_all(b"ACK").await?;
writer.write_all(b"ACK\n").await?;
writer.flush().await?;
println!(
"[+] 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");
writer.write_all(b"FIN").await?;
writer.write_all(b"FIN\n").await?;
writer.flush().await?;
}

81
src/server.rs
View File

@ -9,7 +9,7 @@ use std::{
use tokio::{
self,
fs::File,
io::{AsyncReadExt, AsyncWriteExt, BufReader, BufWriter},
io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader, BufWriter},
net::TcpListener,
time::timeout,
};
@ -36,13 +36,16 @@ pub async fn listen(
match timeout(Duration::from_secs(timeout_duration), listener.accept()).await {
Ok(n) => 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?;
// Send file amount
writer.write_all(file_amount.to_string().as_bytes()).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?;
// 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
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<dyn std::error::Error + Send + Sync>>(())
});
match task.await? {
Ok(_) => continue,
Err(_) => break,
}
}
Ok(())
}
async fn get_metadata(
fileroot: PathBuf,
) -> Result<(Vec<(String, u64)>, usize), Box<dyn std::error::Error + Send + Sync>> {
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();