git/caesar-transfer
1
0
Fork 0
Code Issues Pull requests Projects Releases 2 Packages Wiki Activity Actions

project: first attempt of reimplemantation of caesar with websockets

Browse source
This commit is contained in:
Patryk Hegenberg 2024-04-26 13:03:52 +02:00
parent acfc0a51c7
commit 11531b7d50
19 changed files with 3820 additions and 1066 deletions
Download patch file Download diff file Expand all files Collapse all files

1305
Cargo.lock generated
View file

File diff suppressed because it is too large Load diff

53
Cargo.toml
View file

@ -1,30 +1,45 @@
[package]
name = "caesar"
name = "caesar_test"
version = "0.0.1"
edition = "2021"
authors = ["Manuel Keidel", "Patryk Hegenberg", "Krzysztof Stankiewicz"]
build = "src/build.rs"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[[bin]]
name = "caesar"
path = "src/main.rs"
[dependencies]
clap = { version = "4.5.4", features = ["derive"] }
reqwest = { version = "0.12.3", features = ["blocking", "json"] }
hyper = { version = "1", features = ["full"] }
tokio = { version = "1", features = ["full"] }
http-body-util = "0.1"
hyper-util = { version = "0.1", features = ["full"] }
futures-util = "0.3"
tungstenite = "0.20.1"
tokio = { version = "1.28.1", features = ["full"] }
tokio-tungstenite = { version = "0.20.1" }
serde_json = { version = "1.0" }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
axum = "0.7.5"
rand = "0.8.5"
uuid = { version = "1.3.2", features = ["v4"] }
tracing = "0.1.40"
tracing-subscriber = {version = "0.3.18", features = ["env-filter"]}
env_logger = "0.11.3"
log = "0.4.21"
dotenv = { version = "0.15.0", features = ["clap"] }
tower-http = {version = "0.5.2", features = ["fs"]}
tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
dotenv = { version = "0.15.0", features = ["clap", "cli"] }
clap = { version = "4.5.4", features = ["derive"] }
flume = { git = "https://github.com/zesterer/flume", rev = "80d19c49" }
prost = "0.11.9"
prost-types = "0.11.9"
base64 = "0.21.2"
url = "2.4.0"
p256 = { version = "0.13.2", features = ["ecdh"] }
hmac = "0.12.1"
sha2 = "0.10.7"
rand = { version = "0.8.5", features = ["getrandom"] }
aes-gcm = "0.10.3"
sanitize-filename = "0.4.0"
qr2term = "0.3.1"
axum = { version = "0.7.5", features = ["ws"] }
tower-http = { version = "0.5.2", features = ["fs", "trace"] }
axum-client-ip = "0.6.0"
lazy_static = "1.4.0"
local-ip-address = "0.6.1"
sha2 = "0.10.8"
hex = "0.4.3"
axum-extra = { version = "0.9.3", features = ["typed-header"] }
headers = "0.4"
tower = { version = "0.4", features = ["util"] }
[build-dependencies]
prost-build = "0.5"

42
packets.proto Normal file
View file

@ -0,0 +1,42 @@
syntax = "proto3";
package Packets;
message HandshakePacket {
bytes publicKey = 1;
bytes signature = 2;
}
message HandshakeResponsePacket {
bytes publicKey = 1;
bytes signature = 2;
}
message ListPacket {
message Entry {
uint32 index = 1;
uint64 size = 2;
string name = 3;
}
repeated Entry entries = 1;
}
message ProgressPacket {
uint32 index = 1;
uint32 progress = 2;
}
message ChunkPacket {
uint32 sequence = 1;
bytes chunk = 2;
}
message Packet {
oneof value {
HandshakePacket handshake = 1;
HandshakeResponsePacket handshakeResponse = 2;
ListPacket list = 3;
ProgressPacket progress = 4;
ChunkPacket chunk = 5;
}
}

22
src/build.rs Normal file
View file

@ -0,0 +1,22 @@
// This build script is invoked by cargo when it is building our crate. It is responsible for
// generating code based on other protocol buffer definitions.
//
// Specifically, it generates Rust code from the `packets.proto` file in the root of our
// crate. This generated code is then compiled into our final binary.
//
// The `prost_build` crate is responsible for doing the actual work of generating code from
// protocol buffer definitions. We're passing it the path to our `.proto` file and the root
// directory of our crate.
extern crate prost_build;
fn main() {
// Invoke the `compile_protos` function from the `prost_build` crate. This function takes
// two arguments: a list of `.proto` files to compile and the root directory of our crate.
// It returns a `Result` indicating whether the compilation was successful or not.
//
// The `.unwrap()` method is then called on the `Result` to panic if the compilation
// failed. This is okay in a build script because it will stop the build process and
// prevent our code from being built.
prost_build::compile_protos(&["packets.proto"], &["."]).unwrap();
}

213
src/cli/args.rs
View file

@ -1,14 +1,44 @@
use crate::relay::server;
use crate::{
receiver::client as receiver,
sender::{client as sender, server::serf_file},
};
use crate::receiver;
use crate::relay;
use crate::sender;
use clap::{Parser, Subcommand};
use tracing::{debug, error, info};
use std::env;
use tracing::debug;
/// This struct defines the CLI arguments and subcommands for the caesar command line application.
///
/// The #[derive(Parser, Debug)] macro generates code that:
/// - parses the command line arguments using the clap library
/// - provides a Debug implementation for the struct
///
/// The #[command(version, about, long_about = None)] macro generates code that:
/// - defines the version and about strings for the application
/// - specifies that there is no long about help text
///
/// The #[command(subcommand)] macro generates code that:
/// - defines a subcommand for the caesar command line application.
/// Subcommands are used to break up a large number of options into
/// smaller, more manageable groups.
///
/// The #[command] macro is used to annotate the `command` field of the struct.
/// The `command` field is an Option<Commands> type, which means that the
/// subcommand is optional.
/// If the subcommand is not provided, the program will exit with a status code
/// of 0 and without printing any output.
///
/// The Commands enum defines the possible subcommands for the caesar command
/// line application.
/// See the Commands enum definition for more information about the available
/// subcommands.
#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
#[command(version = env!("CARGO_PKG_VERSION"), about = "Send and receive files securely")]
#[command(long_about = None)]
pub struct Args {
/// The subcommand for the caesar command line application.
/// Subcommands are used to break up a large number of options into smaller,
/// more manageable groups.
/// If no subcommand is provided, the program will exit with a status code
/// of 0 and without printing any output.
#[command(subcommand)]
pub command: Option<Commands>,
}
@ -21,8 +51,8 @@ pub enum Commands {
#[arg(short, long)]
relay: Option<String>,
/// Path to file(s)
#[arg(value_name = "FILE")]
file: Option<String>,
#[arg(value_name = "FILES")]
files: Vec<String>,
},
/// Receives Files from the sender with the matching password
Receive {
@ -36,7 +66,7 @@ pub enum Commands {
/// Name of Transfer to download files
#[arg(value_name = "Transfer_Name")]
name: Option<String>,
name: String,
},
/// Start a relay server
Serve {
@ -47,103 +77,118 @@ pub enum Commands {
#[arg(short, long)]
listen_address: Option<String>,
},
/// Work with your configuration
Config {
/// Show path to config file
#[arg(short, long)]
path: bool,
/// View configured Options
#[arg(short, long)]
show: bool,
/// Edit the config file
#[arg(short, long)]
edit: bool,
/// Reset changed config
#[arg(short, long)]
reset: bool,
},
}
impl Default for Args {
// This function is called by the Default trait when no value is
// provided for a field of type Args. It returns an instance of
// Args that has been created by calling the new() function.
//
// The Default trait is used by various parts of the program to
// provide a sensible default value for a field when no value is
// provided. For example, the clap crate uses the Default trait when
// parsing command line arguments to provide a default value for
// a field.
//
// The new() function is a constructor function for Args that
// creates an instance of Args with default field values.
fn default() -> Self {
Self::new()
}
}
impl Args {
/// Creates a new instance of Args by parsing command line arguments
///
/// This function is a constructor for Args. It uses the clap crate to parse
/// command line arguments and creates an instance of Args with the values
/// provided by the user.
///
/// The clap crate is a command line argument parser that is well tested and
/// widely used. It provides a simple way to define command line
/// arguments and generate helpful documentation for the user.
///
/// The `parse()` function is used to parse the command line arguments and
/// return an instance of Args.
pub fn new() -> Self {
Self::parse()
}
/// Runs the command specified by the user
///
/// This function is called after the command line arguments have been
/// parsed. It matches on the `command` field of the Args struct to determine
/// what command the user wants to run.
///
/// The match statement checks the value of `command` and calls the
/// appropriate function to run the command. The functions that are called
/// are located in other modules of the program.
///
/// The `run()` function is called by the `main()` function of the program.
/// The program's entry point is the `main()` function, which parses the
/// command line arguments and then calls `run()` on the resulting Args
/// instance.
///
/// The `run()` function returns a Result. The error type is `Box<dyn
/// std::error::Error + Send + Sync>`. This means that the error type is a
/// trait object that represents an error that can be sent across threads
/// and sent over a network connection. The `Send` and `Sync` traits are part
/// of the standard library and are used to indicate that the error type can
/// be sent across threads and sent over a network connection.
///
/// The `run()` function does not return anything if the command is `None`.
/// This is because `command` is an `Option<Commands>`. If the user does
/// not specify a command, then `command` is `None`. In this case, there is
/// nothing to run, so `run()` returns early with no error.
pub async fn run(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
debug!("args: {:#?}", self);
match &self.command {
Some(Commands::Send { relay, file }) => {
let _ = match sender::send_info(
relay.as_deref().unwrap_or("http://0.0.0.0:8000"),
file.as_deref().unwrap_or("test.txt"),
)
.await
{
Ok(name) => {
println!("Transfer name: {}", name);
serf_file(file.as_ref().unwrap()).await;
Ok(())
}
Err(err) => Err(err),
};
}
Some(Commands::Receive {
relay,
overwrite,
name,
}) => {
let response = receiver::download_info(
relay.as_deref().unwrap_or("http://0.0.0.0:8000"),
name.as_deref().unwrap_or("None"),
// If the user wants to send files, call `start_sender()` in the
// `sender` module with the list of files that the user wants to
// send.
Some(Commands::Send { relay, files }) => {
sender::start_sender(
relay.as_deref().unwrap_or(
env::var("APP_ORIGIN")
.unwrap_or("ws://0.0.0.0:8000/ws".to_string())
.as_str(),
),
files,
)
.await;
match response {
Ok(res) => {
debug!("The response is: {:#?}", res);
let reachable = receiver::ping_sender(&res.body.ip).await;
match reachable {
Ok(_) => match receiver::download_file(&res, overwrite).await {
Ok(_) => {
info!("Download complete");
receiver::signal_success_relay(
relay.as_deref().unwrap_or("http://0.0.0.0:8000"),
name.as_deref().unwrap_or("None"),
)
.await?;
let _ =
match receiver::signal_success_sender(&res.body.ip).await {
Ok(_) => Ok(()),
Err(err) => Err(err),
};
}
Err(err) => error!(err),
},
Err(err) => error!("Error: {:#?}", err),
}
}
Err(err) => error!("Error: {:#?}", err),
}
}
// If the user wants to receive files, call `start_receiver()` in the
// `receiver` module with the name of the transfer that the user
// wants to download.
Some(Commands::Receive {
relay,
overwrite: _,
name,
}) => {
println!("Receive for {name:?}");
receiver::start_receiver(
relay.as_deref().unwrap_or(
env::var("APP_ORIGIN")
.unwrap_or("ws://0.0.0.0:8000/ws".to_string())
.as_str(),
),
name,
)
.await;
}
// If the user wants to start a relay server, call `start_ws()` in the
// `relay` module with the port and listen address that the user
// specified.
Some(Commands::Serve {
port,
listen_address,
}) => {
server::start_server(port.as_ref(), listen_address.as_ref()).await;
println!("Serve with address '{listen_address:?}' and '{port:?}'");
relay::start_ws(port.as_ref(), listen_address.as_ref()).await;
}
Some(Commands::Config {
path: _,
show: _,
edit: _,
reset: _,
}) => {}
// If the user does not specify a command, return early with no error.
// This is because `command` is an `Option<Commands>`. If the user does
// not specify a command, then `command` is `None`.
None => {}
}
Ok(())

0
src/crypto/mod.rs
View file

52
src/error/mod.rs
View file

@ -1,52 +0,0 @@
use std::error::Error;
use std::fmt;
#[derive(Debug)]
pub struct TransferNotCreatedError {
pub message: String,
}
impl TransferNotCreatedError {
pub fn new(msg: &str) -> Self {
Self {
message: msg.to_string(),
}
}
}
impl fmt::Display for TransferNotCreatedError {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "{}", self.message)
}
}
impl Error for TransferNotCreatedError {
fn description(&self) -> &str {
&self.message
}
}
#[derive(Debug)]
pub struct TransferNotFoundError {
message: String,
}
impl TransferNotFoundError {
pub fn new(msg: &str) -> Self {
Self {
message: msg.to_string(),
}
}
}
impl fmt::Display for TransferNotFoundError {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "{}", self.message)
}
}
impl Error for TransferNotFoundError {
fn description(&self) -> &str {
&self.message
}
}

53
src/http_client.rs
View file

@ -1,53 +0,0 @@
use http_body_util::{BodyExt, Full};
use hyper::{body::Bytes, Request, StatusCode};
use hyper_util::rt::TokioIo;
use tokio::{
io::{self, AsyncWriteExt},
net::TcpStream,
};
pub async fn send_request(
url: &str,
method: &str,
body: Option<String>,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
let url = url.parse::<hyper::Uri>()?;
let host = url.host().expect("uri has no host");
let port = url.port_u16().unwrap_or(80);
let address = format!("{}:{}", host, port);
let stream = TcpStream::connect(address).await?;
let io = TokioIo::new(stream);
let (mut sender, conn) = hyper::client::conn::http1::handshake(io).await?;
tokio::task::spawn(async move {
if let Err(err) = conn.await {
eprintln!("Connection failed: {:?}", err);
}
});
let authority = url.authority().unwrap().clone();
let send_body = match body {
Some(body_str) => Full::<Bytes>::from(Bytes::from(body_str)),
None => Full::<Bytes>::from(Bytes::from("")),
};
let req = Request::builder()
.method(method)
.uri(url)
.header(hyper::header::HOST, authority.as_str())
.header(hyper::header::CONTENT_TYPE, "application/json")
.body(send_body)?;
let mut res = sender.send_request(req).await?;
println!("Response status: {}", res.status());
if res.status() == StatusCode::OK {
while let Some(next) = res.frame().await {
let frame = next?;
if let Some(chunk) = frame.data_ref() {
io::stdout().write_all(chunk).await?;
}
}
}
Ok(())
}

28
src/main.rs
View file

@ -1,25 +1,37 @@
mod cli;
mod error;
mod http_client;
mod receiver;
mod relay;
mod sender;
mod transfer_info;
use crate::cli::args::Args;
use dotenv::dotenv;
use tracing::error;
use tracing_subscriber::filter::EnvFilter;
pub mod cli;
pub mod receiver;
pub mod relay;
pub mod sender;
pub mod shared;
#[tokio::main]
// This is the entrypoint of caesar.
// The #[tokio::main] attribute is required for any async code, and it
// sets up the tokio runtime.
// The async fn main() is the entrypoint of the application, and it's where
// we kick off our program.
async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
// Load environment variables from a .env file if one is present.
dotenv().ok();
// Set up our logging subscriber.
// TheEnvFilter::from_default_env reads the env variable RUST_LOG
// and sets up the logging accordingly.
// The default is INFO level logging.
tracing_subscriber::fmt()
.with_env_filter(EnvFilter::from_default_env())
.init();
// env_logger::init();
// Parse the command line arguments.
let args = Args::new();
// Run the commands based on the parsed arguments.
// If there is an error, print it to the console with the error! macro.
if let Err(e) = args.run().await {
error!("{e}");
}
// Return an Ok result, which just means that our program exited successfully.
Ok(())
}

675
src/receiver/client.rs
View file

@ -1,97 +1,608 @@
use crate::error::TransferNotFoundError;
use crate::transfer_info::transfer_info::TransferInfoRequest;
use hex;
use reqwest::Client;
use sha2::{Digest, Sha256};
use std::fs::File;
use std::io::copy;
use tracing::{debug, error, info};
use std::{fs, io::stdout, path::Path};
type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;
use crate::shared::{
packets::{
packet::Value, ChunkPacket, HandshakePacket, HandshakeResponsePacket, ListPacket, Packet,
ProgressPacket,
},
JsonPacket, JsonPacketResponse, JsonPacketSender, PacketSender, Sender, Socket, Status,
};
pub async fn download_info(relay: &str, name: &str) -> Result<TransferInfoRequest> {
let hashed_name = Sha256::digest(name.as_bytes());
let hashed_string = hex::encode(hashed_name);
use aes_gcm::{aead::Aead, Aes128Gcm, Key};
use base64::{engine::general_purpose, Engine as _};
use futures_util::{future, pin_mut, stream::TryStreamExt, StreamExt};
use hmac::{Hmac, Mac};
use p256::{ecdh::EphemeralSecret, pkcs8::der::Writer, PublicKey};
use prost::Message;
use rand::rngs::OsRng;
use sha2::Sha256;
use tokio_tungstenite::tungstenite::{protocol::Message as WebSocketMessage, Error};
use tracing::error;
match reqwest::get(format!("{}/download/{}", relay, hashed_string)).await {
Ok(resp) => {
let json = resp.json::<TransferInfoRequest>().await?;
debug!("Json Response: {:#?}", json);
if json.message == *"error" {
Err(Box::new(TransferNotFoundError::new(
"no transfer with given name found",
)))
} else {
debug!("Got Positive response");
Ok(json)
const DESTINATION: u8 = 0;
const NONCE_SIZE: usize = 12;
/// This struct represents a file that is being received.
///
/// The struct contains information about the file, such as its name, size,
/// and the handle of the file on the disk.
///
/// The `name` field contains the name of the file, which is the name of the
/// file on the disk.
///
/// The `size` field contains the size of the file in bytes.
///
/// The `progress` field contains the number of bytes that have already been
/// received for the file.
///
/// The `handle` field contains a handle to the file on the disk, which is
/// used to read the contents of the file.
struct File {
name: String,
size: u64,
progress: u64,
handle: fs::File,
}
/// This struct contains the context for the receiver.
///
/// This structure is used to keep track of the state of the receiver, and to
/// pass information between functions.
struct Context {
/// The HMAC key that is used to verify that the packets that are received
/// are authentic.
///
/// The HMAC key is generated by the sender, and is used to verify that the
/// packets that are received are authentic. If the HMAC of a packet does
/// not match the expected HMAC, then the packet is not processed.
hmac: Vec<u8>,
/// The sender that is used to send packets to the server.
///
/// The sender is used to send packets to the server. The sender is also
/// used to receive packets from the server.
sender: Sender,
/// The ephemeral secret key that is used for key exchange with the sender.
///
/// The ephemeral secret key is generated by the receiver, and is used to
/// exchange a shared key with the sender. The shared key is used to
/// encrypt and decrypt packets.
key: EphemeralSecret,
/// The shared key that is used to encrypt and decrypt packets.
///
/// The shared key is established between the receiver and the sender during
/// the key exchange. The shared key is used to encrypt and decrypt packets
/// between the receiver and the sender. If the shared key is `None`, then the
/// packets that are received are not processed.
shared_key: Option<Aes128Gcm>,
/// The files that are being received.
///
/// The files vector contains a list of all the files that are being
/// received. Each file is represented by a `File` struct. The `name` field
/// of the `File` struct contains the name of the file, which is the name of
/// the file on the disk. The `size` field of the `File` struct contains the
/// size of the file in bytes. The `progress` field of the `File` struct
/// contains the number of bytes that have already been received for the
/// file. The `handle` field of the `File` struct contains a handle to the
/// file on the disk, which is used to read the contents of the file.
files: Vec<File>,
/// The sequence number of the next chunk that is expected to be received.
///
/// The sequence number is used to keep track of the sequence of chunks that
/// are received. If a chunk does not have the expected sequence number,
/// then the chunk is not processed.
sequence: u32,
/// The index of the file that is currently being received.
///
/// The index is used to keep track of which file is currently being
/// received. The index is incremented after a file is completely received.
index: usize,
/// The progress of the current file that is being received.
///
/// The progress is used to keep track of the progress of the current file
/// that is being received. The progress is calculated by dividing the
/// number of bytes that have been received by the size of the file. The
/// progress is sent to the server so that the sender knows how much of the
/// file has been received.
progress: u64,
/// The total length of the current file that is being received.
///
/// The length is used to keep track of the total length of the current file
/// that is being received. The length is used to calculate the progress of
/// the file.
length: u64,
}
/// This function is called when the receiver receives a join room packet from
/// the server. The packet contains the size of the list of files that the
/// sender is going to send.
///
/// If the packet does not contain the size of the list, then an error is
/// returned.
///
/// If the packet does contain the size of the list, then a message is printed
/// to the console indicating that the receiver has connected to the room.
///
/// The function does not do anything else. It returns a `Status::Continue`
/// variant to indicate that the event loop should continue processing events.
fn on_join_room(size: Option<usize>) -> Status {
// If the packet does not contain the size of the list, then return an error.
if size.is_none() {
return Status::Err("Invalid join room packet.".into());
}
// If the packet contains the size of the list, then print a message to the
// console indicating that the receiver has connected to the room.
println!("Connected to room.");
// Return a `Status::Continue` variant to indicate that the event loop
// should continue processing events.
Status::Continue()
}
/// This function is called when the event loop receives an error packet from
/// the server. The packet contains a message with a description of the error.
///
/// When an error occurs, the server sends an error packet to the client. The
/// error packet contains a message with a description of the error. This
/// function extracts that message and creates a `Status::Err` variant with it,
/// which is then returned to be handled by the main event loop.
///
/// When the event loop receives a status variant that is an error, it exits
/// with an error message containing the message from the error packet.
///
/// The message from the error packet is the only information that the event
/// loop has about the error, so the message should be descriptive and
/// helpful to the user. The message should not contain technical details
/// about the error or how it occurred. Instead, the message should be
/// written from the perspective of the user and should give the user enough
/// information to understand what went wrong and how they might be able to
/// fix the problem.
///
/// This function takes the message from the error packet and creates a
/// `Status::Err` variant with it. The function returns this variant to be
/// handled by the main event loop.
fn on_error(message: String) -> Status {
Status::Err(message)
}
/// This function is called when the event loop receives a leave room packet from
/// the server. The packet contains the index of the file that was being
/// transferred when the receiver left the room.
///
/// When the receiver receives a leave room packet, it means that the sender
/// has disconnected from the room. In this case, the receiver should check if
/// there are any files that were being transferred but not yet complete. If
/// there are any incomplete files, the receiver should print a message to
/// the user indicating that the transfer was interrupted.
///
/// If there are no incomplete files, then the receiver should exit
/// normally. The `Status::Exit` variant is returned to the main event loop,
/// which will cause the event loop to exit normally.
///
/// This function checks if there are any incomplete files by iterating over
/// the list of files in the context. If there are any files with a progress
/// less than 100%, then the function prints a message to the user and returns
/// an error status.
///
/// If there are no incomplete files, then the function simply returns a
/// `Status::Exit` variant. This will cause the main event loop to exit
/// normally.
fn on_leave_room(context: &mut Context, _: usize) -> Status {
// Check if there are any incomplete files.
if context.files.iter().any(|file| file.progress < 100) {
// If there are any incomplete files, print a message to the user.
println!();
println!("Transfer was interrupted because the host left the room.");
// Return an error status.
Status::Err("Transfer was interrupted because the host left the room.".into())
} else {
// If there are no incomplete files, return a `Status::Exit` variant.
// This will cause the event loop to exit normally.
Status::Exit()
}
}
/// This function is called when the event loop receives a list packet from
/// the server. The packet contains a list of files to be transferred.
///
/// When this function is called, we know that the sender has successfully
/// established a shared key with the receiver. Therefore, we can start
/// receiving encrypted files.
///
/// This function iterates over the list of files in the packet and creates a
/// file on disk for each file in the list. If a file with the same name already
/// exists, an error is returned and the event loop is exited with an error
/// message. This is because the receiver should not overwrite existing files
/// without the user's explicit permission.
///
/// Once all the files have been created, the function initializes the context
/// variables for the event loop. `index` is set to 0 to indicate that we are
/// currently transferring the first file. `progress` is set to 0 to indicate
/// that the progress of the first file is 0%. `sequence` is set to 0 to
/// indicate that the first chunk of data we receive will have a sequence
/// number of 0. `length` is set to 0 to indicate that we have not received
/// any data yet.
///
/// If there is an error creating any of the files, the function returns an
/// error status. This will cause the event loop to exit with an error message.
///
/// If there are no errors, the function returns a `Status::Continue()` variant.
/// This will cause the event loop to continue running and wait for more
/// packets from the sender.
fn on_list(context: &mut Context, list: ListPacket) -> Status {
if context.shared_key.is_none() {
return Status::Err("Invalid list packet: no shared key established".into());
}
// Iterate over the list of files in the packet.
for entry in list.entries {
// Sanitize the file name to remove any characters that are not valid in
// file names on the current platform.
let path = sanitize_filename::sanitize(entry.name.clone());
// Check if a file with the same name already exists.
if Path::new(&path).exists() {
// If the file already exists, return an error and exit the event loop
// with an error message.
return Status::Err(format!("The file '{}' already exists.", path));
}
// Try to create a new file with the sanitized file name.
let handle = match fs::File::create(&path) {
Ok(handle) => handle,
Err(error) => {
// If there is an error creating the file, return an error and
// exit the event loop with an error message.
return Status::Err(format!(
"Error: Failed to create file '{}': {}",
path, error
));
}
}
Err(err) => {
error!("Error: {err}");
Err(Box::new(err))
}
}
}
};
pub async fn download_file(transfer_info: &TransferInfoRequest, overwrite: &bool) -> Result<()> {
if !*overwrite && File::open(&transfer_info.body.files).is_ok() {
return Err(Box::new(std::io::Error::new(
std::io::ErrorKind::AlreadyExists,
format!("File '{} already exists", &transfer_info.body.files),
)));
// Create a new file struct for the file we just created.
let file = File {
name: entry.name,
size: entry.size,
handle,
progress: 0,
};
// Add the new file to the list of files in the context.
context.files.push(file);
}
let resp = reqwest::get(format!(
"http://{}:8100/download_file",
&transfer_info.body.ip
))
.await?;
if !resp.status().is_success() {
return Err(Box::new(std::io::Error::new(
std::io::ErrorKind::Other,
format!("Failed to download file from {}", &transfer_info.body.ip),
)));
// Set the context variables for the event loop.
context.index = 0;
context.progress = 0;
context.sequence = 0;
context.length = 0;
// Return a `Status::Continue()` variant to indicate that the event loop
// should continue running and wait for more packets from the sender.
Status::Continue()
}
/// This function handles a chunk packet received from the sender.
///
/// It checks that the shared key has been established, that the sequence number
/// of the chunk matches the expected sequence number in the context, and that
/// the index of the file in the context is valid.
///
/// If any of these checks fail, an error is returned and the event loop is
/// stopped.
///
/// The function updates the length of the file, increments the sequence number
/// in the context, and writes the contents of the chunk to the file.
///
/// The progress of the file is updated to be the ratio of the number of bytes
/// read so far to the total size of the file.
///
/// If the progress of the file is 100%, or if the difference in progress between
/// this chunk and the last chunk is greater than or equal to 1, or if this is the
/// first chunk, a ProgressPacket is sent to the sender with the index of the file
/// in the context and the progress of the file.
///
/// If the size of the file has been reached, the index of the current file is
/// incremented, the length of the current file is set to 0, the progress of the
/// current file is set to 0, and the sequence number is set to 0.
///
/// Finally, a Status::Continue() variant is returned to indicate that the event
/// loop should continue running and wait for more packets from the sender.
fn on_chunk(context: &mut Context, chunk: ChunkPacket) -> Status {
// Check that the shared key has been established.
if context.shared_key.is_none() {
return Status::Err("Invalid chunk packet: no shared key established".into());
}
let mut dest = File::create(&transfer_info.body.files)?;
let content = resp.text().await?;
copy(&mut content.as_bytes(), &mut dest)?;
info!("Download complete");
Ok(())
}
pub async fn ping_sender(sender: &String) -> Result<bool> {
debug!("Pinging Sender on {:#?}", sender);
match reqwest::get(format!("http://{}:8100/ping", sender)).await {
Ok(resp) => {
debug!("Sender directly reachable");
debug!("Response is: {:#?}", resp);
Ok(true)
}
Err(err) => {
error!("Error: {err}");
Err(Box::new(err))
}
// Check that the sequence number of the chunk matches the expected sequence
// number in the context.
if chunk.sequence != context.sequence {
return Status::Err(format!(
"Expected sequence {}, but got {}.",
context.sequence, chunk.sequence
));
}
// Get a mutable reference to the file in the context at the index of the
// file.
let Some(file) = context.files.get_mut(context.index) else {
// If the index of the file in the context is invalid, return an error and
// stop the event loop.
return Status::Err("Invalid file index.".into());
};
// Update the length of the file.
context.length += chunk.chunk.len() as u64;
// Increment the sequence number in the context.
context.sequence += 1;
// Write the contents of the chunk to the file.
file.handle.write(&chunk.chunk).unwrap();
// Update the progress of the file.
file.progress = (context.length * 100) / file.size;
// If the progress of the file is 100%, or if the difference in progress between
// this chunk and the last chunk is greater than or equal to 1, or if this is the
// first chunk, send a ProgressPacket to the sender.
if file.progress == 100 || file.progress - context.progress >= 1 || chunk.sequence == 0 {
context.progress = file.progress;
let progress = ProgressPacket {
// Convert the index of the file in the context to a u32.
index: context.index.try_into().unwrap(),
// Convert the progress of the file to a u32.
progress: context.progress.try_into().unwrap(),
};
// Send the ProgressPacket to the sender.
context.sender.send_encrypted_packet(
&context.shared_key,
DESTINATION,
Value::Progress(progress),
);
print!("\rTransferring '{}': {}%", file.name, file.progress);
std::io::Write::flush(&mut stdout()).unwrap();
}
// If the size of the file has been reached, increment the index of the
// current file, set the length of the current file to 0, set the progress
// of the current file to 0, and resets the sequence number to 0.
if file.size == context.length {
context.index += 1;
context.length = 0;
context.progress = 0;
context.sequence = 0;
println!();
}
// Return a Status::Continue() variant to indicate that the event loop should
// continue running and wait for more packets from the sender.
Status::Continue()
}
pub async fn signal_success_relay(relay: &str, name: &str) -> Result<()> {
let hashed_name = Sha256::digest(name.as_bytes());
let hashed_string = hex::encode(hashed_name);
debug!("Signaling success to {:#?}", relay);
let client = Client::new();
let _ = client
.post(format!("{}/download_success/{}", relay, hashed_string))
.send()
.await?;
Ok(())
/// This function is called when the Receiver receives a HandshakePacket from the
/// Sender. It verifies the signature of the Sender's public key and generates its own
/// public key. It then generates a shared secret key between the Receiver and the Sender
/// using the Diffie-Hellman key exchange.
///
/// The Receiver sends back a HandshakeResponsePacket to the Sender with its own public
/// key and a signature created using the shared secret key and its own private key.
///
/// The shared secret key is used to encrypt packets sent between the Receiver and the
/// Sender.
fn on_handshake(context: &mut Context, handshake: HandshakePacket) -> Status {
// If a shared key has already been established, this means that the Receiver
// has already performed the handshake, so return an error.
if context.shared_key.is_some() {
return Status::Err("Already performed handshake.".into());
}
// Create a new HMAC using the hmac from the Context struct as the key.
let mut mac = Hmac::<Sha256>::new_from_slice(&context.hmac).unwrap();
// Update the HMAC with the public key from the HandshakePacket.
mac.update(&handshake.public_key);
// Call verify_slice() on the HMAC to verify the signature from the Sender.
// If the signature is invalid, return an error.
let verification = mac.verify_slice(&handshake.signature);
if verification.is_err() {
return Status::Err("Invalid signature from the sender.".into());
}
// Generate the Receiver's public key from the private key.
let public_key = context.key.public_key().to_sec1_bytes().into_vec();
// Create a new HMAC using the hmac from the Context struct as the key.
let mut mac = Hmac::<Sha256>::new_from_slice(&context.hmac).unwrap();
// Update the HMAC with the public key of the Receiver.
mac.update(&public_key);
// Serialize the resulting HMAC into a byte array and use it as the
// signature in the HandshakeResponsePacket.
let signature = mac.finalize().into_bytes().to_vec();
// Create a new shared secret key between the Receiver and the Sender.
let shared_public_key = PublicKey::from_sec1_bytes(&handshake.public_key).unwrap();
let shared_secret = context.key.diffie_hellman(&shared_public_key);
let shared_secret = shared_secret.raw_secret_bytes();
let shared_secret = &shared_secret[0..16];
// Create a new Aes128Gcm key from the shared secret.
let shared_key: &Key<Aes128Gcm> = shared_secret.into();
let shared_key = <Aes128Gcm as aes_gcm::KeyInit>::new(shared_key);
// Create the HandshakeResponsePacket and send it to the Sender.
let handshake_response = HandshakeResponsePacket {
public_key,
signature,
};
context
.sender
.send_packet(DESTINATION, Value::HandshakeResponse(handshake_response));
// Store the shared key in the Context struct.
context.shared_key = Some(shared_key);
// Return a Status::Continue() variant to indicate that the event loop should
// continue running and wait for more packets from the Sender.
Status::Continue()
}
pub async fn signal_success_sender(sender: &String) -> Result<()> {
debug!("Signaling shutdown to {:#?}", sender);
let client = Client::new();
let _ = client
.post(format!("http://{}:8100/shutdown", sender))
.send()
.await?;
Ok(())
/// This function is called when a message is received from the Sender.
///
/// The message can be either text or binary. If it's text, we attempt to
/// parse it as a JsonPacketResponse and match on the type of response it is.
/// If it's binary, we attempt to decrypt it using the shared key (if it
/// exists) and then decode it into a Packet. We then match on the type of
/// value in the Packet and call the appropriate function with the relevant
/// data.
///
/// If the message is not text or binary, we return a Status::Err with an
/// appropriate error message.
fn on_message(context: &mut Context, message: WebSocketMessage) -> Status {
if message.is_text() {
let text = message.into_text().unwrap();
let packet = serde_json::from_str(&text).unwrap();
return match packet {
JsonPacketResponse::Join { size } => on_join_room(size),
JsonPacketResponse::Leave { index } => on_leave_room(context, index),
JsonPacketResponse::Error { message } => on_error(message),
_ => Status::Err(format!("Unexpected json packet: {:?}", packet)),
};
} else if message.is_binary() {
let data = message.into_data();
let data = &data[1..];
let data = if let Some(shared_key) = &context.shared_key {
let nonce = &data[..NONCE_SIZE];
let ciphertext = &data[NONCE_SIZE..];
shared_key.decrypt(nonce.into(), ciphertext).unwrap()
} else {
data.to_vec()
};
let packet = Packet::decode(data.as_ref()).unwrap();
let value = packet.value.unwrap();
return match value {
Value::List(list) => on_list(context, list),
Value::Chunk(chunk) => on_chunk(context, chunk),
Value::Handshake(handshake) => on_handshake(context, handshake),
_ => Status::Err(format!("Unexpected packet: {:?}", value)),
};
}
Status::Err("Invalid message type".into())
}
/// This function takes a websocket connection and an invite code,
/// splits the connection into an outgoing and incoming part,
/// creates a context for the connection, sends a join room packet,
/// and starts two futures to handle incoming and outgoing messages.
///
/// The outgoing future reads from a channel and sends the messages
/// through the outgoing part of the connection. If the sending fails,
/// the future will print an error and exit.
///
/// The incoming future reads from the incoming part of the connection
/// and passes the messages to the `on_message` function. If the message
/// is an exit or an error, the function will print the error and exit.
/// If the message is any other type of packet, it will be handled by the
/// `on_message` function and the future will continue running.
pub async fn start(socket: Socket, fragment: &str) {
// Extract the room id and hmac from the invite code
let Some(index) = fragment.rfind('-') else {
println!("Error: The invite code '{}' is not valid.", fragment);
return;
};
let id = &fragment[..index];
let hmac = &fragment[index + 1..];
let Ok(hmac) = general_purpose::STANDARD.decode(hmac) else {
error!("Error: Invalid base64 inside the invite code.");
return;
};
// Create a new ephemeral key pair
let key = EphemeralSecret::random(&mut OsRng);
// Create a channel for sending messages
let (sender, receiver) = flume::bounded(1000);
// Split the websocket connection into an outgoing and incoming part
let (outgoing, incoming) = socket.split();
// Create a new context for the connection
let mut context = Context {
hmac,
sender,
key,
shared_key: None,
files: vec![],
index: 0,
sequence: 0,
progress: 0,
length: 0,
};
println!("Attempting to join room '{}'...", id);
// Send a join room packet to the server
context
.sender
.send_json_packet(JsonPacket::Join { id: id.to_string() });
// Create futures for handling incoming and outgoing messages
let outgoing_handler = receiver.stream().map(Ok).forward(outgoing);
let incoming_handler = incoming.try_for_each(|message| {
// Call the on_message function to handle the message
match on_message(&mut context, message) {
// If the message is an exit, print a message and exit
Status::Exit() => {
println!("Transfer has completed.");
return future::err(Error::ConnectionClosed);
}
// If the message is an error, print the error and exit
Status::Err(error) => {
println!("Error: {}", error);
return future::err(Error::ConnectionClosed);
}
// If the message is any other type of packet, do nothing
_ => {}
};
// Continue running the future
future::ok(())
});
// Pin the futures to the stack so they can run concurrently
pin_mut!(incoming_handler, outgoing_handler);
// Wait for either future to complete
future::select(incoming_handler, outgoing_handler).await;
}

81
src/receiver/mod.rs
View file

@ -1 +1,82 @@
/// This module is the entry point for the receiver command.
/// It contains a single function, `start_receiver`, which is the
/// entry point for the receiver program.
///
/// The `start_receiver` function takes a `String` which is the URL or
/// invite code for the room that the receiver should join. If the
/// URL is invalid or does not contain an invite code fragment,
/// the function falls back to using the command line arguments to get
/// the file paths to be sent.
///
/// The `start_receiver` function first creates a request to connect
/// to the WebSocket server with a specific origin. This is done to
/// prevent cross-origin requests, which are not allowed by the
/// WebSocket protocol.
///
/// If creating the request succeeds, the function inserts the origin
/// into the request headers. Then, it attempts to connect to the
/// server using the `connect_async` function from the
/// `tokio_tungstenite` crate.
///
/// If the connection attempt succeeds, the function extracts the
/// invite code fragment from the URL and passes it to the `start`
/// function in the `receiver::client` module. The `start` function is
/// defined in the `receiver::client` module and is the function that
/// interacts with the server to receive files.
///
/// If the connection attempt fails or the URL does not contain an
/// invite code fragment, the function falls back to using the command
/// line arguments to get the file paths to be sent. It then calls the
/// `start` function in the `sender::client` module with the
/// WebSocket stream and the file paths. The `start` function in the
/// `sender::client` module is defined in the `sender::client`
/// module and is the function that sends the files over the
/// WebSocket connection.
///
/// The `start` function takes ownership of the WebSocket stream and
/// the file paths, so we pass them by value.
pub mod client;
use crate::receiver::client as receiver;
use tokio_tungstenite::{
connect_async,
tungstenite::{client::IntoClientRequest, http::HeaderValue},
};
use tracing::error;
use url::Url;
pub async fn start_receiver(relay: &str, name: &str) {
let argument = name;
let Ok(mut request) = relay.into_client_request() else {
println!("Error: Failed to create request.");
return;
};
// Insert the origin into the request headers to prevent
// cross-origin requests.
request
.headers_mut()
.insert("Origin", HeaderValue::from_str(relay).unwrap());
println!("Attempting to connect...");
let Ok((socket, _)) = connect_async(request).await else {
error!("Error: Failed to connect.");
return;
};
// If the URL is valid and contains an invite code fragment,
// extract it and pass it to the receiver::client::start
// function. The start function is defined in the
// receiver::client module and is the function that interacts with
// the server to receive files.
if let Ok(url) = Url::parse(argument) {
let Some(fragment) = url.fragment() else {
error!("Error: Missing invite code fragment in url.");
return;
};
receiver::start(socket, fragment).await
}
}

287
src/relay/mod.rs
View file

@ -1 +1,288 @@
/// This function starts the WebSocket server.
///
/// It configures the server to listen on the specified host and port. If
/// these values are not specified in the environment, it falls back to using
/// the defaults of "0.0.0.0" for the host and "8000" for the port.
///
/// It then sets up the application routes for the server. In this case, the
/// only route is for the WebSocket connection.
///
/// The WebSocket route requires a `ConnectInfo` extractor to get the client's
/// IP address, which is then used to store the client in a data structure
/// keyed by their IP address. This allows for efficient lookup of clients by
/// their IP address.
///
/// Finally, it starts the server by binding to the specified host and port,
/// and running the application. If the server fails to bind to the specified
/// host and port, it logs an error and exits.
pub mod server;
use axum::{
extract::{ws::WebSocket, ConnectInfo, State, WebSocketUpgrade},
response::IntoResponse,
routing::get,
Router,
};
use axum_client_ip::SecureClientIpSource;
use futures_util::StreamExt;
use std::{env, net::SocketAddr, sync::Arc};
use tokio::{
net::TcpListener,
signal,
sync::{Mutex, RwLock},
};
use tower_http::trace::{DefaultMakeSpan, TraceLayer};
use tracing::{debug, error, info};
use self::server::Client;
/// This function starts the WebSocket server.
///
/// It retrieves the environment variables that define how the server should
/// be configured. If any of these variables are not defined, it sets a
/// reasonable default value.
///
/// The environment variables are:
///
/// * `APP_ENVIRONMENT`: the environment the server is running in (defaults
/// to "development").
/// * `APP_HOST`: the host the server should listen on (defaults to "0.0.0.0").
/// * `APP_PORT`: the port the server should listen on (defaults to "8000").
/// * `APP_DOMAIN`: the domain the server is accessible at (defaults to "").
///
/// It then sets up the application routes for the server. In this case, the
/// only route is for the WebSocket connection.
///
/// The WebSocket route requires a `ConnectInfo` extractor to get the client's
/// IP address, which is then used to store the client in a data structure
/// keyed by their IP address. This allows for efficient lookup of clients by
/// their IP address.
///
/// Finally, it starts the server by binding to the specified host and port,
/// and running the application. If the server fails to bind to the specified
/// host and port, it logs an error and exits.
pub async fn start_ws(port: Option<&i32>, listen_addr: Option<&String>) {
// Retrieve environment variables and set defaults if necessary.
let app_environemt = env::var("APP_ENVIRONMENT").unwrap_or("development".to_string());
let app_host = match listen_addr {
Some(address) => address.to_string(),
None => env::var("APP_HOST").unwrap_or("0.0.0.0".to_string()),
};
let app_port = match port {
Some(port) => port.to_string(),
None => env::var("APP_PORT").unwrap_or("8000".to_string()),
};
let app_domain = env::var("APP_DOMAIN").unwrap_or("".to_string());
// Log information about the server's configuration.
debug!(
"Server configured to accept connections on host {app_host}...",
);
debug!(
"Server configured to listen connections on port {app_port}...",
);
debug!(
"Server configured to listen connections on port {app_domain}...",
);
// Based on the environment variable, set the logging level.
match app_environemt.as_str() {
"development" => {
debug!("Running in development mode");
}
"production" => {
debug!("Running in production mode");
}
_ => {
debug!("Running in development mode");
}
}
// Create a new server data structure.
let server = server::Server::new();
// Set up the application routes.
let app = Router::new()
.route("/ws", get(ws_handler))
.with_state(server)
.layer(SecureClientIpSource::ConnectInfo.into_extension())
.layer(
TraceLayer::new_for_http()
.make_span_with(DefaultMakeSpan::default().include_headers(true)),
);
// Attempt to bind to the specified host and port.
if let Ok(listener) = TcpListener::bind(&format!("{}:{}", app_host, app_port)).await {
// Log successful binding.
info!("Listening on: {}", listener.local_addr().unwrap());
// Run the server.
axum::serve(
listener,
app.into_make_service_with_connect_info::<SocketAddr>(),
)
.with_graceful_shutdown(shutdown_signal())
.await
.unwrap();
} else {
// Log binding failure and exit.
error!("Failed to listen on: {}:{}", app_host, app_port);
}
}
/// This function is an endpoint for the WebSocket route.
///
/// This function is called whenever a client makes a WebSocket request to
/// the `/ws` endpoint.
///
/// The function takes two arguments:
///
/// - `ws`: This is the WebSocketUpgrade object, which is used to upgrade the
/// HTTP connection to a WebSocket connection.
/// - `State(shared_state)`: This is the state of the server, which is stored
/// in a read-write lock. The state is shared between all WebSocket
/// connections.
/// - `ConnectInfo(addr)`: This is the information about the client that
/// connected to the server. The function uses this information to log the
/// address of the client that connected to the server.
///
/// The function upgrades the HTTP connection to a WebSocket connection using
/// the `ws` argument. It then passes the upgraded WebSocket connection, along
/// with the state of the server, to the `handle_socket` function.
///
/// The `handle_socket` function is defined in the `src/relay/mod.rs` file. It
/// is the function that handles the WebSocket connection.
///
/// The `handle_socket` function takes three arguments:
///
/// - `socket`: This is the WebSocket connection that it should handle.
/// - `who`: This is the address of the client that connected to the server.
/// - `rooms`: This is the state of the server, which is stored in a read-write
/// lock. The state is shared between all WebSocket connections.
///
/// The `handle_socket` function handles the WebSocket connection by calling
/// the `handle_message` function on a `Client` object that it creates. The
/// `handle_message` function is defined in the `src/relay/client.rs` file. The
/// `handle_message` function handles incoming messages from the client and
/// takes care of sending the appropriate response back to the client.
async fn ws_handler(
ws: WebSocketUpgrade,
State(shared_state): State<Arc<RwLock<server::Server>>>,
ConnectInfo(addr): ConnectInfo<SocketAddr>,
) -> impl IntoResponse {
debug!("Got Request on Websocket route");
debug!("WebSocket connection established from:{}", addr.to_string());
debug!("Upgrading Connection");
ws.on_upgrade(move |socket| handle_socket(socket, addr, shared_state))
}
/// This function is called when a new WebSocket connection is established.
/// The function takes three arguments:
///
/// - `socket`: This is the WebSocket connection that it should handle.
/// - `who`: This is the address of the client that connected to the server.
/// - `rooms`: This is the state of the server, which is stored in a read-write
/// lock. The state is shared between all WebSocket connections.
///
/// The function creates a `Client` object, which will handle the WebSocket
/// connection. The `Client` object is created with an Arc-wrapped Mutex
/// containing the `sender` of the WebSocket connection. The `sender` is used to
/// send messages to the client.
///
/// The function then creates a new `split` of the WebSocket connection, which
/// is a pair of a `sender` and a `receiver`. The `sender` is used to send
/// messages to the client, and the `receiver` is used to receive messages from
/// the client. The `receiver` is wrapped in a `Stream` (which is an async
/// iterator) so that the function can use the `next` method to receive messages
/// from the client.
///
/// The function then enters a loop that receives incoming messages from the
/// client and handles them. For each received message, the function calls the
/// `handle_message` method on the `Client` object that it created. The
/// `handle_message` method is defined in the `src/relay/client.rs` file. The
/// `handle_message` method handles incoming messages from the client and
/// takes care of sending the appropriate response back to the client.
///
/// If the function encounters an error while reading a message from the
/// client, it logs the error and breaks out of the loop.
///
/// After the loop finishes (either because an error occurred or because the
/// client disconnected), the function calls the `handle_close` method on the
/// `Client` object that it created. The `handle_close` method is defined in the
/// `src/relay/client.rs` file. The `handle_close` method handles the close event
/// from the client.
async fn handle_socket(socket: WebSocket, who: SocketAddr, rooms: Arc<RwLock<server::Server>>) {
let (sender, mut receiver) = socket.split();
let sender = Arc::new(Mutex::new(sender));
let mut client = Client::new(sender.clone());
while let Some(message) = receiver.next().await {
match message {
Ok(message) => {
client.handle_message(&rooms, message).await;
}
Err(error) => {
error!("Failed to read message from client {}: {}", who, error);
break;
}
}
}
// Handle the close event from the client.
client.handle_close(&rooms).await
}
/// This function sets up a signal handler for SIGINT (Ctrl+C) and SIGTERM
/// (terminate) on Unix platforms. It does nothing on non-Unix platforms.
///
/// The function installs two signal handlers: one for SIGINT and one for
/// SIGTERM. When either of these signals is received, the signal handler
/// simply resolves the future with `()`. This allows the main function to
/// wait for the signal handler to trigger a shutdown.
///
/// The function uses the `tokio::select!` macro to wait for either of the
/// signal handlers to resolve. When the future returned by `tokio::select!`
/// resolves, the function simply drops the value and does nothing else.
///
/// The function does not actually do anything itself. It simply waits for
/// one of the signal handlers to trigger a shutdown.
async fn shutdown_signal() {
// Install a signal handler for SIGINT (Ctrl+C). This future resolves
// when the user presses Ctrl+C.
let ctrl_c = async {
signal::ctrl_c()
.await
.expect("failed to install Ctrl+C handler");
};
// Install a signal handler for SIGTERM (terminate). This future
// resolves when the operating system sends a SIGTERM signal to the
// program.
#[cfg(unix)]
let terminate = async {
signal::unix::signal(signal::unix::SignalKind::terminate())
.expect("failed to install signal handler")
.recv()
.await;
};
// If we are not on a Unix platform, we don't need to install a signal
// handler for SIGTERM. Instead, we create a future that never resolves.
#[cfg(not(unix))]
let terminate = std::future::pending::<()>();
// Wait for either of the two signal handlers to resolve. When one of them
// resolves, the other one may still be waiting, but it doesn't matter
// because we don't need to do anything else.
tokio::select! {
// If the Ctrl+C signal handler resolves, drop the value and do
// nothing else.
_ = ctrl_c => {},
// If the terminate signal handler resolves, drop the value and do
// nothing else.
_ = terminate => {},
}
}

839
src/relay/server.rs
View file

@ -1,206 +1,663 @@
use crate::transfer_info::transfer_info::{TransferInfoBody, TransferInfoRequest};
use axum::{
extract::{connect_info::ConnectInfo, Json, Path, State},
http::StatusCode,
response::IntoResponse,
routing::{get, post},
Router,
};
use axum_client_ip::SecureClientIpSource;
use serde_json::json;
use std::{
env,
net::SocketAddr,
sync::{Arc, Mutex},
};
use tokio::signal;
use tracing::{debug, info, warn};
use axum::extract::ws::Message;
use futures_util::{future::join_all, stream::SplitSink, SinkExt};
use serde::{Deserialize, Serialize};
use std::{collections::HashMap, sync::Arc, vec};
use tokio::{sync::Mutex, sync::RwLock};
use tracing::error;
#[derive(Debug, Clone)]
struct AppState {
data: Arc<Mutex<Vec<TransferInfoRequest>>>,
}
use uuid::Uuid;
pub async fn start_server(port: Option<&i32>, listen_addr: Option<&String>) {
info!("Server starting...");
let shared_state = AppState {
data: Arc::new(Mutex::new(Vec::new())),
};
let app_environemt = env::var("APP_ENVIRONMENT").unwrap_or("development".to_string());
let app_host = match listen_addr {
Some(address) => address.to_string(),
None => env::var("APP_HOST").unwrap_or("0.0.0.0".to_string()),
};
let app_port = match port {
Some(port) => port.to_string(),
None => env::var("APP_PORT").unwrap_or("8000".to_string()),
};
// A type alias for a sender to a WebSocket connection.
//
// The sender is a mutex-guarded, split sink of a WebSocket stream and Message
// values. It is used to send messages to a client.
//
// The Mutex is used to ensure that only one thread can send a message at a
// time. This is because the SplitSink is not thread-safe, and sending a
// message from multiple threads could result in the messages being sent
// out of order.
//
// The SplitSink is used to send messages to a client. It is the part of the
// WebSocket stream that handles the sending of messages.
//
// The WebSocket stream is the underlying connection to the client. It is used
// to send and receive messages.
//
// The Message value is the type of data that is sent over the WebSocket
// connection. It is a struct that contains the data that is being sent.
//
// The type alias is used so that the type is not mentioned every time it is
// used. This makes the code easier to read and understand.
type Sender = Arc<Mutex<SplitSink<axum::extract::ws::WebSocket, Message>>>;
debug!("Server configured to accept connections on host {app_host}...");
debug!("Server configured to listen connections on port {app_port}...");
match app_environemt.as_str() {
"development" => {
debug!("Running in development mode");
}
"production" => {
debug!("Running in production mode");
}
_ => {
debug!("Running in development mode");
}
}
let app = Router::new()
.route("/status", get(status))
.route("/upload", post(upload_info))
.route("/download/:name", get(download_info))
.route("/download_success/:name", post(download_success))
.with_state(shared_state)
.layer(SecureClientIpSource::ConnectInfo.into_extension());
let listener = tokio::net::TcpListener::bind(format!("{}:{}", app_host, app_port).to_string())
.await
.unwrap();
axum::serve(
listener,
app.into_make_service_with_connect_info::<SocketAddr>(),
)
.with_graceful_shutdown(shutdown_signal())
.await
.unwrap();
}
async fn download_info(
State(shared_state): State<AppState>,
ConnectInfo(addr): ConnectInfo<SocketAddr>,
Path(name): Path<String>,
) -> impl IntoResponse {
debug!("Get new download request from: {}", addr.ip().to_string());
let data = shared_state.data.lock().unwrap();
match data.iter().find(|request| request.body.name == name) {
Some(request) => {
debug!("Found transfer name.");
(StatusCode::OK, Json(request.clone()))
}
None => {
warn!("couldn't find transfer-name: {}", name);
(
StatusCode::NOT_FOUND,
Json(TransferInfoRequest {
message: "error".to_string(),
body: TransferInfoBody {
keyword: "".to_string(),
files: "".to_string(),
ip: "".to_string(),
name: "".to_string(),
},
}),
)
}
}
}
async fn upload_info(
State(shared_state): State<AppState>,
ConnectInfo(addr): ConnectInfo<SocketAddr>,
Json(payload): Json<TransferInfoBody>,
) -> impl IntoResponse {
debug!("Got upload request from {}", addr.ip().to_string());
let mut data = shared_state.data.lock().unwrap();
let t_request = TransferInfoRequest {
message: "created".to_string(),
body: TransferInfoBody {
keyword: payload.keyword,
files: payload.files,
ip: payload.ip,
name: payload.name,
},
};
data.push(t_request.clone());
debug!("New TransferRequest created");
debug!("Actual AppState is {:#?}", *data);
(StatusCode::CREATED, Json(t_request))
}
async fn status() -> impl IntoResponse {
let version = env!("CARGO_PKG_VERSION");
let response = json!({
"data": {
"version": version,
#[derive(Debug, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "camelCase")]
// This enum is used to represent the different types of requests that a client
// can send to the server.
//
// The requests that a client can send are:
//
// * Join: A request to join a room. The request contains the ID of the room
// that the client wants to join.
// * Create: A request to create a new room.
// * Leave: A request to leave the current room.
pub enum RequestPacket {
Join {
// The ID of the room that the client wants to join.
id: String,
},
"message": "Service is running..."
});
(StatusCode::OK, Json(response))
Create,
Leave,
}
async fn download_success(
State(shared_state): State<AppState>,
Path(name): Path<String>,
) -> impl IntoResponse {
let mut data = shared_state.data.lock().unwrap();
if let Some(index) = data.iter().position(|request| request.body.name == name) {
debug!("Found Transfer by name '{name}'");
data.remove(index);
debug!("Transfer deleted");
(
StatusCode::OK,
Json(json!({
"message": "transfer deleted"
})),
)
} else {
warn!("couldn't find transfer-name: {}", name);
(
StatusCode::NOT_FOUND,
Json(json!({
"message": "transfer not found"
})),
)
/// This enum is used to represent the different types of responses that the
/// server can send to the client.
///
/// The responses that the server can send are:
///
/// * Join: A response to a `Join` request from the client. If the client
/// successfully joined a room, the `size` field will be `Some` and contain
/// the size of the room. If the client could not join a room, the `size` field
/// will be `None`.
/// * Create: A response to a `Create` request from the client. If the server
/// successfully created a room, the `id` field will contain the ID of the
/// room. If the server could not create a room, the `id` field will be empty.
/// * Leave: A response to a `Leave` request from the client. If the client
/// successfully left a room, the `index` field will contain the index of the
/// client that left the room. If the client could not leave a room, the
/// `index` field will be 0.
/// * Error: A response to indicate that an error occurred. The `message`
/// field will contain a description of the error.
#[derive(Debug, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "camelCase")]
pub enum ResponsePacket {
Join {
/// The size of the room that the client joined. If the client could
/// not join a room, this field will be `None`.
#[serde(skip_serializing_if = "Option::is_none")]
size: Option<usize>,
},
Create {
/// The ID of the room that the server created. If the server could
/// not create a room, this field will be empty.
id: String,
},
Leave {
/// The index of the client that left the room. If the client could not
/// leave a room, this field will be 0.
index: usize,
},
Error {
/// A description of the error that occurred.
message: String,
},
}
/// A `Room` is a collection of clients that are connected to each other.
///
/// Each room has a set of clients, represented by a `Vec` of `Sender`
/// instances. The `Sender` instances are used to send messages to the
/// clients in the room.
///
/// The `senders` field is the list of senders that are connected to each
/// other. Each sender is a mutex-guarded, split sink of a WebSocket
/// stream and Message values. This is explained in more detail in the
/// documentation for the `Sender` type alias in the `packets` module.
///
/// The `size` field is the maximum number of clients that a room can have.
/// When a room reaches its maximum size, no more clients can join the room.
/// This is used to prevent rooms from getting too full and causing the
/// server to run out of memory.
#[derive(Debug)]
pub struct Room {
senders: Vec<Sender>,
size: usize,
}
impl Room {
/// The default size of a room.
///
/// This is the size that a room will have when it is created.
const DEFAULT_ROOM_SIZE: usize = 2;
/// Creates a new `Room` with the given size.
///
/// The `size` parameter is the maximum number of clients that can join the
/// room. If `size` is 0, then the room will not be able to hold any
/// clients.
///
/// The `senders` field of the returned `Room` is an empty vector.
///
/// The `size` field of the returned `Room` is `size`.
fn new(size: usize) -> Room {
Room {
// Initialize the list of senders to be empty.
senders: Vec::new(),
// Set the size of the room.
size,
}
}
// match data.iter().find(|request| request.body.name == name) {
// Some(request) => {
// debug!("Found transfer name.");
// return (
// StatusCode::OK,
// Json(json!({
// "message" : "transfer deleted"
// })),
// );
// }
// None => {
// warn!("couldn't find transfer-name: {}", name);
// return (
// StatusCode::NOT_FOUND,
// Json(json!({
// "message" : "transfer not found"
// })),
// );
// }
// }
}
async fn shutdown_signal() {
let ctrl_c = async {
signal::ctrl_c()
/// A struct that holds all of the rooms that the server knows about.
///
/// The rooms are stored in a `HashMap` with the room ID as the key and the
/// room as the value. This means that looking up a room by its ID is an O(1)
/// operation, which is very fast.
#[derive(Debug)]
pub struct Server {
pub rooms: HashMap<String, Room>,
}
impl Server {
/// Creates a new `Server` with an empty list of rooms.
///
/// The `rooms` field of the returned `Server` is an empty `HashMap`.
/// This means that the server will not have any rooms when it is first
/// created.
///
/// This function returns an `Arc<RwLock<Server>>` because the server
/// needs to be shared between different parts of the program. The
/// `Arc` makes it so that the server can be shared by multiple threads,
/// and the `RwLock` makes it so that the server can be read from and
/// written to from multiple threads at the same time.
///
/// The `Arc` and `RwLock` are both parts of the `tokio` library, which
/// provides asynchronous programming tools for Rust.
///
/// The `Arc` and `RwLock` are used together to create a Mutex-like
/// object that can be shared between threads. The main difference
/// between a Mutex and an `Arc<RwLock<T>>` is that a Mutex can only be
/// locked by one thread at a time, while an `Arc<RwLock<T>>` can be
/// locked by multiple threads at the same time.
///
/// This function is used to create a new `Server` and share it between
/// different parts of the program. The `Server` is shared because it
/// needs to be able to handle connections from multiple clients at the
/// same time.
pub fn new() -> Arc<RwLock<Server>> {
// Create a new `Server` instance.
Arc::new(RwLock::new(Server {
// Initialize the list of rooms to be empty.
rooms: HashMap::new(),
}))
}
}
/// This struct represents a single client connection to the server.
///
/// A `Client` instance holds a `Sender` and a `room_id`.
///
/// The `Sender` is a type alias for a `tokio::sync::mpsc::Sender<Message>`.
/// It is used to send messages to the client.
///
/// The `room_id` is an `Option<String>`. It is used to keep track of which
/// room the client is currently in. If the `room_id` is `None`, then the
/// client is not in any room. If the `room_id` is `Some(id)`, where `id` is a
/// `String`, then the client is in the room with the ID `id`.
///
/// The `room_id` is used to keep track of which room the client is in so
/// that the server knows which room to send messages to. When a client
/// joins a room, their `room_id` is set to the ID of the room that they
/// joined. When a client leaves a room, their `room_id` is set to `None`.
///
/// The `Client` struct is used to keep track of which room each client is
/// in. It is used by the `Server` to determine which room to send messages
/// to.
///
#[derive(Debug)]
pub struct Client {
sender: Sender,
room_id: Option<String>,
}
impl Client {
/// Creates a new `Client` instance.
///
/// The `sender` argument is a `Sender` for sending messages to the client.
/// It is used by the `Server` to send messages to the client.
///
/// The `room_id` field of the `Client` instance is set to `None` initially.
/// This is because the client is not in any room when they first connect
/// to the server.
///
/// The `sender` field of the `Client` instance is used to send messages to
/// the client. When the server wants to send a message to the client, it
/// uses the `sender` to send the message.
///
/// The `Client` instance is used by the `Server` to keep track of which
/// room each client is in. It is used by the `Server` to determine which
/// room to send messages to.
pub fn new(sender: Sender) -> Client {
Client {
sender,
room_id: None,
}
}
/// Sends a message to a client.
///
/// This function takes a `sender` argument, which is a `Mutex` guard
/// for a WebSocket connection. The `sender` is used to send a message
/// to the client.
///
/// The `message` argument is the message that is sent to the client. It
/// is a WebSocket message.
///
/// This function locks the `sender` Mutex to ensure that only one thread
/// can send a message at a time. This is because the SplitSink that the
/// `sender` mutex guards is not thread-safe, and sending a message from
/// multiple threads could result in the messages being sent out of order.
///
/// If sending the message fails, this function logs an error message.
async fn send(&self, sender: Sender, message: Message) {
let mut sender = sender.lock().await;
if let Err(error) = sender.send(message).await {
error!("Failed to send message to the client: {}", error);
}
}
/// Sends a packet to a client.
///
/// This function takes a `sender` argument, which is a `Mutex` guard
/// for a WebSocket connection. The `sender` is used to send a message
/// to the client.
///
/// The `packet` argument is the packet that is sent to the client. It
/// is a struct that contains the data that is being sent.
///
/// This function serializes the `packet` using serde_json and sends it
/// to the client as a WebSocket Text message.
///
/// This function locks the `sender` Mutex to ensure that only one thread
/// can send a message at a time. This is because the SplitSink that the
/// `sender` mutex guards is not thread-safe, and sending a message from
/// multiple threads could result in the messages being sent out of order.
async fn send_packet(&self, sender: Sender, packet: ResponsePacket) {
let serialized_packet = serde_json::to_string(&packet).unwrap();
self.send(sender, Message::Text(serialized_packet)).await;
}
/// Sends an error packet to a client.
///
/// This function takes a `sender` argument, which is a `Mutex` guard
/// for a WebSocket connection. The `sender` is used to send a message
/// to the client.
///
/// The `message` argument is the message that is sent to the client. It
/// is a string that describes the error.
///
/// This function creates an error packet with the `message` and sends it
/// to the client using the `send_packet` function.
///
/// This function locks the `sender` Mutex to ensure that only one thread
/// can send a message at a time. This is because the SplitSink that the
/// `sender` mutex guards is not thread-safe, and sending a message from
/// multiple threads could result in the messages being sent out of order.
async fn send_error_packet(&self, sender: Sender, message: String) {
let error_packet = ResponsePacket::Error { message };
self.send_packet(sender, error_packet).await
}
/// Handles a CreateRoom request from a client.
///
/// This function is called when a client sends a CreateRoom request to
/// the server. The server will create a new room with the specified
/// size and return the room's identifier to the client.
///
/// This function takes a `server` argument, which is a `RwLock`
/// guard for the server's state. The `server` is used to check if the
/// current client is already in a room, and to insert the new room into
/// the server's state.
///
/// If the current client is already in a room, this function returns
/// without doing anything. This is to prevent a client from being in
/// multiple rooms at the same time.
///
/// If there is already a room with the same identifier as the one that
/// is being created, this function sends an error packet to the client
/// and returns.
///
/// If there is no existing room with the same identifier, this function
/// creates a new room with the specified size and inserts it into the
/// server's state. It then sends a CreateRoom response packet to the
/// client with the room's identifier.
///
/// This function locks the `server` RwLock to ensure that only one
/// thread can access the server's state at a time. This is because the
/// server's state is not thread-safe, and accessing it from multiple
/// threads could result in undefined behavior.
async fn handle_create_room(&mut self, server: &RwLock<Server>) {
let mut server = server.write().await;
// If the current client is already in a room, do nothing.
if server.rooms.iter().any(|(_, room)| {
room.senders
.iter()
.any(|sender| Arc::ptr_eq(sender, &self.sender))
}) {
return;
}
// Generate a new room identifier.
let size = Room::DEFAULT_ROOM_SIZE;
let room_id = Uuid::new_v4().to_string();
// If there is already a room with the same identifier, send an error
// packet to the client and return.
if server.rooms.contains_key(&room_id) {
drop(server);
return self
.send_error_packet(
self.sender.clone(),
"A room with that identifier already exists.".to_string(),
)
.await;
}
// Create a new room with the specified size and insert it into the
// server's state.
let mut room = Room::new(size);
room.senders.push(self.sender.clone());
server.rooms.insert(room_id.clone(), room);
// Set the client's room ID to the new room's identifier.
self.room_id = Some(room_id.clone());
drop(server);
// Send a CreateRoom response packet to the client with the room's
// identifier.
self.send_packet(self.sender.clone(), ResponsePacket::Create { id: room_id })
.await
.expect("failed to install Ctrl+C handler");
};
}
#[cfg(unix)]
let terminate = async {
signal::unix::signal(signal::unix::SignalKind::terminate())
.expect("failed to install signal handler")
.recv()
.await;
};
/// This function is called when the client sends a JoinRoom packet.
///
/// If the client is already in a room, then this function does nothing.
///
/// If the client is not in a room, then the function checks if the room
/// specified in the packet exists. If the room does not exist, an error
/// packet is sent to the client with a message indicating that the room
/// does not exist.
///
/// If the room does exist, then the function checks if the room is full.
/// If the room is full, an error packet is sent to the client with a
/// message indicating that the room is full.
///
/// If the room is not full, then the client is added to the room and the
/// function sends a JoinRoom response packet to the client with the size
/// of the room (excluding the client itself) and a `size` field set to
/// `None`. The response packet is sent to all other clients in the room.
async fn handle_join_room(&mut self, server: &RwLock<Server>, room_id: String) {
let mut server = server.write().await;
#[cfg(not(unix))]
let terminate = std::future::pending::<()>();
// If the client is already in a room, do nothing.
if server.rooms.iter().any(|(_, room)| {
room.senders
.iter()
.any(|sender| Arc::ptr_eq(sender, &self.sender))
}) {
return;
}
tokio::select! {
_ = ctrl_c => {},
_ = terminate => {},
// Get a mutable reference to the room specified in the packet.
// If the room does not exist, return an error to the client.
let Some(room) = server.rooms.get_mut(&room_id) else {
drop(server);
return self
.send_error_packet(self.sender.clone(), "The room does not exist.".to_string())
.await;
};
// If the room is full, return an error to the client.
if room.senders.len() >= room.size {
drop(server);
return self
.send_error_packet(self.sender.clone(), "The room is full.".to_string())
.await;
}
// Add the client to the room and set the client's room ID to the new
// room's identifier.
room.senders.push(self.sender.clone());
self.room_id = Some(room_id);
// Create a list of futures to send JoinRoom response packets to all
// other clients in the room. The `size` field of the response packet is
// set to `None` if the client sending the packet is the one joining the
// room. Otherwise, the `size` field is set to the number of clients in
// the room minus one (to exclude the client joining the room).
let mut futures = vec![];
for sender in &room.senders {
if Arc::ptr_eq(sender, &self.sender) {
futures.push(self.send_packet(
sender.clone(),
ResponsePacket::Join {
size: Some(room.senders.len() - 1),
},
));
} else {
futures.push(self.send_packet(sender.clone(), ResponsePacket::Join { size: None }));
}
}
drop(server);
join_all(futures).await;
}
/// Handles a request to leave a room.
///
/// This function is called when a client sends a `LeaveRoom` request
/// packet. The function obtains a write lock on the server's state and
/// does the following:
///
/// 1. Gets the room ID of the client who sent the request. If the client is
/// not in a room, the function returns early.
/// 2. Tries to get a mutable reference to the room with the obtained room
/// ID. If the room does not exist, the function returns early.
/// 3. Finds the index of the client's sender in the room's list of senders.
/// If the client is not in the room, the function returns early.
/// 4. Removes the client's sender from the room's list of senders.
/// 5. Sets the client's room ID to `None`.
/// 6. Creates a list of futures to send `LeaveRoom` response packets to
/// all other clients in the room. The `index` field of the response
/// packet is set to the index of the client's sender in the room's list
/// of senders.
/// 7. If the room is now empty, removes the room from the server's list
/// of rooms.
/// 8. Drops the write lock on the server's state.
/// 9. Waits for all futures to complete.
async fn handle_leave_room(&mut self, server: &RwLock<Server>) {
// Obtain a write lock on the server's state.
let mut server = server.write().await;
// Get the room ID of the client who sent the request.
let Some(room_id) = self.room_id.clone() else {
// If the client is not in a room, return early.
return;
};
// Try to get a mutable reference to the room with the obtained room ID.
let Some(room) = server.rooms.get_mut(&room_id) else {
// If the room does not exist, return early.
return;
};
// Find the index of the client's sender in the room's list of senders.
let Some(index) = room
.senders
.iter()
.position(|sender| Arc::ptr_eq(sender, &self.sender))
else {
// If the client is not in the room, return early.
return;
};
// Remove the client's sender from the room's list of senders.
room.senders.remove(index);
// Set the client's room ID to `None`.
self.room_id = None;
// Create a list of futures to send `LeaveRoom` response packets to
// all other clients in the room. The `index` field of the response
// packet is set to the index of the client's sender in the room's list
// of senders.
let mut futures = vec![];
for sender in &room.senders {
futures.push(self.send_packet(sender.clone(), ResponsePacket::Leave { index }));
}
// If the room is now empty, removes the room from the server's list
// of rooms.
if room.senders.is_empty() {
server.rooms.remove(&room_id);
}
// Drop the write lock on the server's state.
drop(server);
// Wait for all futures to complete.
join_all(futures).await;
}
/// This function handles an incoming message from a client.
///
/// The message can be one of four types: `Text`, `Binary`, `Ping`, or `Close`.
///
/// If the message is `Text`, the function parses the message as a `RequestPacket` and
/// calls the appropriate function to handle the request. If the message cannot be
/// parsed as a `RequestPacket`, the function does nothing and returns early.
///
/// If the message is `Binary`, the function first acquires a read lock on the server's
/// state. If the client is not currently in a room, the function drops the read lock and
/// returns early. If the client is not in a room, or if the room does not exist, the
/// function drops the read lock and returns early.
///
/// The function then finds the index of the client's sender in the room's list of
/// senders. If the client's sender is not in the room's list of senders, the function
/// drops the read lock and returns early.
///
/// The function then gets the binary data from the message and sets the first byte to
/// the index of the client's sender in the room's list of senders. If there is no
/// binary data in the message, the function drops the read lock and returns early.
///
/// The function then determines where to send the message. If the first byte of the
/// message is less than the number of clients in the room, the function sends the message
/// to the client at that index in the room's list of senders. If the first byte of the
/// message is equal to the number of clients in the room plus one, the function sends the
/// message to all clients in the room, excluding the client that sent the message.
///
/// If the first byte of the message is any other value, the function drops the read
/// lock and returns early.
///
/// Finally, the function drops the read lock and waits for all futures to complete.
///
/// If the message is `Ping`, the function prints a message to stdout.
///
/// If the message is `Pong`, the function prints a message to stdout.
///
/// If the message is `Close`, the function prints a message to stdout and calls the
/// `handle_close` function.
pub async fn handle_message(&mut self, server: &RwLock<Server>, message: Message) {
match message {
Message::Text(text) => {
let packet = match serde_json::from_str(&text) {
Ok(packet) => packet,
Err(_) => return,
};
match packet {
RequestPacket::Create => self.handle_create_room(server).await,
RequestPacket::Join { id } => self.handle_join_room(server, id).await,
RequestPacket::Leave => self.handle_leave_room(server).await,
}
}
Message::Binary(_) => {
// Acquire a read lock on the server's state.
let server = server.read().await;
// If the client is not currently in a room, return early.
let Some(room_id) = &self.room_id else {
drop(server);
return;
};
// If the room does not exist, return early.
let Some(room) = server.rooms.get(room_id) else {
drop(server);
return;
};
// Find the index of the client's sender in the room's list of senders.
let Some(index) = room
.senders
.iter()
.position(|sender| Arc::ptr_eq(sender, &self.sender))
else {
drop(server);
return;
};
// Get the binary data from the message and set the first byte to
// the index of the client's sender in the room's list of senders.
let mut data = message.into_data();
if data.is_empty() {
drop(server);
return;
}
let source = u8::try_from(index).unwrap();
// Determine where to send the message.
let destination = usize::from(data[0]);
data[0] = source;
// Send the message to the client at the destination index in the
// room's list of senders.
if destination < room.senders.len() {
let sender = room.senders[destination].clone();
drop(server);
return self.send(sender, Message::Binary(data)).await;
}
// Send the message to all clients in the room, excluding the
// client that sent the message.
if destination == usize::from(u8::MAX) {
let mut futures = vec![];
for sender in &room.senders {
if Arc::ptr_eq(sender, &self.sender) {
continue;
}
futures.push(self.send(sender.clone(), Message::Binary(data.clone())));
}
drop(server);
join_all(futures).await;
}
}
Message::Ping(_) => {
println!("Got Message Type Ping");
}
Message::Pong(_) => {
println!("Got Message Type Pong");
}
Message::Close(_) => {
println!("Got Message Type Close");
self.handle_close(server).await;
}
}
}
pub async fn handle_close(&mut self, server: &RwLock<Server>) {
self.handle_leave_room(server).await
}
}

746
src/sender/client.rs
View file

@ -1,67 +1,707 @@
use crate::error::TransferNotCreatedError;
use crate::transfer_info::transfer_info::TransferInfoRequest;
use hex;
use local_ip_address;
use rand::{seq::SliceRandom, thread_rng};
use reqwest::{Client, StatusCode};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use tracing::debug;
use crate::shared::{
packets::{
list_packet, packet::Value, ChunkPacket, HandshakePacket, HandshakeResponsePacket,
ListPacket, Packet, ProgressPacket,
},
JsonPacket, JsonPacketResponse, JsonPacketSender, PacketSender, Sender, Socket, Status,
};
type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;
use aes_gcm::{aead::Aead, Aes128Gcm, Key};
use base64::{engine::general_purpose, Engine as _};
use futures_util::{future, pin_mut, stream::TryStreamExt, StreamExt};
use hmac::{Hmac, Mac};
use p256::{ecdh::EphemeralSecret, PublicKey};
use prost::Message;
use rand::{rngs::OsRng, RngCore};
use sha2::Sha256;
use std::{
fs,
io::{stdout, Write},
path::Path,
time::Duration,
};
use tokio::{io::AsyncReadExt, task::JoinHandle, time::sleep};
use tokio_tungstenite::tungstenite::{protocol::Message as WebSocketMessage, Error};
use tracing::{debug, error};
pub async fn send_info(relay: &str, file: &str) -> Result<String> {
let sender_ip = local_ip_address::local_ip().unwrap();
let ip_str = sender_ip.to_owned().to_string();
let rand_name = generate_random_name();
let hashed_name = Sha256::digest(rand_name.as_bytes());
let hashed_string = hex::encode(hashed_name);
debug!("local ip is: {}", sender_ip);
debug!("Send Request to: {:?}", relay.to_string());
let mut map = HashMap::new();
map.insert("keyword", "test");
map.insert("files", file);
map.insert("ip", ip_str.as_str());
map.insert("name", hashed_string.as_str());
const DESTINATION: u8 = 1;
const NONCE_SIZE: usize = 12;
const MAX_CHUNK_SIZE: isize = u16::MAX as isize;
const DELAY: Duration = Duration::from_millis(750);
const URL: &str = "https://0.0.0.0:8000/#";
let client = Client::new();
let res = client
.post(format!("{}/upload", relay))
.json(&map)
.send()
.await?;
if res.status() == StatusCode::CREATED {
let transfer_info: TransferInfoRequest = res.json().await?;
debug!("Json Response: {:#?}", transfer_info);
Ok(rand_name)
} else {
Err(Box::new(TransferNotCreatedError::new(
"Transfer could not be created.",
)))
/// A file that is to be sent.
///
/// This structure contains all the information about a file that is to be
/// sent. It is used to keep track of the files that a user wants to send.
#[derive(Clone)]
struct File {
/// The path to the file on the file system.
///
/// This is the path to the file on the user's file system. The path is
/// used to open the file and read its contents.
path: String,
/// The name of the file.
///
/// This is the name that the file will have when it is received by the
/// receiver. This name is used when creating the file on the receiver's
/// file system.
name: String,
/// The size of the file in bytes.
///
/// This is the size of the file in bytes. The size is used to calculate
/// the number of chunks that the file will be split into, and is also
/// used to keep track of the progress of the file being sent.
size: u64,
}
/// The context for the sender.
///
/// This structure contains all the information that the sender needs in order
/// to function properly. It is used to keep track of the state of the
/// sender, and to pass information between functions.
struct Context {
/// The HMAC key for the sender.
///
/// This is the key that is used to sign packets. The key is also used to
/// generate a URL that the receiver can use to join the session.
hmac: Vec<u8>,
/// The sender that is used to send packets to the receiver.
///
/// This sender is used to send handshake packets, list packets, chunk
/// packets, and progress packets to the receiver.
sender: Sender,
/// The ephemeral keypair that is used to establish a shared key with the
/// receiver.
///
/// This key is used to establish a shared key between the sender and
/// receiver. The key is ephemeral, meaning that it is only used once in
/// the session. The key is generated when the sender is created, and is
/// then discarded after the session is complete.
key: EphemeralSecret,
/// The files that the sender wants to send.
///
/// This vec contains all the information about the files that the sender
/// wants to send. The vec is filled when the user specifies the files to
/// send using the command line arguments.
files: Vec<File>,
/// The shared key that is used to encrypt packets.
///
/// This value is set to `None` initially, and is set to `Some` when the
/// shared key is established with the receiver. The shared key is used to
/// encrypt packets that are sent to the receiver.
shared_key: Option<Aes128Gcm>,
/// The task that is running in the background to send chunks of files to
/// the receiver.
///
/// This task is created when the sender is created, and is used to send
/// chunks of files to the receiver in the background. The task is
/// initially set to `None`, but is set to `Some` when the task is spawned.
/// The task is used to cancel the background task when the sender is
/// dropped.
task: Option<JoinHandle<()>>,
}
/// This function is called when the client receives a create room packet
/// from the server. The function is responsible for printing a URL to the
/// console that the user can use to join the room.
///
/// The function first generates a base64 string from the hmac value that is
/// used to verify the integrity of the room. The base64 string is then
/// 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) -> Status {
let base64 = general_purpose::STANDARD.encode(&context.hmac);
let url = format!("{}{}-{}", URL, id, base64);
// Print a newline to the console to separate the output from the command
// line.
println!();
// Try to generate a QR code from the URL. If the function fails for some
// reason, print an error message to the console.
if let Err(error) = qr2term::print_qr(&url) {
error!("Failed to generate QR code: {}", error);
}
// Print a newline to the console to separate the output from the command
// line.
println!();
// Print a message to the console with the URL.
println!("Created room: {}", url);
// Continue the event loop.
Status::Continue()
}
/// This function is called when the client receives a join room packet from
/// the server. The function is responsible for sending a handshake packet to
/// the server containing the client's public key and a signature generated
/// using the client's private key and the room's hmac value.
///
/// The function first generates the client's public key from the private key.
/// The public key is then serialized into a byte array.
///
/// Next, the function creates a HMAC object with the room's hmac value and
/// updates it with the serialized public key. The resulting HMAC is then
/// serialized into a byte array and used as the signature in the handshake
/// packet.
///
/// Finally, the function sends the handshake packet to the server using the
/// sender object.
fn on_join_room(context: &Context, size: Option<usize>) -> Status {
if size.is_some() {
return Status::Err("Invalid join room packet.".into());
}
// Generate the client's public key from the private key.
let public_key = context.key.public_key().to_sec1_bytes().into_vec();
// Create a HMAC object with the room's hmac value and update
// it with the serialized public key.
let mut mac = Hmac::<Sha256>::new_from_slice(&context.hmac).unwrap();
mac.update(&public_key);
// Serialize the resulting HMAC into a byte array and use it as the
// signature in the handshake packet.
let signature = mac.finalize().into_bytes().to_vec();
// Create the handshake packet and send it to the server.
let handshake = HandshakePacket {
public_key,
signature,
};
context
.sender
.send_packet(DESTINATION, Value::Handshake(handshake));
Status::Continue()
}
/// This function is called when an error packet is received from the
/// server. It creates a `Status::Err` variant containing the error
/// message from the server and returns it to be handled by the main
/// event loop.
///
/// When an error occurs, the server sends an error packet to the
/// client. The error packet contains a message with a description of
/// the error. This function extracts that message and creates a
/// `Status::Err` variant with it, which is then returned to be handled
/// by the main event loop.
///
/// The main event loop checks the status of the client and performs
/// the necessary actions based on its value. If the status is
/// `Status::Err`, the event loop exits with an error message
/// containing the error message from the server.
///
/// This function is called from the event loop when an error packet is
/// received from the server.
fn on_error(message: String) -> Status {
Status::Err(message)
}
/// This function is called when the server sends a leave room packet to
/// the client. It is responsible for aborting the file transfer task,
/// generating a new ECDH key pair for the next handshake, and setting the
/// shared key to `None`.
///
/// When the server sends a leave room packet to the client, it means that
/// the receiver has disconnected from the room. In this case, the client
/// should abort the file transfer task and print an error message to the
/// user.
///
/// If the client is currently transferring files, it should abort the task
/// by calling `AbortHandle::abort` on the task handle.
///
/// After that, the client should generate a new ECDH key pair using the
/// `EphemeralSecret::random` function from the `p256` crate. This key pair
/// will be used for the next handshake with the server.
///
/// Finally, the client should set the shared key to `None` to indicate that
/// there is no shared key established for the current room.
///
/// This function is called from the event loop when a leave room packet is
/// received from the server.
fn on_leave_room(context: &mut Context, _: usize) -> Status {
if let Some(task) = &context.task {
// If the client is currently transferring files, abort the task
// by calling `AbortHandle::abort` on the task handle.
task.abort();
}
// Generate a new ECDH key pair for the next handshake.
context.key = EphemeralSecret::random(&mut OsRng);
// Set the shared key to `None` to indicate that there is no shared key
// established for the current room.
context.shared_key = None;
// Set the task handle to `None` to indicate that there is no task
// running.
context.task = None;
// Print an error message to the user indicating that the transfer was
// interrupted because the receiver disconnected.
println!();
error!("Transfer was interrupted because the receiver disconnected.");
// Continue the event loop.
Status::Continue()
}
/// This function is called by the event loop when a progress packet is
/// received from the server.
///
/// The progress packet contains the index of the file that is being
/// transferred and the current progress of that file as a percentage.
///
/// If the client does not have a shared key established with the server,
/// the function returns an error and does not continue. This indicates
/// that the event loop should exit with an error message.
///
/// The function then retrieves the file at the index specified by the
/// progress packet from the context. If the index is out of bounds, the
/// function returns an error and does not continue. This indicates that
/// the event loop should exit with an error message.
///
/// The function then prints a message to the console indicating which file
/// is currently being transferred and what its progress is. The progress
/// message is printed to the same line as a carriage return (`\r`) so that
/// it overwrites the previous message.
///
/// If the progress of the file is 100%, the function prints a newline
/// (`\n`) to the console to move the cursor to the next line.
///
/// If the progress of the last file is 100%, the function returns
/// `Status::Exit()`. This indicates that the event loop should exit
/// successfully.
///
/// If any other condition is met, the function returns `Status::Continue()`.
/// This indicates that the event loop should continue running.
fn on_progress(context: &Context, progress: ProgressPacket) -> Status {
if context.shared_key.is_none() {
return Status::Err("Invalid progress packet: no shared key established".into());
}
let file = match context.files.get(progress.index as usize) {
Some(file) => file,
None => return Status::Err("Invalid index in progress packet.".into()),
};
print!("\rTransferring '{}': {}%", file.name, progress.progress);
stdout().flush().unwrap();
if progress.progress == 100 {
println!();
if progress.index as usize == context.files.len() - 1 {
return Status::Exit();
}
}
Status::Continue()
}
/// This function reads a file in chunks, sends each chunk to the receiver over
/// the WebSocket connection, and then sleeps for a short amount of time
/// before sending the next chunk.
///
/// The function takes the sender, the shared key, and a vector of files to
/// transfer as arguments.
///
/// For each file in the vector of files, the function reads the file in
/// chunks, sends each chunk to the receiver over the WebSocket connection,
/// and then sleeps for a short amount of time before sending the next chunk.
///
/// The chunk size is set to the maximum chunk size. If the number of bytes
/// left to read in the file is less than the chunk size, the chunk size is set
/// to the number of bytes left to read.
///
/// The function opens the file for reading using the tokio::fs::File::open
/// function. If there is an error opening the file, the function prints an
/// error message to the console and returns.
///
/// The function reads the file in chunks using the read_exact function from
/// the tokio::io::AsyncReadExt trait. If there is an error reading from the
/// file, the function prints an error message to the console and returns.
///
/// The function sends each chunk to the receiver over the WebSocket
/// connection using the send_encrypted_packet function from the Sender struct.
/// The function also increments the sequence number for each chunk that is
/// sent.
///
/// After sending all of the chunks for a file, the function sleeps for a short
/// amount of time using the tokio::time::sleep function. This helps to prevent
/// the sender from overwhelming the receiver with too many messages.
///
/// The function repeats this process for all of the files in the vector of
/// files.
async fn on_chunk(sender: Sender, shared_key: Option<Aes128Gcm>, files: Vec<File>) {
for file in files {
// Initialize a sequence number for the chunks of this file
let mut sequence = 0;
// Set the chunk size to the maximum chunk size
let mut chunk_size = MAX_CHUNK_SIZE;
// Set the number of bytes left to read in the file
let mut size = file.size as isize;
// Open the file for reading
let mut handle = match tokio::fs::File::open(file.path).await {
Ok(handle) => handle,
Err(error) => {
println!("Error: Unable to open file '{}': {}", file.name, error);
return;
}
};
while size > 0 {
// If the number of bytes left to read in the file is less than the
// chunk size, set the chunk size to the number of bytes left to read
if size < chunk_size {
chunk_size = size;
}
// Create a vector to hold the chunk of data to be read from the file
let mut chunk = vec![0u8; chunk_size.try_into().unwrap()];
// Read a chunk of data from the file into the vector
handle.read_exact(&mut chunk).await.unwrap();
// Send the chunk to the receiver over the WebSocket connection
sender.send_encrypted_packet(
&shared_key,
DESTINATION,
Value::Chunk(ChunkPacket { sequence, chunk }),
);
// Increment the sequence number for the next chunk
sequence += 1;
// Decrement the number of bytes left to read in the file
size -= chunk_size;
}
// Sleep for a short amount of time to prevent overwhelming the receiver
// with too many messages
sleep(DELAY).await;
}
}
fn generate_random_name() -> String {
let mut rng = thread_rng();
let adjective = adjectives().choose(&mut rng).unwrap();
// let adjective = adjectives().sample(&mut rng).unwrap();
let noun1 = nouns1().choose(&mut rng).unwrap();
let noun2 = nouns2().choose(&mut rng).unwrap();
/// This function sends a ListPacket to the receiver containing the list of
/// files to be transferred. The ListPacket contains a vector of Entry structs,
/// each of which represents one file.
///
/// The function creates a vector of Entry structs from the vector of File structs
/// in the Context struct. Each Entry struct contains the index, name, and size
/// of the corresponding File struct.
///
/// The function then sends the ListPacket to the receiver using the send_encrypted_packet
/// function from the Sender struct.
///
/// After sending the ListPacket, the function spawns a task using tokio::spawn to
/// call the on_chunk function with the Sender, shared_key, and vector of File
/// structs as arguments. The on_chunk function will send each chunk of data for
/// each file to the receiver.
///
/// The function returns Status::Continue(), which tells the main loop to continue
/// running until all of the files have been transferred.
fn on_handshake_finalize(context: &mut Context) -> Status {
let mut entries = vec![];
format!("{adjective}-{noun1}-{noun2}")
for (index, file) in context.files.iter().enumerate() {
let entry = list_packet::Entry {
// The index of the file in the vector of Files in the Context struct
index: index.try_into().unwrap(),
// The name of the file
name: file.name.clone(),
// The size of the file in bytes
size: file.size,
};
entries.push(entry);
}
context.sender.send_encrypted_packet(
&context.shared_key,
DESTINATION,
Value::List(ListPacket { entries }),
);
context.task = Some(tokio::spawn(on_chunk(
context.sender.clone(),
context.shared_key.clone(),
context.files.clone(),
)));
Status::Continue()
}
fn adjectives() -> &'static [&'static str] {
static ADJECTIVES: &[&str] = &["funny", "smart", "creative", "friendly", "great"];
ADJECTIVES
/// Handshake function that is called when the Sender receives a HandshakeResponsePacket
/// from the Receiver. This function verifies the signature from the Receiver and if
/// successful, creates a shared key using the from the PublicKey struct.
///
/// The shared key is used to encrypt and decrypt packets sent between the Sender
/// and the Receiver.
///
/// This function is called by the main loop in client.rs.
fn on_handshake(context: &mut Context, handshake_response: HandshakeResponsePacket) -> Status {
if context.shared_key.is_some() {
// If the shared key is already established, this means that the Sender
// has already performed the handshake, so return an error.
return Status::Err("Already performed handshake.".into());
}
// Create a new HMAC using the hmac from the Context struct as the key.
let mut mac = Hmac::<Sha256>::new_from_slice(&context.hmac).unwrap();
// Update the HMAC with the public key from the HandshakeResponsePacket.
mac.update(&handshake_response.public_key);
// Call verify_slice() on the HMAC to verify the signature from the Receiver.
// If the signature is invalid, return an error.
let verification = mac.verify_slice(&handshake_response.signature);
if verification.is_err() {
return Status::Err("Invalid signature from the receiver.".into());
}
// Create a new PublicKey struct from the public key bytes in the
// HandshakeResponsePacket.
let shared_public_key = PublicKey::from_sec1_bytes(&handshake_response.public_key).unwrap();
// Use the diffie_hellman() method from the PublicKey struct to create a shared
// secret key between the Sender and the Receiver. The shared secret key is a
// 16 byte long slice of bytes.
let shared_secret = context.key.diffie_hellman(&shared_public_key);
let shared_secret = shared_secret.raw_secret_bytes();
let shared_secret = &shared_secret[0..16];
// Create a new Key struct from the shared secret key. The Key<Aes128Gcm> type
// is used to encrypt and decrypt packets.
let shared_key: &Key<Aes128Gcm> = shared_secret.into();
let shared_key = <Aes128Gcm as aes_gcm::KeyInit>::new(shared_key);
// Set the shared_key field of the Context struct to the shared key.
context.shared_key = Some(shared_key);
// Call on_handshake_finalize() to start the transfer of files between the
// Sender and the Receiver.
on_handshake_finalize(context)
}
fn nouns1() -> &'static [&'static str] {
static NOUNS1: &[&str] = &["dog", "cat", "flower", "tree", "house"];
NOUNS1
/// This function is called by the `Sender` when a new message is received over
/// the WebSocket connection. The message could be a text message or a binary
/// message. If it is a text message, it will be deserialized into a
/// `JsonPacketResponse` enum. If it is a binary message, it will be decrypted
/// if necessary and then deserialized into a `Packet` struct.
///
/// The `JsonPacketResponse` enum will have one of the following variants:
///
/// * `Create { id }`: The Receiver has created a new room with the given ID.
/// * `Join { size }`: The Receiver has joined a room with `size` number of
/// files.
/// * `Leave { index }`: The Receiver has left a room.
/// * `Error { message }`: The Receiver has encountered an error.
///
/// If the message is a binary message, the `Packet` struct will have a
/// `Value` variant that will have one of the following variants:
///
/// * `HandshakeResponse`: The Receiver has responded to the Sender's
/// `Handshake` packet.
/// * `Progress`: The Receiver has sent progress information for one of the
/// files in the room.
///
/// This function does the following:
///
/// * If the message is a text message, it is deserialized into a
/// `JsonPacketResponse` enum and then matched on to call the appropriate
/// function.
/// * If the message is a binary message, it is decrypted if necessary and then
/// deserialized into a `Packet` struct. The `Value` variant of the `Packet`
/// struct is then matched on to call the appropriate function.
///
/// If the message is invalid, an error is returned.
fn on_message(context: &mut Context, message: WebSocketMessage) -> 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),
JsonPacketResponse::Join { size } => on_join_room(context, size),
JsonPacketResponse::Leave { index } => on_leave_room(context, index),
JsonPacketResponse::Error { message } => on_error(message),
};
} else if message.is_binary() {
let data = message.into_data();
let data = &data[1..];
let data = if let Some(shared_key) = &context.shared_key {
let nonce = &data[..NONCE_SIZE];
let ciphertext = &data[NONCE_SIZE..];
shared_key.decrypt(nonce.into(), ciphertext).unwrap()
} else {
data.to_vec()
};
let packet = Packet::decode(data.as_ref()).unwrap();
let value = packet.value.unwrap();
return match value {
Value::HandshakeResponse(handshake_response) => {
on_handshake(context, handshake_response)
}
Value::Progress(progress) => on_progress(context, progress),
_ => Status::Err(format!("Unexpected packet: {:?}", value)),
};
}
Status::Err("Invalid message type".into())
}
fn nouns2() -> &'static [&'static str] {
static NOUNS2: &[&str] = &["cookie", "cake", "frosting"];
NOUNS2
/// Starts the sender client. This function will attempt to create a room with a size of 2
/// (the number of clients that will be joining the room) and then it will open a file for
/// each of the paths provided. It will then read chunks of data from each file and send them
/// to the server.
///
/// This function takes two arguments:
/// 1. `socket`: A `Socket` that represents the connection to the server.
/// 2. `paths`: A `Vec` of `String`s that represent the paths to the files that will be sent
/// to the server.
///
/// When the function is finished, it will exit and the transfer will be complete. If there
/// is an error during the transfer, the function will print an error message to stdout and
/// exit.
pub async fn start(socket: Socket, paths: Vec<String>) {
// Create a vector to store metadata about each file that will be sent.
let mut files = vec![];
// For each path in the `paths` vector:
for path in paths {
// Attempt to open the file at the given path.
let handle = match fs::File::open(&path) {
// If the file is successfully opened, store it in the `handle` variable.
Ok(handle) => handle,
// If there is an error, print an error message to stdout and exit the function.
Err(error) => {
error!("Error: Failed to open file '{}': {}", path, error);
return;
}
};
// Get the metadata for the file.
let metadata = handle.metadata().unwrap();
// If the file is a directory, print an error message to stdout and exit the function.
if metadata.is_dir() {
error!("Error: The path '{}' does not point to a file.", path);
return;
}
// Get the file name from the path.
let name = Path::new(&path).file_name().unwrap().to_str().unwrap();
// Get the file size from the metadata.
let size = metadata.len();
// If the file is empty, print an error message to stdout and exit the function.
if size == 0 {
error!("Error: The file '{}' is empty and cannot be sent.", name);
return;
}
// Add the file metadata to the `files` vector.
files.push(File {
name: name.to_string(),
path,
size,
});
}
// Generate a random key for HMAC.
let mut hmac = [0u8; 32];
OsRng.fill_bytes(&mut hmac);
// Generate a random key for AES-GCM.
let key = EphemeralSecret::random(&mut OsRng);
// Create a channel to send packets to the server.
let (sender, receiver) = flume::bounded(1000);
// Split the socket into separate send and receive streams.
let (outgoing, incoming) = socket.split();
// Create a context that will be used throughout the transfer.
let mut context = Context {
// Store the sender half of the channel to send packets to the server.
sender,
// Store the ephemeral key for AES-GCM.
key,
// Store the files that will be sent to the server.
files,
// Store the HMAC key.
hmac: hmac.to_vec(),
// Set the shared key to None.
shared_key: None,
// Set the current task to None.
task: None,
};
// Print a message to stdout indicating that the client is attempting to create a room.
debug!("Attempting to create room...");
// Send a JSON packet to the server to create a room with a size of 2.
context.sender.send_json_packet(JsonPacket::Create);
// Create a future that handles the outgoing stream of messages from the client to the
// server.
let outgoing_handler = receiver.stream().map(Ok).forward(outgoing);
// Create a future that handles the incoming stream of messages from the server to the
// 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) {
// If the status is `Status::Exit`, the transfer is complete. Print a message to
// stdout and exit the function.
Status::Exit() => {
// TODO: Signal Exit to the server
println!("Transfer has completed.");
// Exit the function with a `Result` of `Err`.
return future::err(Error::ConnectionClosed);
}
// If the status is `Status::Err`, there was an error. Print an error message to
// stdout and exit the function.
Status::Err(error) => {
error!("Error: {}", error);
// Exit the function with a `Result` of `Err`.
return future::err(Error::ConnectionClosed);
}
// Otherwise, the message was handled successfully.
_ => {}
};
// Continue handling the incoming messages.
future::ok(())
});
// Pin the `incoming_handler` and `outgoing_handler` futures so that they do not move.
pin_mut!(incoming_handler, outgoing_handler);
// Wait for either the `incoming_handler` or `outgoing_handler` to complete. If the
// `incoming_handler` completes, return the result of the `incoming_handler`. If the
// `outgoing_handler` completes, return the result of the `outgoing_handler`.
future::select(incoming_handler, outgoing_handler).await;
}

70
src/sender/mod.rs
View file

@ -1,2 +1,70 @@
/// Connects to the WebSocket server at `ws://0.0.0.0:8000/ws` with an
/// `Origin` header of `ws://0.0.0.0:8000/ws`. This is the URL that the
/// sender and receiver clients will connect to.
///
/// The `start_sender` function takes a reference to a vector of strings,
/// which are the paths to the files that the sender will send over the
/// WebSocket connection.
///
/// The function first creates a WebSocket request using the `IntoClientRequest`
/// trait from `tungstenite`, which is defined on the `IntoClientRequest` struct.
/// This struct is a type that represents a request to a WebSocket server.
///
/// The `into_client_request` function returns a `Result` because it may fail
/// to create the request. In this case, we do not handle the error, so we just
/// return if the result is an error.
///
/// Once we have a request, we insert the `Origin` header into the headers of
/// the request. This is necessary because the WebSocket protocol requires the
/// `Origin` header to be present in the handshake.
///
/// After that, we print out a message to the console indicating that we are
/// attempting to connect to the server.
///
/// Next, we call the `connect_async` function from `tokio_tungstenite` which
/// takes our request and attempts to connect to the server. This function
/// returns a `Future` that resolves to a tuple of a `WebSocketStream` and a
/// `Response` from the server. The `WebSocketStream` is a stream of
/// WebSocket messages from the server, and the `Response` is the response
/// from the server to our handshake request.
///
/// If connecting to the server fails, we print out an error message and
/// return.
///
/// If connecting to the server succeeds, we pass the `WebSocketStream` and
/// the paths to the files to the `start` function from the `sender` module.
/// The `start` function is defined in the `sender` module, and it is the
/// function that sends the files over the WebSocket connection.
///
/// The `start` function takes ownership of the `WebSocketStream` and the file
/// paths, so we pass it the `paths` vector by value.
pub mod client;
pub mod server;
use crate::sender::client as sender;
use tokio_tungstenite::{
connect_async,
tungstenite::{client::IntoClientRequest, http::HeaderValue},
};
use tracing::{debug, error};
// const ORIGIN: &str = "ws://0.0.0.0:8000/ws";
pub async fn start_sender(relay: &str, files: &[String]) {
let Ok(mut request) = relay.into_client_request() else {
println!("Error: Failed to create request.");
return;
};
request
.headers_mut()
.insert("Origin", HeaderValue::from_str(relay).unwrap());
debug!("Attempting to connect...");
let Ok((socket, _)) = connect_async(request).await else {
error!("Error: Failed to connect.");
return;
};
let paths = files.to_vec();
sender::start(socket, paths).await
}

57
src/sender/server.rs
View file

@ -1,57 +0,0 @@
use axum::{
extract::Json,
http::StatusCode,
response::IntoResponse,
routing::{get, post},
Router,
};
use lazy_static::lazy_static;
use serde_json::json;
use std::{net::SocketAddr, sync::Arc, time::Duration};
use tokio::sync::Mutex;
use tower_http::services::ServeFile;
use tracing::debug;
lazy_static! {
static ref SHUTDOWN_SIGNAL: Arc<Mutex<bool>> = Arc::new(Mutex::new(false));
}
pub async fn serf_file(path: &String) {
debug!("Sender starting...");
let app_host = "0.0.0.0".to_string();
let app_port = "8100".to_string();
debug!("Server configured to accept connections on host {app_host}...");
debug!("Server configured to listen connections on port {app_port}...");
let app = Router::new()
.route_service("/download_file", ServeFile::new(path))
.route("/ping", get(ping))
.route("/shutdown", post(shutdown));
let listener = tokio::net::TcpListener::bind(format!("{}:{}", app_host, app_port).to_string())
.await
.unwrap();
axum::serve(
listener,
app.into_make_service_with_connect_info::<SocketAddr>(),
)
.with_graceful_shutdown(async {
while !*SHUTDOWN_SIGNAL.lock().await {
tokio::time::sleep(Duration::from_millis(100)).await;
}
})
.await
.unwrap();
}
async fn ping() -> impl IntoResponse {
let response = json!({
"message": "pong"
});
(StatusCode::OK, Json(response))
}
async fn shutdown() -> impl IntoResponse {
debug!("Initiating server shutdown...");
*SHUTDOWN_SIGNAL.lock().await = true;
debug!("Server is shutting down...");
(StatusCode::OK, "Server is shutting down...")
}

333
src/shared.rs Normal file
View file

@ -0,0 +1,333 @@
pub mod packets {
include!(concat!(env!("OUT_DIR"), "/packets.rs"));
}
use aes_gcm::{
aead::{Aead, AeadCore},
Aes128Gcm,
};
use packets::Packet;
use prost::Message;
use rand::rngs::OsRng;
use serde::{Deserialize, Serialize};
use tokio::net::TcpStream;
use tokio_tungstenite::tungstenite::protocol::Message as WebSocketMessage;
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
/// This struct is used to serialize/deserialize JSON packets sent
/// between the client and the server.
///
/// The `type` field is used to specify the type of packet that is being sent.
/// The possible values for this field are listed as variants of the enum.
#[derive(Debug, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "camelCase")]
pub enum JsonPacket {
/// Sent from the client to ask to join a room.
///
/// The `id` field specifies the ID of the room that the client wants
/// to join.
Join {
/// The ID of the room that the client wants to join.
id: String,
},
/// Sent from the client to ask to create a new room.
Create,
/// Sent from the client to ask to leave the current room.
Leave,
}
/// This struct is used to serialize/deserialize JSON packets sent
/// from the server to the client.
///
/// The `type` field is used to specify the type of packet that is being
/// sent. The possible values for this field are listed as variants of the
/// enum.
#[derive(Debug, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "camelCase")]
pub enum JsonPacketResponse {
/// Sent from the server to inform the client of the result of a `Join`
/// packet.
///
/// If the client successfully joined a room, the `size` field will be
/// `Some` and contain the size of the room. If the client could not join
/// a room, the `size` field will be `None`.
Join {
/// The size of the room that the client joined. If the client could
/// not join a room, this field will be `None`.
#[serde(skip_serializing_if = "Option::is_none")]
size: Option<usize>,
},
/// Sent from the server to inform the client of the result of a `Create`
/// packet.
///
/// If the server successfully created a room, the `id` field will
/// contain the ID of the room. If the server could not create a room,
/// the `id` field will be empty.
Create {
/// The ID of the room that the server created. If the server could
/// not create a room, this field will be empty.
id: String,
},
/// Sent from the server to inform the client of the result of a `Leave`
/// packet.
///
/// If the client successfully left a room, the `index` field will
/// contain the index of the client that left the room. If the client
/// could not leave a room, the `index` field will be 0.
Leave {
/// The index of the client that left the room. If the client could
/// not leave a room, this field will be 0.
index: usize,
},
/// Sent from the server to inform the client of an error.
///
/// The `message` field contains a description of the error.
Error {
/// A description of the error that occurred.
message: String,
},
}
/// This enum represents the result of processing an event in the event loop.
///
/// The `Status` enum has three variants:
///
/// * `Continue` - This variant indicates that the event loop should
/// continue processing events. This is the most common result and is used
/// when the event loop has nothing special to do.
///
/// * `Exit` - This variant indicates that the event loop should exit. This
/// is used when the event loop should exit because of an error or
/// because the user has requested that the program exit.
///
/// * `Err` - This variant indicates that the event loop encountered an
/// error. When the event loop receives a `Status::Err` variant, it should
/// exit with an error message containing the message from the error packet.
/// The message from the error packet is the only information that the event
/// loop has about the error, so the message should be descriptive and
/// helpful to the user. The message should not contain technical details
/// about the error or how it occurred. Instead, the message should be
/// written from the perspective of the user and should give the user enough
/// information to understand what went wrong and how they might be able to
/// fix the problem.
pub enum Status {
/// Indicates that the event loop should continue processing events.
Continue(),
/// Indicates that the event loop should exit.
Exit(),
/// Indicates that the event loop encountered an error.
Err(String),
}
/// A trait for sending JSON packets.
///
/// This trait provides a single method, `send_json_packet`, which sends a
/// JSON packet over some underlying transport.
pub trait JsonPacketSender {
/// Sends a JSON packet.
///
/// This method takes a single argument, `packet`, which is the JSON packet
/// to send. The packet will be serialized into a JSON string and then sent
/// over the underlying transport.
///
/// Note that the exact semantics of what it means to "send a JSON packet"
/// will depend on the specific implementation of this trait. However, in
/// general, the packet will be sent as a single message over the
/// transport, and the transport will be responsible for ensuring that the
/// packet is delivered to the intended recipient.
///
/// # Errors
///
/// If there is an error serializing the JSON packet, or if there is an
/// error sending the serialized packet over the transport, this method
/// may return an error. The exact semantics of what constitutes an error
/// will depend on the specific implementation of this trait.
fn send_json_packet(&self, packet: JsonPacket);
}
/// A trait for sending Protocol Buffers packets over some underlying transport.
///
/// This trait provides two methods for sending Protocol Buffers packets:
///
/// * `send_packet` sends a packet in the clear (i.e., not encrypted).
/// * `send_encrypted_packet` sends a packet encrypted using the AES-GCM
/// algorithm with a 128-bit key.
///
/// The exact semantics of what it means to "send a packet" will depend on the
/// specific implementation of this trait. However, in general, the packet will
/// be serialized into a binary message using the Protocol Buffers wire format,
/// and then sent over the underlying transport.
///
/// The `destination` argument specifies which recipient should receive the
/// packet. This is a 1-byte field that is prepended to the serialized packet
/// before it is sent.
///
/// The `key` argument is an optional AES-GCM key. If a key is provided, the
/// packet will be encrypted before being sent. If no key is provided, the
/// packet will be sent in the clear.
///
/// # Errors
///
/// If there is an error serializing the Protocol Buffers packet, or if there
/// is an error sending the serialized packet over the transport, either of
/// these methods may return an error. The exact semantics of what constitutes
/// an error will depend on the specific implementation of this trait.
pub trait PacketSender {
/// Sends a Protocol Buffers packet in the clear.
///
/// The packet will be serialized into a binary message using the Protocol
/// Buffers wire format, and then sent over the underlying transport.
fn send_packet(&self, destination: u8, packet: packets::packet::Value);
/// Sends a Protocol Buffers packet encrypted using AES-GCM.
///
/// The packet will be serialized into a binary message using the Protocol
/// Buffers wire format, encrypted using AES-GCM with a 128-bit key, and
/// then sent over the underlying transport.
///
/// If no key is provided, the packet will be sent in the clear.
fn send_encrypted_packet(
&self,
key: &Option<Aes128Gcm>,
destination: u8,
value: packets::packet::Value,
);
}
impl JsonPacketSender for Sender {
/// Serializes the given JSON packet into a string, and then sends it as a
/// text message over the underlying transport.
///
/// The `JsonPacket` type is defined in the `serde_json` crate, and it is a
/// simple wrapper around a JSON object with string keys and values. This
/// trait method is responsible for taking a `JsonPacket` and sending it
/// over the WebSocket connection.
///
/// The `serde_json::to_string` function is used to serialize the packet
/// into a JSON string. If this function returns an error, we panic
/// because there is no reasonable recovery behavior in this case.
///
/// Once we have the JSON string, we wrap it in a `WebSocketMessage::Text`
/// enum variant and send it over the WebSocket connection using the
/// `send` method. If this method returns an error, we panic because there
/// is no reasonable recovery behavior in this case.
fn send_json_packet(&self, packet: JsonPacket) {
let serialized_packet =
serde_json::to_string(&packet).expect("Failed to serialize JSON packet.");
self.send(WebSocketMessage::Text(serialized_packet))
.expect("Failed to send JSON packet.");
}
}
impl PacketSender for Sender {
/// Serializes the given packet value into a binary message, and then
/// sends it over the underlying transport.
///
/// The `destination` parameter specifies which client should receive
/// this message. The value of this parameter should be a byte that
/// represents the client's index in the list of connected clients.
///
/// The `value` parameter specifies the actual data that should be sent
/// to the client. This will be serialized into a `Packet` struct using
/// the Protocol Buffers wire format.
///
/// This function will first encode the `Packet` struct into a vector of
/// bytes using the Protocol Buffers wire format. It will then insert the
/// `destination` byte as the first element of the vector, so that the
/// receiving client knows which client this message is intended for.
///
/// Finally, this function will send the serialized packet over the
/// underlying transport, which is assumed to be a WebSocket connection.
/// If this send operation fails, this function will panic because there
/// is no reasonable recovery behavior in this case.
fn send_packet(&self, destination: u8, value: packets::packet::Value) {
let packet = Packet { value: Some(value) };
let mut serialized_packet = packet.encode_to_vec();
serialized_packet.insert(0, destination);
self.send(WebSocketMessage::Binary(serialized_packet))
.expect("Failed to send Packet.");
}
/// Similar to `send_packet`, but the message is encrypted using AES-GCM
/// with a 128-bit key.
///
/// If no key is provided (i.e., if `key` is `None`), then the message will
/// be sent in the clear.
///
/// This function works by generating a random 12-byte nonce using the
/// `rand::OsRng` PRNG, encrypting the message using AES-GCM with the
/// provided key and nonce, and then prepending the nonce to the ciphertext
/// before sending it over the WebSocket connection. The receiving client
/// will use the same key and nonce to decrypt the message.
///
/// Note that this function does not actually check whether the provided
/// key is valid. If an invalid key is provided, the encryption will fail
/// and the receiver will not be able to decrypt the message.
fn send_encrypted_packet(
&self,
key: &Option<Aes128Gcm>,
destination: u8,
value: packets::packet::Value,
) {
let packet = Packet { value: Some(value) };
let nonce = Aes128Gcm::generate_nonce(&mut OsRng);
let plaintext = packet.encode_to_vec();
let mut ciphertext = key
.as_ref()
.unwrap()
.encrypt(&nonce, plaintext.as_ref())
.expect("Failed to encrypt Packet.");
let mut serialized_packet = nonce.to_vec();
serialized_packet.append(&mut ciphertext);
serialized_packet.insert(0, destination);
self.send(WebSocketMessage::Binary(serialized_packet))
.expect("Failed to send encrypted Packet.");
}
}
/// A sender is a type that allows us to send messages to a WebSocket client.
///
/// In this case, a sender is a channel that allows us to send WebSocket
/// messages to a client. The messages can be any type that implements the
/// `Into<WebSocketMessage>`.
///
/// The `WebSocketMessage` type represents any message that can be sent over a
/// WebSocket connection. It can be a binary message, a text message, or a
/// close message.
///
/// The `MaybeTlsStream` type is a stream that may or may not be encrypted.
/// If the connection is encrypted (e.g., via TLS), then the stream will be
/// encrypted. If the connection is not encrypted, then the stream will be
/// unencrypted.
///
/// The `TcpStream` type is a stream that is used to connect to a remote
/// server over a TCP connection.
///
/// The `WebSocketStream` type is a stream that is used to connect to a remote
/// WebSocket server. It is a wrapper around the `MaybeTlsStream` stream that
/// adds WebSocket-specific functionality.
pub type Sender = flume::Sender<WebSocketMessage>;
/// A socket is a type that represents a WebSocket connection.
///
/// In this case, a socket is a wrapper around a `MaybeTlsStream` stream that
/// adds WebSocket-specific functionality.
///
/// The `MaybeTlsStream` type is a stream that may or may not be encrypted.
/// If the connection is encrypted (e.g., via TLS), then the stream will be
/// encrypted. If the connection is not encrypted, then the stream will be
/// unencrypted.
///
/// The `TcpStream` type is a stream that is used to connect to a remote
/// server over a TCP connection.
///
/// The `WebSocketStream` type is a stream that is used to connect to a remote
/// WebSocket server. It is a wrapper around the `MaybeTlsStream` stream that
/// adds WebSocket-specific functionality.
pub type Socket = WebSocketStream<MaybeTlsStream<TcpStream>>;

1
src/transfer_info/mod.rs
View file

@ -1 +0,0 @@
pub mod transfer_info;

29
src/transfer_info/transfer_info.rs
View file

@ -1,29 +0,0 @@
use serde::{Deserialize, Serialize};
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct TransferInfoRequest {
pub message: String,
pub body: TransferInfoBody,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct TransferInfoBody {
pub keyword: String,
pub files: String,
pub ip: String,
pub name: String,
}
impl TransferInfoRequest {
pub fn new() -> Self {
Self {
message: "".to_string(),
body: TransferInfoBody {
keyword: "".to_string(),
files: "".to_string(),
ip: "".to_string(),
name: "".to_string(),
},
}
}
}