From 19f6b24e0f08a4eb3ad94c7d5d1e4fb77df837bc Mon Sep 17 00:00:00 2001 From: Patryk Hegenberg Date: Wed, 1 May 2024 22:38:18 +0200 Subject: [PATCH] feature(sender,receiver): added better functionality to separate room-ids --- src/cli/args.rs | 1 - src/receiver/mod.rs | 12 +++--------- src/relay/server.rs | 32 +++++++++++++++++++++++++------- src/relay/transfer.rs | 30 ++++++++++++++++++++++-------- src/sender/client.rs | 37 ++++++++++++++++++++++++++++--------- src/sender/http_client.rs | 13 ++++++++++--- src/sender/mod.rs | 4 ++++ 7 files changed, 92 insertions(+), 37 deletions(-) diff --git a/src/cli/args.rs b/src/cli/args.rs index 894086c..c377826 100644 --- a/src/cli/args.rs +++ b/src/cli/args.rs @@ -3,7 +3,6 @@ use crate::relay; use crate::sender; use clap::{Parser, Subcommand}; use std::{env, sync::Arc}; -use tokio::sync::mpsc; use tracing::debug; /// This struct defines the CLI arguments and subcommands for the caesar command line application. diff --git a/src/receiver/mod.rs b/src/receiver/mod.rs index 5377490..db2f615 100644 --- a/src/receiver/mod.rs +++ b/src/receiver/mod.rs @@ -51,16 +51,10 @@ pub async fn start_receiver(relay: &str, name: &str) { debug!("Got room_id from Server: {:?}", res); let res_ip = res.ip + ":9000"; - if let Err(e) = start_ws_com(res_ip.as_str(), res.room_id[0].as_str()).await { + if let Err(e) = start_ws_com(res_ip.as_str(), res.local_room_id.as_str()).await { debug!("Failed to connect local with first room_id: {e}"); - if let Err(e) = start_ws_com(res_ip.as_str(), res.room_id[1].as_str()).await { - debug!("Failed to connect local with second room_id: {e}"); - if let Err(e) = start_ws_com(relay, res.room_id[0].as_str()).await { - debug!("Failed to connect remote with first room_id: {e}"); - if let Err(e) = start_ws_com(relay, res.room_id[1].as_str()).await { - error!("Failed to accomplish transfer with error: {e}"); - } - } + if let Err(e) = start_ws_com(relay, res.relay_room_id.as_str()).await { + debug!("Failed to connect remote with first room_id: {e}"); } } } diff --git a/src/relay/server.rs b/src/relay/server.rs index e4e4034..749ea8e 100644 --- a/src/relay/server.rs +++ b/src/relay/server.rs @@ -292,17 +292,34 @@ pub async fn upload_info( .find(|request| request.name == payload.name) { Some(request) => { - request.room_id.push(payload.room_id); - debug!("Found Transfer and updated"); - (StatusCode::OK, Json(request.clone())) + debug!("Found Transfer"); + debug!("Request is: {:?}", request); + if request.relay_room_id.is_empty() { + request.relay_room_id = payload.relay_room_id; + debug!("Found Transfer and updated"); + debug!("request is: {:#?}", request); + (StatusCode::OK, Json(request.clone())) + } else { + request.local_room_id = payload.local_room_id; + debug!("Found Transfer and updated"); + debug!("request is: {:#?}", request); + (StatusCode::OK, Json(request.clone())) + } } None => { - let mut t_request = TransferResponse { + let mut local = String::from(""); + let mut relay = String::from(""); + if payload.relay_room_id.is_empty() { + local = payload.local_room_id; + } else { + relay = payload.relay_room_id; + } + let t_request = TransferResponse { name: payload.name, ip: payload.ip, - room_id: Vec::new(), + local_room_id: local, + relay_room_id: relay, }; - t_request.room_id.push(payload.room_id); data.transfers.push(t_request.clone()); debug!("New TransferRequest created"); @@ -330,7 +347,8 @@ pub async fn download_info( Json(TransferResponse { name: String::from(""), ip: String::from(""), - room_id: vec![String::from("")], + local_room_id: String::from(""), + relay_room_id: String::from(""), }), ) } diff --git a/src/relay/transfer.rs b/src/relay/transfer.rs index a69329d..30da55d 100644 --- a/src/relay/transfer.rs +++ b/src/relay/transfer.rs @@ -4,11 +4,17 @@ use serde::{Deserialize, Serialize}; pub struct TransferRequest { pub name: String, pub ip: String, - pub room_id: String, + pub local_room_id: String, + pub relay_room_id: String, } impl TransferRequest { - pub fn new(name: String, ip: String, room_id: String) -> Self { - Self { name, ip, room_id } + pub fn new(name: String, ip: String, local_room_id: String, relay_room_id: String) -> Self { + Self { + name, + ip, + local_room_id, + relay_room_id, + } } } @@ -16,12 +22,18 @@ impl TransferRequest { pub struct TransferResponse { pub name: String, pub ip: String, - pub room_id: Vec, + pub local_room_id: String, + pub relay_room_id: String, } impl TransferResponse { - pub fn new(name: String, ip: String, room_id: Vec) -> Self { - Self { name, ip, room_id } + pub fn new(name: String, ip: String, local_room_id: String, relay_room_id: String) -> Self { + Self { + name, + ip, + local_room_id, + relay_room_id, + } } } #[cfg(test)] @@ -33,13 +45,15 @@ mod tests { let transfer = TransferResponse { name: "Test".to_string(), ip: "127.0.0.1".to_string(), - room_id: vec!["This_is_a_test_room_id".to_string()], + local_room_id: "This_is_a_test_room_id".to_string(), + relay_room_id: "This_is_a_test_room_id".to_string(), }; assert_eq!( TransferResponse::new( "Test".to_string(), "127.0.0.1".to_string(), - vec!["This_is_a_test_room_id".to_string()], + "This_is_a_test_room_id".to_string(), + "This_is_a_test_room_id".to_string(), ), transfer ) diff --git a/src/sender/client.rs b/src/sender/client.rs index c00afaa..4c57d2c 100644 --- a/src/sender/client.rs +++ b/src/sender/client.rs @@ -1,5 +1,5 @@ use crate::sender::http_client::send_info; -use crate::sender::util::{generate_random_name, hash_random_name}; +use crate::sender::util::hash_random_name; use crate::shared::{ packets::{ list_packet, packet::Value, ChunkPacket, HandshakePacket, HandshakeResponsePacket, @@ -126,7 +126,13 @@ impl Context { /// appended to the room id to create a URL. The URL is then printed to the /// console using the qr2term library. Finally, the function prints a /// message to the console with the URL. -fn on_create_room(context: &Context, id: String, relay: String, transfer_name: String) -> Status { +fn on_create_room( + context: &Context, + id: String, + relay: String, + transfer_name: String, + is_local: bool, +) -> Status { debug!("Creating room on: {relay}"); let base64 = general_purpose::STANDARD.encode(&context.hmac); let url = format!("{}-{}", id, base64); @@ -142,7 +148,7 @@ fn on_create_room(context: &Context, id: String, relay: String, transfer_name: S .enable_all() .build() .unwrap() - .block_on(send_info(&server_url, &h_name, send_url.as_str())) + .block_on(send_info(&server_url, &h_name, send_url.as_str(), is_local)) }) .join() .unwrap(); @@ -562,13 +568,16 @@ fn on_message( message: WebSocketMessage, relay: String, transfer_name: String, + is_local: bool, ) -> Status { if message.is_text() { let text = message.into_text().unwrap(); let packet = serde_json::from_str(&text).unwrap(); return match packet { - JsonPacketResponse::Create { id } => on_create_room(context, id, relay, transfer_name), + JsonPacketResponse::Create { id } => { + on_create_room(context, id, relay, transfer_name, is_local) + } JsonPacketResponse::Join { size } => on_join_room(context, size), JsonPacketResponse::Leave { index } => on_leave_room(context, index), JsonPacketResponse::Error { message } => on_error(message), @@ -621,6 +630,7 @@ pub async fn start( room_id: Option, relay: String, transfer_name: String, + is_local: bool, ) { // Create a vector to store metadata about each file that will be sent. let mut files = vec![]; @@ -715,7 +725,13 @@ pub async fn start( // client. let incoming_handler = incoming.try_for_each(|message| { // Call the `on_message` function to handle the incoming message. - match on_message(&mut context, message, relay.clone(), transfer_name.clone()) { + match on_message( + &mut context, + message, + relay.clone(), + transfer_name.clone(), + is_local, + ) { // If the status is `Status::Exit`, the transfer is complete. Print a message to // stdout and exit the function. Status::Exit() => { @@ -839,7 +855,8 @@ mod tests { "b531e87d-e51a-4507-94f4-335cbe2d32f3-Nc5skZReq7qJN7INwckyAZLWEEbxsrFfH/692tUNgkM=" .to_string(), String::from("0.0.0.0:8000"), - String::from("Test") + String::from("Test"), + true, ), Status::Continue() ); @@ -926,11 +943,12 @@ mod tests { &mut context, WebSocketMessage::Text(r#"{"type":"leave","index":5}"#.to_string()), String::from("0.0.0.0:8000"), - String::from("Test") + String::from("Test"), + true, ), Status::Continue() ); - assert_eq!(on_message(&mut context, WebSocketMessage::Text(r#"{"type":"create","id":"b531e87d-e51a-4507-94f4-335cbe2d32f3-Nc5skZReq7qJN7INwckyAZLWEEbxsrFfH/692tUNgkM="}"#.to_string()), String::from("0.0.0.0:8000"), String::from("Test")), Status::Continue()); + assert_eq!(on_message(&mut context, WebSocketMessage::Text(r#"{"type":"create","id":"b531e87d-e51a-4507-94f4-335cbe2d32f3-Nc5skZReq7qJN7INwckyAZLWEEbxsrFfH/692tUNgkM="}"#.to_string()), String::from("0.0.0.0:8000"), String::from("Test"), true), Status::Continue()); assert_eq!( on_message( &mut context, @@ -938,7 +956,8 @@ mod tests { r#"{"type":"error","message":"Error Message: Test"}"#.to_string() ), String::from("0.0.0.0:8000"), - String::from("Test") + String::from("Test"), + true ), Status::Err("Error Message: Test".to_string()) ); diff --git a/src/sender/http_client.rs b/src/sender/http_client.rs index 053dc6f..8582c3f 100644 --- a/src/sender/http_client.rs +++ b/src/sender/http_client.rs @@ -1,5 +1,5 @@ use std::collections::HashMap; -use tracing::error; +use tracing::{debug, error}; use local_ip_address::{local_ip, local_ipv6}; use reqwest::blocking::Client; @@ -7,7 +7,7 @@ use tokio::task; type Result = std::result::Result>; -pub async fn send_info(relay: &str, name: &str, room_id: &str) -> Result { +pub async fn send_info(relay: &str, name: &str, room_id: &str, is_local: bool) -> Result { let url = relay.to_string(); let sender_ip = match local_ipv6() { Ok(ip) => ip, @@ -24,11 +24,18 @@ pub async fn send_info(relay: &str, name: &str, room_id: &str) -> Result let mut map = HashMap::new(); map.insert("name", String::from(name)); map.insert("ip", ip_str); - map.insert("room_id", String::from(room_id)); + if is_local { + map.insert("local_room_id", String::from(room_id)); + map.insert("relay_room_id", String::from("")); + } else { + map.insert("relay_room_id", String::from(room_id)); + map.insert("local_room_id", String::from("")); + } map }; let room_id = room_id.to_string(); + debug!("Trying to send Request."); let result: Result = task::spawn_blocking(move || { let client = Client::new(); client diff --git a/src/sender/mod.rs b/src/sender/mod.rs index e64b044..a57d52a 100644 --- a/src/sender/mod.rs +++ b/src/sender/mod.rs @@ -79,6 +79,7 @@ pub async fn start_sender(relay: Arc, files: Arc>) { relay.clone(), Arc::new(rand_name.clone()), tx.clone(), + false, ) .await }); @@ -90,6 +91,7 @@ pub async fn start_sender(relay: Arc, files: Arc>) { local_relay.clone(), Arc::new(local_rand_name.clone()), local_tx.clone(), + true, ) .await }); @@ -142,6 +144,7 @@ async fn connect_to_server( message_server: Arc, transfer_name: Arc, tx: mpsc::Sender<()>, + is_local: bool, ) { let url = format!("ws://{}/ws", relay); let message_relay = format!("{}", message_server); @@ -167,6 +170,7 @@ async fn connect_to_server( Some(room_id), message_relay.to_string(), transfer_name.clone(), + is_local, ) .await; tx.send(()).await.unwrap();