DevComp/ruck
1
0
Fork 0
mirror of https://github.com/CompeyDev/ruck.git synced 2025-01-07 11:29:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge pull request #2 from rictorlome/encrypt

Browse source 
Encrypt
This commit is contained in:
rictorlome 2022-02-13 11:10:43 -05:00 committed by GitHub
commit 9479653364
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 522 additions and 87 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

82
Cargo.lock generated
View file

@ -75,6 +75,15 @@ version = "1.3.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
[[package]]
name = "blake2"
version = "0.10.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b94ba84325db59637ffc528bbe8c7f86c02c57cff5c0e2b9b00f9a851f42f309"
dependencies = [
"digest 0.10.2",
]
[[package]]
name = "block-buffer"
version = "0.10.2"
@ -179,7 +188,7 @@ checksum = "0b9fdf9972b2bd6af2d913799d9ebc165ea4d2e65878e329d9c6b372c4491b61"
dependencies = [
"byteorder",
"digest 0.9.0",
"rand_core",
"rand_core 0.5.1",
"subtle",
"zeroize",
]
@ -337,7 +346,18 @@ checksum = "8fc3cb4d91f53b50155bdcfd23f6a4c39ae1969c2ae85982b135750cccaf5fce"
dependencies = [
"cfg-if",
"libc",
"wasi",
"wasi 0.9.0+wasi-snapshot-preview1",
]
[[package]]
name = "getrandom"
version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "418d37c8b1d42553c93648be529cb70f920d3baf8ef469b74b9638df426e0b4c"
dependencies = [
"cfg-if",
"libc",
"wasi 0.10.2+wasi-snapshot-preview1",
]
[[package]]
@ -605,6 +625,12 @@ dependencies = [
"universal-hash",
]
[[package]]
name = "ppv-lite86"
version = "0.2.16"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "eb9f9e6e233e5c4a35559a617bf40a4ec447db2e84c20b55a6f83167b7e57872"
[[package]]
name = "proc-macro-error"
version = "1.0.4"
@ -647,13 +673,53 @@ dependencies = [
"proc-macro2",
]
[[package]]
name = "rand"
version = "0.8.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2e7573632e6454cf6b99d7aac4ccca54be06da05aca2ef7423d22d27d4d4bcd8"
dependencies = [
"libc",
"rand_chacha",
"rand_core 0.6.3",
"rand_hc",
]
[[package]]
name = "rand_chacha"
version = "0.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88"
dependencies = [
"ppv-lite86",
"rand_core 0.6.3",
]
[[package]]
name = "rand_core"
version = "0.5.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "90bde5296fc891b0cef12a6d03ddccc162ce7b2aff54160af9338f8d40df6d19"
dependencies = [
"getrandom",
"getrandom 0.1.16",
]
[[package]]
name = "rand_core"
version = "0.6.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d34f1408f55294453790c48b2f1ebbb1c5b4b7563eb1f418bcfcfdbb06ebb4e7"
dependencies = [
"getrandom 0.2.4",
]
[[package]]
name = "rand_hc"
version = "0.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d51e9f596de227fda2ea6c84607f5558e196eeaf43c986b724ba4fb8fdf497e7"
dependencies = [
"rand_core 0.6.3",
]
[[package]]
@ -672,9 +738,11 @@ dependencies = [
"aes-gcm",
"anyhow",
"bincode",
"blake2",
"bytes",
"clap",
"futures",
"rand",
"serde",
"spake2",
"tokio",
@ -749,7 +817,7 @@ checksum = "7139ade210556eca57dfc299ec2454846ab6be09232eb1139a36e285ae7fd48e"
dependencies = [
"curve25519-dalek",
"hkdf",
"rand_core",
"rand_core 0.5.1",
"sha2",
]
@ -897,6 +965,12 @@ version = "0.9.0+wasi-snapshot-preview1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cccddf32554fecc6acb585f82a32a72e28b48f8c4c1883ddfeeeaa96f7d8e519"
[[package]]
name = "wasi"
version = "0.10.2+wasi-snapshot-preview1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fd6fbd9a79829dd1ad0cc20627bf1ed606756a7f77edff7b66b7064f9cb327c6"
[[package]]
name = "winapi"
version = "0.3.9"

2
Cargo.toml
View file

@ -7,10 +7,12 @@ edition = "2021"
aes-gcm = "0.9.4"
anyhow = "1.0"
blake2 = "0.10.2"
bytes = { version = "1", features = ["serde"] }
bincode = "1.3.3"
clap = { version = "3.0.14", features = ["derive"] }
futures = { version = "0.3.0", features = ["thread-pool"]}
rand = "0.8.4"
serde = { version = "1.0", features = ["derive"] }
spake2 = "0.3.1"
tokio = { version = "1.16.1", features = ["full"] }

24
README.md
View file

@ -1,13 +1,17 @@
https://github.com/mcginty/snow/tree/master/examples
https://github.com/black-binary/snowstorm/tree/master/examples
https://cryptoservices.github.io/cryptography/protocols/2016/04/27/noise-protocol.html
- https://github.com/mcginty/snow/tree/master/examples
- https://github.com/black-binary/snowstorm/tree/master/examples
- https://cryptoservices.github.io/cryptography/protocols/2016/04/27/noise-protocol.html
https://github.com/libp2p
- https://github.com/libp2p
https://docs.rs/spake2/latest/spake2/
https://crates.io/crates/chacha20poly1305
https://briansmith.org/rustdoc/ring/aead/index.html
https://libsodium.gitbook.io/doc/secret-key_cryptography/secretstream
- https://docs.rs/spake2/latest/spake2/
- https://crates.io/crates/chacha20poly1305
- https://briansmith.org/rustdoc/ring/aead/index.html
- https://libsodium.gitbook.io/doc/secret-key_cryptography/secretstream
https://kerkour.com/rust-symmetric-encryption-aead-benchmark/
https://docs.rs/aes-gcm/latest/aes_gcm/
- https://kerkour.com/rust-symmetric-encryption-aead-benchmark/
- https://docs.rs/aes-gcm/latest/aes_gcm/
- https://github.com/rust-lang/flate2-rs
- https://crates.io/crates/async-compression

152
src/client.rs
View file

@ -1,51 +1,138 @@
use crate::crypto::handshake;
use crate::message::{Message, MessageStream};
use crate::file::{to_size_string, FileHandle, FileInfo};
use crate::message::{EncryptedMessage, FileNegotiationPayload, Message, MessageStream};
use anyhow::Result;
use bytes::{BufMut, Bytes, BytesMut};
use aes_gcm::Aes256Gcm;
use anyhow::{anyhow, Result};
use blake2::{Blake2s256, Digest};
use bytes::{Bytes, BytesMut};
use futures::future::try_join_all;
use futures::prelude::*;
use std::path::PathBuf;
use tokio::io;
use tokio::net::TcpStream;
use tokio_util::codec::{FramedRead, LinesCodec};
fn pass_to_bytes(password: &String) -> Bytes {
let mut hasher = Blake2s256::new();
hasher.update(password.as_bytes());
let res = hasher.finalize();
BytesMut::from(&res[..]).freeze()
}
pub async fn send(file_paths: &Vec<PathBuf>, password: &String) -> Result<()> {
// Fail early if there are problems generating file handles
let handles = get_file_handles(file_paths).await?;
// Establish connection to server
let socket = TcpStream::connect("127.0.0.1:8080").await?;
let mut stream = Message::to_stream(socket);
let (stream, key) = handshake(
// Complete handshake, returning cipher used for encryption
let (stream, cipher) = handshake(
&mut stream,
true,
Bytes::from(password.to_string()),
Bytes::from("id123"),
pass_to_bytes(password),
)
.await?;
return upload_encrypted_files(stream, file_paths, key).await;
// Complete file negotiation
let handles = negotiate_files_up(handles, stream, &cipher).await?;
// Send the value
// for path in paths.iter() {
// let b = path.to_str().unwrap().as_bytes();
// let mut buf = BytesMut::with_capacity(1024);
// buf.put(&b[..]);
// let body = buf.freeze();
// let m = Message {
// key: "abc".to_string(),
// from_sender: true,
// body: body,
// };
// stream.send(m).await.unwrap();
// }
// Upload negotiated files
// Exit
Ok(())
}
pub async fn receive(password: &String) -> Result<()> {
let socket = TcpStream::connect("127.0.0.1:8080").await?;
let mut stream = Message::to_stream(socket);
let (stream, key) = handshake(
let (stream, cipher) = handshake(
&mut stream,
false,
Bytes::from(password.to_string()),
Bytes::from("id123"),
pass_to_bytes(password),
)
.await?;
let files = negotiate_files_down(stream, &cipher).await?;
return Ok(());
}
pub async fn get_file_handles(file_paths: &Vec<PathBuf>) -> Result<Vec<FileHandle>> {
let tasks = file_paths
.into_iter()
.map(|path| FileHandle::new(path.to_path_buf()));
let handles = try_join_all(tasks).await?;
Ok(handles)
}
pub async fn negotiate_files_up(
file_handles: Vec<FileHandle>,
stream: &mut MessageStream,
cipher: &Aes256Gcm,
) -> Result<Vec<FileHandle>> {
let files = file_handles.iter().map(|fh| fh.to_file_info()).collect();
let msg = EncryptedMessage::FileNegotiationMessage(FileNegotiationPayload { files });
let server_msg = msg.to_encrypted_message(cipher)?;
stream.send(server_msg).await?;
let reply_payload = match stream.next().await {
Some(Ok(msg)) => match msg {
Message::EncryptedMessage(response) => response,
_ => return Err(anyhow!("Expecting encrypted message back")),
},
_ => {
return Err(anyhow!("No response to negotiation message"));
}
};
let plaintext_reply = EncryptedMessage::from_encrypted_message(cipher, &reply_payload)?;
let requested_paths: Vec<PathBuf> = match plaintext_reply {
EncryptedMessage::FileNegotiationMessage(fnm) => {
fnm.files.into_iter().map(|f| f.path).collect()
}
_ => return Err(anyhow!("Expecting file negotiation message back")),
};
Ok(file_handles
.into_iter()
.filter(|fh| requested_paths.contains(&fh.path))
.collect())
}
pub async fn negotiate_files_down(stream: &mut MessageStream, cipher: &Aes256Gcm) -> Result<()> {
let file_offer = match stream.next().await {
Some(Ok(msg)) => match msg {
Message::EncryptedMessage(response) => response,
_ => return Err(anyhow!("Expecting encrypted message back")),
},
_ => {
return Err(anyhow!("No response to negotiation message"));
}
};
let plaintext_offer = EncryptedMessage::from_encrypted_message(cipher, &file_offer)?;
let requested_infos: Vec<FileInfo> = match plaintext_offer {
EncryptedMessage::FileNegotiationMessage(fnm) => fnm.files,
_ => return Err(anyhow!("Expecting file negotiation message back")),
};
let mut stdin = FramedRead::new(io::stdin(), LinesCodec::new());
let mut files = vec![];
for path in requested_infos.into_iter() {
let mut reply = prompt_user_input(&mut stdin, &path).await;
while reply.is_none() {
reply = prompt_user_input(&mut stdin, &path).await;
}
match reply {
Some(true) => files.push(path),
_ => {}
}
}
let msg = EncryptedMessage::FileNegotiationMessage(FileNegotiationPayload { files });
let server_msg = msg.to_encrypted_message(cipher)?;
stream.send(server_msg).await?;
Ok(())
}
pub async fn upload_encrypted_files(
stream: &mut MessageStream,
file_paths: &Vec<PathBuf>,
@ -53,3 +140,26 @@ pub async fn upload_encrypted_files(
) -> Result<()> {
return Ok(());
}
pub async fn prompt_user_input(
stdin: &mut FramedRead<io::Stdin, LinesCodec>,
file_info: &FileInfo,
) -> Option<bool> {
let prompt_name = file_info.path.file_name().unwrap();
println!(
"Do you want to download {:?}? It's {:?}. (Y/n)",
prompt_name,
to_size_string(file_info.size)
);
match stdin.next().await {
Some(Ok(line)) => match line.as_str() {
"" | "Y" | "y" | "yes" | "Yes" | "YES" => Some(true),
"N" | "n" | "NO" | "no" | "No" => Some(false),
_ => {
println!("Invalid input. Please enter one of the following characters: [YyNn]");
return None;
}
},
_ => None,
}
}

53
src/crypto.rs
View file

@ -1,23 +1,29 @@
use crate::message::{HandshakeMessage, Message, MessageStream};
use crate::message::{EncryptedPayload, HandshakePayload, Message, MessageStream};
use aes_gcm::aead::{Aead, NewAead};
use aes_gcm::{Aes256Gcm, Key, Nonce}; // Or `Aes128Gcm`
use anyhow::{anyhow, Result};
use bytes::Bytes;
use futures::prelude::*;
use rand::{thread_rng, Rng};
use spake2::{Ed25519Group, Identity, Password, Spake2};
pub async fn handshake(
stream: &mut MessageStream,
up: bool,
password: Bytes,
id: Bytes,
) -> Result<(&mut MessageStream, Bytes)> {
) -> Result<(&mut MessageStream, Aes256Gcm)> {
let (s1, outbound_msg) =
Spake2::<Ed25519Group>::start_symmetric(&Password::new(password), &Identity::new(&id));
stream
.send(Message::HandshakeMessage(HandshakeMessage {
id,
msg: Bytes::from(outbound_msg),
}))
.await?;
println!("client - sending handshake msg");
let handshake_msg = Message::HandshakeMessage(HandshakePayload {
up,
id,
msg: Bytes::from(outbound_msg),
});
println!("client - handshake msg, {:?}", handshake_msg);
stream.send(handshake_msg).await?;
let first_message = match stream.next().await {
Some(Ok(msg)) => match msg {
Message::HandshakeMessage(response) => response.msg,
@ -27,10 +33,39 @@ pub async fn handshake(
return Err(anyhow!("No response to handshake message"));
}
};
println!("client - handshake msg responded to");
let key = match s1.finish(&first_message[..]) {
Ok(key_bytes) => key_bytes,
Err(e) => return Err(anyhow!(e.to_string())),
};
println!("Handshake successful. Key is {:?}", key);
return Ok((stream, Bytes::from(key)));
return Ok((stream, new_cipher(&key)));
}
pub fn new_cipher(key: &Vec<u8>) -> Aes256Gcm {
let key = Key::from_slice(&key[..]);
Aes256Gcm::new(key)
}
pub const NONCE_SIZE_IN_BYTES: usize = 96 / 8;
pub fn encrypt(cipher: &Aes256Gcm, body: &Vec<u8>) -> Result<EncryptedPayload> {
let mut arr = [0u8; NONCE_SIZE_IN_BYTES];
thread_rng().try_fill(&mut arr[..])?;
let nonce = Nonce::from_slice(&arr);
let plaintext = body.as_ref();
match cipher.encrypt(nonce, plaintext) {
Ok(body) => Ok(EncryptedPayload {
nonce: arr.to_vec(),
body,
}),
Err(_) => Err(anyhow!("Encryption error")),
}
}
pub fn decrypt(cipher: &Aes256Gcm, payload: &EncryptedPayload) -> Result<Bytes> {
let nonce = Nonce::from_slice(payload.nonce.as_ref());
match cipher.decrypt(nonce, payload.body.as_ref()) {
Ok(payload) => Ok(Bytes::from(payload)),
Err(_) => Err(anyhow!("Decryption error")),
}
}

48
src/file.rs Normal file
View file

@ -0,0 +1,48 @@
use anyhow::{anyhow, Result};
use serde::{Deserialize, Serialize};
use std::fs::Metadata;
use std::path::PathBuf;
use tokio::fs::File;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct FileInfo {
pub path: PathBuf,
pub size: u64,
}
pub struct FileHandle {
pub file: File,
pub md: Metadata,
pub path: PathBuf,
}
impl FileHandle {
pub async fn new(path: PathBuf) -> Result<FileHandle> {
let file = File::open(&path).await?;
let md = file.metadata().await?;
let fh = FileHandle { file, md, path };
return Ok(fh);
}
pub fn to_file_info(&self) -> FileInfo {
FileInfo {
path: self.path.clone(),
size: self.md.len(),
}
}
}
const SUFFIX: [&'static str; 9] = ["B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"];
// Stolen: https://gitlab.com/forkbomb9/human_bytes-rs/-/blob/master/src/lib.rs
pub fn to_size_string(size: u64) -> String {
let size = size as f64;
let base = size.log10() / 1024_f64.log10();
let mut result = format!("{:.1}", 1024_f64.powf(base - base.floor()),)
.trim_end_matches(".0")
.to_owned();
// Add suffix
result.push(' ');
result.push_str(SUFFIX[base.floor() as usize]);
result
}

6
src/main.rs
View file

@ -1,9 +1,13 @@
mod cli;
mod client;
mod crypto;
mod file;
mod message;
mod server;
#[cfg(test)]
mod tests;
use clap::Parser;
use cli::{Cli, Commands};
use client::{receive, send};
@ -16,7 +20,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
match &args.command {
Commands::Send { paths, password } => {
println!("Sending {:?}", paths);
send(&paths, password).await?;
send(paths, password).await?;
}
Commands::Receive { password } => {
println!("Receiving password {}", password);

76
src/message.rs
View file

@ -1,20 +1,92 @@
use crate::crypto::{decrypt, encrypt};
use crate::file::FileInfo;
use aes_gcm::Aes256Gcm; // Or `Aes128Gcm`
use anyhow::{anyhow, Result};
use bincode::config;
use bytes::Bytes;
use serde::{Deserialize, Serialize};
use std::error::Error;
use std::fmt;
use tokio::net::TcpStream;
use tokio_serde::{formats::SymmetricalBincode, SymmetricallyFramed};
use tokio_util::codec::{Framed, LengthDelimitedCodec};
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum Message {
HandshakeMessage(HandshakeMessage),
HandshakeMessage(HandshakePayload),
EncryptedMessage(EncryptedPayload),
ErrorMessage(RuckError),
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct HandshakeMessage {
pub struct HandshakePayload {
pub up: bool,
pub id: Bytes,
pub msg: Bytes,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct EncryptedPayload {
pub nonce: Vec<u8>,
pub body: Vec<u8>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum EncryptedMessage {
FileNegotiationMessage(FileNegotiationPayload),
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct FileNegotiationPayload {
pub files: Vec<FileInfo>,
}
impl EncryptedMessage {
pub fn from_encrypted_message(cipher: &Aes256Gcm, payload: &EncryptedPayload) -> Result<Self> {
let raw = decrypt(cipher, payload)?;
let res = match bincode::deserialize(raw.as_ref()) {
Ok(result) => result,
Err(e) => {
println!("deserialize error {:?}", e);
return Err(anyhow!("deser error"));
}
};
Ok(res)
}
pub fn to_encrypted_message(&self, cipher: &Aes256Gcm) -> Result<Message> {
let raw = match bincode::serialize(&self) {
Ok(result) => result,
Err(e) => {
println!("serialize error {:?}", e);
return Err(anyhow!("serialize error"));
}
};
let payload = encrypt(cipher, &raw)?;
Ok(Message::EncryptedMessage(payload))
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum RuckError {
NotHandshake,
SenderNotConnected,
SenderAlreadyConnected,
PairDisconnected,
}
impl fmt::Display for RuckError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "RuckError is here!")
}
}
impl Error for RuckError {
fn source(&self) -> Option<&(dyn Error + 'static)> {
Some(self)
}
}
impl Message {
pub fn to_stream(stream: TcpStream) -> MessageStream {
tokio_serde::SymmetricallyFramed::new(

155
src/server.rs
View file

@ -1,66 +1,125 @@
use crate::message::{Message, MessageStream};
use crate::message::{HandshakePayload, Message, MessageStream, RuckError};
use anyhow::{anyhow, Result};
use bytes::Bytes;
use futures::prelude::*;
use std::collections::HashMap;
use std::io;
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::{mpsc, Mutex};
use tokio::sync::{mpsc, oneshot, Mutex};
type Tx = mpsc::UnboundedSender<Message>;
type Rx = mpsc::UnboundedReceiver<Message>;
pub struct Shared {
rooms: HashMap<String, RoomInfo>,
handshakes: HashMap<Bytes, Rx>,
senders: HashMap<Bytes, Tx>,
receivers: HashMap<Bytes, Tx>,
}
type State = Arc<Mutex<Shared>>;
struct RoomInfo {
sender_tx: Tx,
}
struct Client {
is_sender: bool,
messages: MessageStream,
struct Client<'a> {
up: bool,
id: Bytes,
messages: &'a mut MessageStream,
rx: Rx,
}
impl Shared {
fn new() -> Self {
Shared {
rooms: HashMap::new(),
handshakes: HashMap::new(),
senders: HashMap::new(),
receivers: HashMap::new(),
}
}
// async fn broadcast(&mut self, sender: SocketAddr, message: Message) {
// for peer in self.peers.iter_mut() {
// if *peer.0 != sender {
// let _ = peer.1.send(message.clone());
// }
// }
// }
async fn relay<'a>(&self, client: &Client<'a>, message: Message) -> Result<()> {
println!("in relay - got client={:?}, msg {:?}", client.id, message);
match client.up {
true => match self.receivers.get(&client.id) {
Some(tx) => {
tx.send(message)?;
}
None => {
return Err(anyhow!(RuckError::PairDisconnected));
}
},
false => match self.senders.get(&client.id) {
Some(tx) => {
tx.send(message)?;
}
None => {
return Err(anyhow!(RuckError::PairDisconnected));
}
},
}
Ok(())
}
}
impl Client {
async fn new(is_sender: bool, state: State, messages: MessageStream) -> io::Result<Client> {
impl<'a> Client<'a> {
async fn new(
up: bool,
id: Bytes,
state: State,
messages: &'a mut MessageStream,
) -> Result<Client<'a>> {
let (tx, rx) = mpsc::unbounded_channel();
let room_info = RoomInfo { sender_tx: tx };
state
.lock()
.await
.rooms
.insert("abc".to_string(), room_info);
println!("server - creating client up={:?}, id={:?}", up, id);
let shared = &mut state.lock().await;
match shared.senders.get(&id) {
Some(_) if up => {
messages
.send(Message::ErrorMessage(RuckError::SenderAlreadyConnected))
.await?;
}
Some(_) => {
println!("server - adding client to receivers");
shared.receivers.insert(id.clone(), tx);
}
None if up => {
println!("server - adding client to senders");
shared.senders.insert(id.clone(), tx);
}
None => {
messages
.send(Message::ErrorMessage(RuckError::SenderNotConnected))
.await?;
}
}
Ok(Client {
is_sender,
up,
id,
messages,
rx,
})
}
async fn complete_handshake(&mut self, state: State, msg: Message) -> Result<()> {
match self.up {
true => {
let (tx, rx) = mpsc::unbounded_channel();
tx.send(msg)?;
state.lock().await.handshakes.insert(self.id.clone(), rx);
}
false => {
let shared = &mut state.lock().await;
if let Some(tx) = shared.senders.get(&self.id) {
tx.send(msg)?;
}
if let Some(mut rx) = shared.handshakes.remove(&self.id) {
drop(shared);
if let Some(msg) = rx.recv().await {
self.messages.send(msg).await?;
}
}
}
}
Ok(())
}
}
pub async fn serve() -> Result<(), Box<dyn std::error::Error>> {
pub async fn serve() -> Result<()> {
let addr = "127.0.0.1:8080".to_string();
let listener = TcpListener::bind(&addr).await?;
let state = Arc::new(Mutex::new(Shared::new()));
@ -70,8 +129,8 @@ pub async fn serve() -> Result<(), Box<dyn std::error::Error>> {
let state = Arc::clone(&state);
tokio::spawn(async move {
match handle_connection(state, stream, address).await {
Ok(_) => println!("ok"),
Err(_) => println!("err"),
Ok(_) => println!("Connection complete!"),
Err(err) => println!("Error handling connection! {:?}", err),
}
});
}
@ -81,19 +140,33 @@ pub async fn handle_connection(
state: Arc<Mutex<Shared>>,
socket: TcpStream,
addr: SocketAddr,
) -> Result<(), Box<dyn std::error::Error>> {
) -> Result<()> {
let mut stream = Message::to_stream(socket);
let first_message = match stream.next().await {
Some(Ok(msg)) => {
println!("first msg: {:?}", msg);
msg
println!("server - new conn from {:?}", addr);
let handshake_payload = match stream.next().await {
Some(Ok(Message::HandshakeMessage(payload))) => payload,
Some(Ok(_)) => {
stream
.send(Message::ErrorMessage(RuckError::NotHandshake))
.await?;
return Ok(());
}
_ => {
println!("no first message");
println!("No first message");
return Ok(());
}
};
let mut client = Client::new(true, state.clone(), stream).await?;
println!("server - received msg from {:?}", addr);
let mut client = Client::new(
handshake_payload.up,
handshake_payload.id.clone(),
state.clone(),
&mut stream,
)
.await?;
client
.complete_handshake(state.clone(), Message::HandshakeMessage(handshake_payload))
.await?;
// add client to state here
loop {
tokio::select! {
@ -104,6 +177,8 @@ pub async fn handle_connection(
result = client.messages.next() => match result {
Some(Ok(msg)) => {
println!("GOT: {:?}", msg);
let state = state.lock().await;
state.relay(&client, msg).await?;
}
Some(Err(e)) => {
println!("Error {:?}", e);

11
src/tests.rs Normal file
View file

@ -0,0 +1,11 @@
use crate::crypto::new_cipher;
use crate::file::FileHandle;
use std::path::PathBuf;
#[tokio::test]
async fn test_file_handle_nonexistent_file() {
let pb = PathBuf::new();
let fh = FileHandle::new(pb).await;
assert!(fh.is_err());
}