compressed & more concise structure (untested)
This commit is contained in:
parent
572c1e7468
commit
15579b016e
117
src/client.rs
Normal file
117
src/client.rs
Normal file
@ -0,0 +1,117 @@
|
||||
use std::{error::Error, net::SocketAddr, path::PathBuf};
|
||||
|
||||
use tokio::{io::AsyncWriteExt, net::TcpStream};
|
||||
|
||||
use crate::{
|
||||
crypto::{self, Crypto},
|
||||
sockets::SocketHandler,
|
||||
util::{new_file, FileInfo},
|
||||
};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Client {
|
||||
addr: SocketAddr,
|
||||
key: String,
|
||||
output: PathBuf,
|
||||
}
|
||||
|
||||
impl Client {
|
||||
pub fn new(addr: SocketAddr, key: String, output: PathBuf) -> Self {
|
||||
Self { addr, key, output }
|
||||
}
|
||||
|
||||
pub async fn connection(&self) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
let mut socket = TcpStream::connect(self.addr).await?;
|
||||
let mut handler = SocketHandler::new(&mut socket);
|
||||
let crypto = Crypto::new(&mut handler, true).await?;
|
||||
handler.set_crypto(crypto);
|
||||
|
||||
if !self.authorize(&mut handler).await? {
|
||||
// log: invalid access key '<self.key>'
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let metadata = self.metadata(&mut handler).await?;
|
||||
self.requests(&mut handler, metadata).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn authorize(
|
||||
&self,
|
||||
handler: &mut SocketHandler<'_>,
|
||||
) -> Result<bool, Box<dyn Error + Send + Sync>> {
|
||||
let msg = self.key.as_bytes().to_vec();
|
||||
handler.send(&msg).await?;
|
||||
|
||||
let buf = handler.recv().await?;
|
||||
let msg = String::from_utf8(buf)?;
|
||||
let msg = msg.trim();
|
||||
|
||||
if msg == "DISCONNECT" {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn metadata(
|
||||
&self,
|
||||
handler: &mut SocketHandler<'_>,
|
||||
) -> Result<Vec<FileInfo>, Box<dyn Error + Send + Sync>> {
|
||||
let buf = handler.recv().await?;
|
||||
let amt = String::from_utf8(buf.clone())?.parse::<usize>()?;
|
||||
handler.send(&buf).await?; // confirmation
|
||||
|
||||
let mut metadata = Vec::new();
|
||||
|
||||
while metadata.len() < amt {
|
||||
let buf = handler.recv().await?;
|
||||
let data = String::from_utf8(buf)?;
|
||||
|
||||
let split = data.split(':').collect::<Vec<&str>>();
|
||||
let name = split[0].trim().to_string();
|
||||
let size = split[1].trim().parse::<u64>()?;
|
||||
let hash = split[2].trim().to_string();
|
||||
|
||||
let info = FileInfo::new(name, size, hash);
|
||||
|
||||
metadata.push(info);
|
||||
}
|
||||
|
||||
Ok(metadata)
|
||||
}
|
||||
|
||||
async fn requests(
|
||||
&self,
|
||||
handler: &mut SocketHandler<'_>,
|
||||
metadata: Vec<FileInfo>,
|
||||
) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
for file in metadata {
|
||||
let (mut handle, path) = new_file(self.output.clone(), &file.name).await?;
|
||||
let msg = file.hash.as_bytes().to_vec();
|
||||
handler.send(&msg).await?;
|
||||
|
||||
// log: downloading file to <path>
|
||||
|
||||
let mut remaining = file.size;
|
||||
|
||||
while remaining != 0 {
|
||||
let buf = handler.recv().await?;
|
||||
handle.write_all(&buf).await?;
|
||||
handle.flush().await?;
|
||||
remaining -= buf.len() as u64;
|
||||
}
|
||||
|
||||
let check_hash = crypto::try_hash(&path)?;
|
||||
let msg = check_hash.as_bytes().to_vec();
|
||||
handler.send(&msg).await?;
|
||||
|
||||
if check_hash != file.hash {
|
||||
return Err("Unsuccessful file transfer, hashes don't match".into());
|
||||
} // else: log that the transfer was successful
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
231
src/connector.rs
231
src/connector.rs
@ -1,231 +0,0 @@
|
||||
use std::{collections::HashMap, error::Error, net::SocketAddr, path::PathBuf};
|
||||
|
||||
use tokio::{
|
||||
fs::File,
|
||||
io::{AsyncWriteExt, BufWriter},
|
||||
net::TcpStream,
|
||||
sync::mpsc,
|
||||
};
|
||||
|
||||
use crate::{
|
||||
common::{Connection, Message},
|
||||
comms, crypto,
|
||||
};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct Request {
|
||||
pub name: String,
|
||||
pub size: u64,
|
||||
pub hash: String,
|
||||
}
|
||||
|
||||
impl Request {
|
||||
pub fn new(name: String, metadata: &HashMap<String, (u64, String)>) -> Option<Self> {
|
||||
let (size, hash) = metadata.get(&name)?.clone();
|
||||
Some(Self { name, size, hash })
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Connector {
|
||||
target_addr: SocketAddr,
|
||||
access_key: String,
|
||||
output_path: PathBuf,
|
||||
}
|
||||
|
||||
impl Connector {
|
||||
pub fn new(target_addr: SocketAddr, access_key: String, output_path: PathBuf) -> Self {
|
||||
Self {
|
||||
target_addr,
|
||||
access_key,
|
||||
output_path,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn connect(
|
||||
self,
|
||||
tx: mpsc::Sender<Message>,
|
||||
mut rx: mpsc::Receiver<Message>,
|
||||
) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
let mut socket = TcpStream::connect(self.target_addr).await?;
|
||||
let mut connection = Connection::new(&mut socket).await?;
|
||||
|
||||
self.authorize(&mut connection).await?;
|
||||
let metadata = self.metadata(&mut connection).await?;
|
||||
tx.send(Message::Metadata(metadata.clone())).await?;
|
||||
self.request_handler(&mut connection, &mut rx, &metadata)
|
||||
.await?;
|
||||
|
||||
let msg = b"FIN".to_vec();
|
||||
comms::send(
|
||||
&mut connection.writer,
|
||||
Some(&mut connection.cipher),
|
||||
Some(&mut connection.rng),
|
||||
&msg,
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn authorize(
|
||||
&self,
|
||||
conn: &mut Connection<'_>,
|
||||
) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
let msg = self.access_key.to_string().as_bytes().to_vec();
|
||||
comms::send(
|
||||
&mut conn.writer,
|
||||
Some(&mut conn.cipher),
|
||||
Some(&mut conn.rng),
|
||||
&msg,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let buf = comms::recv(&mut conn.reader, Some(&mut conn.cipher)).await?;
|
||||
let msg = String::from_utf8(buf)?;
|
||||
|
||||
if msg == "FIN" {
|
||||
return Err("Incorrect access key".into());
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn metadata(
|
||||
&self,
|
||||
conn: &mut Connection<'_>,
|
||||
) -> Result<HashMap<String, (u64, String)>, Box<dyn Error + Send + Sync>> {
|
||||
let buf = comms::recv(&mut conn.reader, Some(&mut conn.cipher)).await?;
|
||||
let amt: usize = String::from_utf8(buf)?.parse()?;
|
||||
|
||||
let msg = b"AMT".to_vec();
|
||||
comms::send(
|
||||
&mut conn.writer,
|
||||
Some(&mut conn.cipher),
|
||||
Some(&mut conn.rng),
|
||||
&msg,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let mut metadata = HashMap::new();
|
||||
|
||||
while metadata.len() < amt {
|
||||
let buf = comms::recv(&mut conn.reader, Some(&mut conn.cipher)).await?;
|
||||
let msg = String::from_utf8(buf)?;
|
||||
|
||||
let split: Vec<&str> = msg.split(':').collect();
|
||||
let name = split[0].trim().to_string();
|
||||
let size: u64 = split[1].trim().parse()?;
|
||||
let hash = split[2].trim().to_string();
|
||||
|
||||
metadata.insert(name, (size, hash));
|
||||
}
|
||||
|
||||
Ok(metadata)
|
||||
}
|
||||
|
||||
async fn new_handle(
|
||||
&self,
|
||||
filename: &str,
|
||||
) -> Result<(BufWriter<File>, PathBuf), Box<dyn Error + Send + Sync>> {
|
||||
let mut path = self.output_path.clone();
|
||||
path.push(filename);
|
||||
let filehandle = File::create(&path).await?;
|
||||
|
||||
Ok((BufWriter::new(filehandle), path))
|
||||
}
|
||||
|
||||
async fn request(
|
||||
&self,
|
||||
conn: &mut Connection<'_>,
|
||||
req: Request,
|
||||
) -> Result<bool, Box<dyn Error + Send + Sync>> {
|
||||
let (mut handle, path) = self.new_handle(&req.name).await?;
|
||||
let msg = req.hash.as_bytes().to_vec();
|
||||
comms::send(
|
||||
&mut conn.writer,
|
||||
Some(&mut conn.cipher),
|
||||
Some(&mut conn.rng),
|
||||
&msg,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let mut remaining = req.size;
|
||||
|
||||
while remaining != 0 {
|
||||
let buf = comms::recv(&mut conn.reader, Some(&mut conn.cipher)).await?;
|
||||
handle.write_all(&buf).await?;
|
||||
handle.flush().await?;
|
||||
remaining -= buf.len() as u64;
|
||||
}
|
||||
|
||||
let new_hash = crypto::try_hash(&path)?;
|
||||
|
||||
let msg: Vec<u8> = if new_hash == req.hash {
|
||||
b"OK".to_vec()
|
||||
} else {
|
||||
b"ERROR".to_vec()
|
||||
};
|
||||
|
||||
comms::send(
|
||||
&mut conn.writer,
|
||||
Some(&mut conn.cipher),
|
||||
Some(&mut conn.rng),
|
||||
&msg,
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn request_handler(
|
||||
&self,
|
||||
conn: &mut Connection<'_>,
|
||||
rx: &mut mpsc::Receiver<Message>,
|
||||
metadata: &HashMap<String, (u64, String)>,
|
||||
) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
loop {
|
||||
let msg = rx.recv().await;
|
||||
if msg.is_none() {
|
||||
continue;
|
||||
}
|
||||
|
||||
match self.msg_handler(msg.unwrap(), conn, metadata).await {
|
||||
Ok(true) => continue,
|
||||
Ok(false) => break,
|
||||
Err(e) => return Err(e),
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn msg_handler(
|
||||
&self,
|
||||
msg: Message,
|
||||
conn: &mut Connection<'_>,
|
||||
metadata: &HashMap<String, (u64, String)>,
|
||||
) -> Result<bool, Box<dyn Error + Send + Sync>> {
|
||||
match msg {
|
||||
Message::ClientReq(name) => {
|
||||
let req = Request::new(name, metadata).unwrap();
|
||||
self.request(conn, req).await?;
|
||||
|
||||
Ok(true)
|
||||
}
|
||||
Message::Shutdown => {
|
||||
let msg = b"DISCONNECT".to_vec();
|
||||
comms::send(
|
||||
&mut conn.writer,
|
||||
Some(&mut conn.cipher),
|
||||
Some(&mut conn.rng),
|
||||
&msg,
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(false)
|
||||
}
|
||||
_ => Ok(true),
|
||||
}
|
||||
}
|
||||
}
|
@ -41,11 +41,11 @@ impl Crypto {
|
||||
let msg = own_pbk.as_bytes().to_vec();
|
||||
|
||||
if go_first {
|
||||
handler.send(&msg).await?;
|
||||
buf = handler.recv().await?;
|
||||
handler.send_raw(&msg).await?;
|
||||
buf = handler.recv_raw().await?;
|
||||
} else {
|
||||
buf = handler.recv().await?;
|
||||
handler.send(&msg).await?;
|
||||
buf = handler.recv_raw().await?;
|
||||
handler.send_raw(&msg).await?;
|
||||
}
|
||||
|
||||
let slice: [u8; DH_PBK_SIZE] = buf[..DH_PBK_SIZE].try_into()?;
|
||||
|
@ -1,7 +1,7 @@
|
||||
pub mod cli;
|
||||
pub mod connector;
|
||||
pub mod client;
|
||||
pub mod crypto;
|
||||
pub mod listener;
|
||||
pub mod parser;
|
||||
pub mod server;
|
||||
pub mod sockets;
|
||||
pub mod util;
|
||||
|
236
src/listener.rs
236
src/listener.rs
@ -1,236 +0,0 @@
|
||||
use std::{collections::HashMap, error::Error, net::SocketAddr, path::PathBuf, sync::Arc};
|
||||
|
||||
use tokio::{
|
||||
fs::File,
|
||||
io::AsyncReadExt,
|
||||
net::{TcpListener, TcpStream},
|
||||
sync::mpsc,
|
||||
};
|
||||
|
||||
use crate::{
|
||||
common::{Connection, Message},
|
||||
comms, crypto,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Listener {
|
||||
host_addr: SocketAddr,
|
||||
access_key: String,
|
||||
chunksize: usize,
|
||||
}
|
||||
|
||||
// TODO: impl Drop (?)
|
||||
|
||||
impl Listener {
|
||||
pub fn new(
|
||||
host_addr: SocketAddr,
|
||||
access_key: String,
|
||||
chunksize: usize,
|
||||
) -> Result<Arc<Self>, Box<dyn Error>> {
|
||||
Ok(Arc::new(Self {
|
||||
host_addr,
|
||||
access_key,
|
||||
chunksize,
|
||||
}))
|
||||
}
|
||||
|
||||
pub async fn start(
|
||||
self: Arc<Self>,
|
||||
tx: mpsc::Sender<Message>,
|
||||
mut kill: mpsc::Receiver<Message>,
|
||||
files: Vec<PathBuf>,
|
||||
) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
tokio::select! {
|
||||
_ = self.listen(tx, files) => Ok(()),
|
||||
_ = kill.recv() => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn listen(
|
||||
self: Arc<Self>,
|
||||
tx: mpsc::Sender<Message>,
|
||||
files: Vec<PathBuf>,
|
||||
) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
let listener = TcpListener::bind(self.host_addr).await?;
|
||||
|
||||
loop {
|
||||
let files = files.clone();
|
||||
let (mut socket, addr) = listener.accept().await?;
|
||||
tx.send(Message::ClientConnect(addr)).await?;
|
||||
let conn_tx = tx.clone();
|
||||
let err_tx = tx.clone();
|
||||
let conn_self = Arc::clone(&self);
|
||||
|
||||
match tokio::spawn(async move {
|
||||
conn_self
|
||||
.connection(&mut socket, addr, conn_tx, &files)
|
||||
.await
|
||||
})
|
||||
.await
|
||||
{
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
let err_msg = format!("{}: {}", addr, e);
|
||||
err_tx.send(Message::Error(err_msg)).await?;
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
async fn connection(
|
||||
&self,
|
||||
socket: &mut TcpStream,
|
||||
addr: SocketAddr,
|
||||
tx: mpsc::Sender<Message>,
|
||||
files: &Vec<PathBuf>,
|
||||
) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
let mut connection = Connection::new(socket).await?;
|
||||
|
||||
if !self.authorize(&mut connection).await? {
|
||||
return Ok::<(), Box<dyn Error + Send + Sync>>(());
|
||||
}
|
||||
|
||||
let index = self.metadata_handler(&mut connection, files).await?;
|
||||
tx.send(Message::ConnectionReady).await?;
|
||||
self.request_handler(&mut connection, &index).await?;
|
||||
tx.send(Message::ClientDisconnect(addr)).await?;
|
||||
|
||||
Ok::<(), Box<dyn Error + Send + Sync>>(())
|
||||
}
|
||||
|
||||
async fn authorize(
|
||||
&self,
|
||||
conn: &mut Connection<'_>,
|
||||
) -> Result<bool, Box<dyn Error + Send + Sync>> {
|
||||
let buf = comms::recv(&mut conn.reader, Some(&mut conn.cipher)).await?;
|
||||
let key = String::from_utf8(buf)?;
|
||||
let msg: Vec<u8>;
|
||||
let res: bool;
|
||||
|
||||
if key != self.access_key {
|
||||
res = false;
|
||||
msg = b"DISCONNECT".to_vec();
|
||||
} else {
|
||||
res = true;
|
||||
msg = b"OK".to_vec();
|
||||
}
|
||||
|
||||
comms::send(
|
||||
&mut conn.writer,
|
||||
Some(&mut conn.cipher),
|
||||
Some(&mut conn.rng),
|
||||
&msg,
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(res)
|
||||
}
|
||||
|
||||
async fn metadata(
|
||||
&self,
|
||||
files: &Vec<PathBuf>,
|
||||
) -> Result<
|
||||
(usize, Vec<(String, u64, String)>, HashMap<String, PathBuf>),
|
||||
Box<dyn Error + Send + Sync>,
|
||||
> {
|
||||
let mut metadata: Vec<(String, u64, String)> = Vec::new();
|
||||
let mut index = HashMap::new();
|
||||
|
||||
for path in files {
|
||||
let split: Vec<&str> = path.to_str().unwrap().split('/').collect();
|
||||
let name = split[split.len() - 1].to_string();
|
||||
let handle = File::open(path).await?;
|
||||
let size = handle.metadata().await?.len();
|
||||
let hash = crypto::try_hash(path)?;
|
||||
|
||||
if size > 0 {
|
||||
metadata.push((name, size, hash.clone()));
|
||||
index.insert(hash, path.clone());
|
||||
}
|
||||
}
|
||||
|
||||
Ok((metadata.len(), metadata, index))
|
||||
}
|
||||
|
||||
async fn metadata_handler(
|
||||
&self,
|
||||
conn: &mut Connection<'_>,
|
||||
files: &Vec<PathBuf>,
|
||||
) -> Result<HashMap<String, PathBuf>, Box<dyn Error + Send + Sync>> {
|
||||
let (amt, metadata, index) = self.metadata(files).await?;
|
||||
let msg = amt.to_string().as_bytes().to_vec();
|
||||
|
||||
comms::send(
|
||||
&mut conn.writer,
|
||||
Some(&mut conn.cipher),
|
||||
Some(&mut conn.rng),
|
||||
&msg,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let buf = comms::recv(&mut conn.reader, Some(&mut conn.cipher)).await?;
|
||||
let msg = String::from_utf8(buf)?;
|
||||
|
||||
if msg != "AMT" {
|
||||
return Err("Broken message sequence".into());
|
||||
}
|
||||
|
||||
for file in metadata {
|
||||
let msg = format!("{}:{}:{}", file.0, file.1, file.2)
|
||||
.as_bytes()
|
||||
.to_vec();
|
||||
|
||||
comms::send(
|
||||
&mut conn.writer,
|
||||
Some(&mut conn.cipher),
|
||||
Some(&mut conn.rng),
|
||||
&msg,
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
|
||||
Ok(index)
|
||||
}
|
||||
|
||||
async fn request_handler(
|
||||
&self,
|
||||
conn: &mut Connection<'_>,
|
||||
index: &HashMap<String, PathBuf>,
|
||||
) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
loop {
|
||||
let buf = comms::recv(&mut conn.reader, Some(&mut conn.cipher)).await?;
|
||||
let cmd = String::from_utf8(buf)?;
|
||||
|
||||
if cmd == "DISCONNECT" {
|
||||
break;
|
||||
}
|
||||
|
||||
let mut file = File::open(index[&cmd].clone()).await?;
|
||||
let mut remaining = file.metadata().await?.len();
|
||||
let mut send_buf = vec![0u8; self.chunksize];
|
||||
|
||||
while remaining != 0 {
|
||||
let n = file.read(&mut send_buf).await?;
|
||||
|
||||
comms::send(
|
||||
&mut conn.writer,
|
||||
Some(&mut conn.cipher),
|
||||
Some(&mut conn.rng),
|
||||
&send_buf[..n].to_vec(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
remaining -= n as u64;
|
||||
}
|
||||
|
||||
let buf = comms::recv(&mut conn.reader, Some(&mut conn.cipher)).await?;
|
||||
let msg = String::from_utf8(buf)?;
|
||||
|
||||
if msg == "ERROR" {
|
||||
return Err("Incomplete file request (hashes don't match)".into());
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
162
src/server.rs
Normal file
162
src/server.rs
Normal file
@ -0,0 +1,162 @@
|
||||
use std::{collections::HashMap, error::Error, net::SocketAddr, path::PathBuf, sync::Arc};
|
||||
|
||||
use tokio::{
|
||||
fs::File,
|
||||
io::AsyncReadExt,
|
||||
net::{TcpListener, TcpStream},
|
||||
sync::mpsc,
|
||||
};
|
||||
|
||||
use crate::{crypto::Crypto, sockets::SocketHandler, util::FileInfo};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Server {
|
||||
addr: SocketAddr,
|
||||
key: String,
|
||||
chunksize: usize,
|
||||
metadata: Vec<FileInfo>,
|
||||
index: HashMap<String, PathBuf>,
|
||||
}
|
||||
|
||||
impl Server {
|
||||
pub fn new(
|
||||
addr: SocketAddr,
|
||||
key: String,
|
||||
chunksize: usize,
|
||||
metadata: Vec<FileInfo>,
|
||||
index: HashMap<String, PathBuf>,
|
||||
) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
addr,
|
||||
key,
|
||||
chunksize,
|
||||
metadata,
|
||||
index,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn start(
|
||||
self: Arc<Self>,
|
||||
mut kill: mpsc::Receiver<()>,
|
||||
) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
tokio::select! {
|
||||
_ = self.listen() => Ok(()),
|
||||
_ = kill.recv() => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn listen(self: Arc<Self>) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
let listener = TcpListener::bind(self.addr).await?;
|
||||
|
||||
loop {
|
||||
let this_self = self.clone();
|
||||
let (mut socket, addr) = listener.accept().await?;
|
||||
|
||||
// log: new client connected: <addr>
|
||||
|
||||
match tokio::spawn(async move { this_self.connection(&mut socket).await }).await {
|
||||
Ok(_) => {}
|
||||
Err(e) => eprintln!("Error during connection ({}): {}", addr, e),
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
async fn connection(&self, socket: &mut TcpStream) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
let mut handler = SocketHandler::new(socket);
|
||||
let crypto = Crypto::new(&mut handler, true).await?;
|
||||
handler.set_crypto(crypto);
|
||||
|
||||
if !self.authorize(&mut handler).await? {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
self.metadata(&mut handler).await?;
|
||||
self.requests(&mut handler).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn authorize(
|
||||
&self,
|
||||
handler: &mut SocketHandler<'_>,
|
||||
) -> Result<bool, Box<dyn Error + Send + Sync>> {
|
||||
let buf = handler.recv().await?;
|
||||
let key = String::from_utf8(buf)?;
|
||||
|
||||
let is_valid: bool;
|
||||
let res_msg: Vec<u8>;
|
||||
|
||||
if key != self.key {
|
||||
is_valid = false;
|
||||
res_msg = b"DISCONNECT".to_vec();
|
||||
} else {
|
||||
is_valid = true;
|
||||
res_msg = b"VALID".to_vec();
|
||||
}
|
||||
|
||||
handler.send(&res_msg).await?;
|
||||
|
||||
Ok(is_valid)
|
||||
}
|
||||
|
||||
async fn metadata(
|
||||
&self,
|
||||
handler: &mut SocketHandler<'_>,
|
||||
) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
let amt = self.metadata.len();
|
||||
let msg = amt.to_string().as_bytes().to_vec();
|
||||
|
||||
handler.send(&msg).await?;
|
||||
|
||||
let buf = handler.recv().await?;
|
||||
let res_amt = String::from_utf8(buf)?.trim().parse::<usize>()?;
|
||||
|
||||
if res_amt != amt {
|
||||
return Err("Broken message sequence during metadata exchange".into());
|
||||
}
|
||||
|
||||
for file in &self.metadata {
|
||||
let msg = format!("{}:{}:{}", file.name, file.size, file.hash)
|
||||
.as_bytes()
|
||||
.to_vec();
|
||||
handler.send(&msg).await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn requests(
|
||||
&self,
|
||||
handler: &mut SocketHandler<'_>,
|
||||
) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
loop {
|
||||
let buf = handler.recv().await?;
|
||||
let hash = String::from_utf8(buf)?;
|
||||
let hash = hash.trim();
|
||||
|
||||
if hash == "DISCONNECT" {
|
||||
break;
|
||||
}
|
||||
|
||||
let mut file = File::open(self.index[hash].clone()).await?;
|
||||
let mut remaining = file.metadata().await?.len();
|
||||
let mut sendbuf = vec![0u8; self.chunksize];
|
||||
|
||||
while remaining != 0 {
|
||||
let n = file.read(&mut sendbuf).await?;
|
||||
handler.send(&sendbuf[..n].to_vec()).await?;
|
||||
remaining -= n as u64;
|
||||
}
|
||||
|
||||
let buf = handler.recv().await?;
|
||||
let confirmation = String::from_utf8(buf)?;
|
||||
let confirmation = confirmation.trim();
|
||||
|
||||
if confirmation != hash {
|
||||
return Err("Unsuccessful file transfer, hashes don't match".into());
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
@ -36,26 +36,26 @@ impl<'a> SocketHandler<'a> {
|
||||
self.crypto = Some(crypto);
|
||||
}
|
||||
|
||||
pub async fn sender(&mut self, data: &[u8]) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
pub async fn send(&mut self, data: &[u8]) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
let data = match &self.crypto {
|
||||
Some(c) => c.encrypt(data).await?,
|
||||
None => data.to_vec(),
|
||||
};
|
||||
|
||||
self.send(&data).await?;
|
||||
self.send_raw(&data).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn send(&mut self, data: &[u8]) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
pub async fn send_raw(&mut self, data: &[u8]) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
self.writer.write_all(data).await?;
|
||||
self.writer.flush().await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn receiver(&mut self) -> Result<Vec<u8>, Box<dyn Error + Send + Sync>> {
|
||||
let mut buf = self.recv().await?;
|
||||
pub async fn recv(&mut self) -> Result<Vec<u8>, Box<dyn Error + Send + Sync>> {
|
||||
let mut buf = self.recv_raw().await?;
|
||||
buf.pop();
|
||||
buf = general_purpose::STANDARD_NO_PAD.decode(&buf)?.to_vec();
|
||||
|
||||
@ -67,7 +67,7 @@ impl<'a> SocketHandler<'a> {
|
||||
Ok(data)
|
||||
}
|
||||
|
||||
pub async fn recv(&mut self) -> Result<Vec<u8>, Box<dyn Error + Send + Sync>> {
|
||||
pub async fn recv_raw(&mut self) -> Result<Vec<u8>, Box<dyn Error + Send + Sync>> {
|
||||
let mut buf = Vec::new();
|
||||
let n = self.reader.read_until(b':', &mut buf).await?;
|
||||
|
||||
|
55
src/util.rs
55
src/util.rs
@ -1,4 +1,8 @@
|
||||
use std::{error::Error, fs, net::SocketAddr, path::PathBuf};
|
||||
use std::{collections::HashMap, error::Error, fs, net::SocketAddr, path::PathBuf};
|
||||
|
||||
use tokio::{fs::File, io::BufWriter};
|
||||
|
||||
use crate::crypto;
|
||||
|
||||
const PUBLIC_IPV4: &str = "https://ipinfo.io/ip";
|
||||
const PUBLIC_IPV6: &str = "https://ipv6.icanhazip.com";
|
||||
@ -15,9 +19,9 @@ impl Ip {
|
||||
let addr = match self {
|
||||
Ip::V4 => PUBLIC_IPV4,
|
||||
Ip::V6 => PUBLIC_IPV6,
|
||||
Local => {
|
||||
Ip::Local => {
|
||||
let addr_str = format!("127.0.0.1:{}", port);
|
||||
return addr_str.parse::<SocketAddr>();
|
||||
return Ok(addr_str.parse::<SocketAddr>()?);
|
||||
}
|
||||
};
|
||||
|
||||
@ -28,6 +32,19 @@ impl Ip {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct FileInfo {
|
||||
pub name: String,
|
||||
pub size: u64,
|
||||
pub hash: String,
|
||||
}
|
||||
|
||||
impl FileInfo {
|
||||
pub fn new(name: String, size: u64, hash: String) -> Self {
|
||||
Self { name, size, hash }
|
||||
}
|
||||
}
|
||||
|
||||
fn filepaths(
|
||||
infile: Option<PathBuf>,
|
||||
files: Option<Vec<PathBuf>>,
|
||||
@ -49,3 +66,35 @@ fn filepaths(
|
||||
|
||||
Ok(filepaths)
|
||||
}
|
||||
|
||||
pub async fn metadata(
|
||||
files: &Vec<PathBuf>,
|
||||
) -> Result<(Vec<(String, u64, String)>, HashMap<String, PathBuf>), Box<dyn Error + Send + Sync>> {
|
||||
let mut metadata = Vec::new();
|
||||
let mut index = HashMap::new();
|
||||
|
||||
for path in files {
|
||||
let split = path.to_str().unwrap().split('/').collect::<Vec<&str>>();
|
||||
let name = split[split.len() - 1].to_string();
|
||||
let handle = File::open(path).await?;
|
||||
let size = handle.metadata().await?.len();
|
||||
let hash = crypto::try_hash(path)?;
|
||||
|
||||
if size > 0 {
|
||||
metadata.push((name, size, hash.clone()));
|
||||
index.insert(hash, path.clone());
|
||||
}
|
||||
}
|
||||
|
||||
Ok((metadata, index))
|
||||
}
|
||||
|
||||
pub async fn new_file(
|
||||
mut path: PathBuf,
|
||||
name: &str,
|
||||
) -> Result<(BufWriter<File>, PathBuf), Box<dyn Error + Send + Sync>> {
|
||||
path.push(name);
|
||||
let handle = File::create(&path).await?;
|
||||
|
||||
Ok((BufWriter::new(handle), path))
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user