I am new to Rust, so I created a mini library/wrapper for creating multi-cliented Tcp Servers. This is a starter project for me to learn Rust.
How the module works is a user can pass in a function or closure to the start fn to use to handle each new client's connection, the function has to accept HandleClientType as a parameter, which is a type alias for (TcpStream, BubbleServer), so a tuple that contains the newly connected TcpStream and a reference to the server that is calling this function.
I'd like to know how I could improve it's performance and improving it's overall design (some details in the comments) and if there is anyways to simplify my code.
If you wanted to compile this for yourself, a full crate for the module and the main example (also a crate for the TCP client) is provided here
The Module in question:
buble_host/mod.rs:
use std::io;
use std::io::{Error, ErrorKind};
use std::net::{Shutdown, SocketAddr, TcpListener, TcpStream};
use std::sync::mpsc::{channel, Sender};
use std::sync::*;
use std::thread;
use std::vec::Vec;
#[cfg(feature = "bubble_dbg")]
macro_rules! bubble_print{
($($args:expr), *)=> {println!($($args), *);}
}
#[cfg(not(feature = "bubble_dbg"))]
macro_rules! bubble_print {
($($args:expr), *) => {};
}
//Events for debugging and getting information from the server.
pub enum ServerEvent {
None,
Disconnection(SocketAddr),
}
type ErrorSender = Option<Sender<ServerEvent>>;
//Clients List Thread Safe
pub type ClientsListTS = Arc<Mutex<Vec<TcpStream>>>;
pub type ClientsVec = Vec<TcpStream>;
pub type HandleClientCallBack = dyn Fn(HandleClientType) + Send + Sync + 'static;
//Helper struct
struct ClientsList(ClientsListTS);
impl ClientsList {
pub fn new() -> ClientsListTS {
Arc::new(Mutex::new(Vec::<TcpStream>::new()))
}
}
///Data that is sent to user provided FnMut `set in fn start` when a client connects to a started(fn start) server
pub type HandleClientType<'a> = (&'a mut TcpStream, &'a mut BubbleServer);
#[derive(Clone)]
pub struct BubbleServer {
ip: String,
error_sender: ErrorSender,
clients: ClientsListTS,
}
impl BubbleServer {
pub fn new(ip: String) -> Self {
BubbleServer {
ip,
clients: ClientsList::new(),
error_sender: None,
}
}
fn get_sock_index(
clients: &ClientsVec,
socket_addr: &std::net::SocketAddr,
) -> io::Result<usize> {
clients
.iter()
.position(|x| &x.peer_addr().unwrap() == socket_addr)
.ok_or_else(|| {
Error::new(
ErrorKind::InvalidInput,
format!("Could not find socket address: '{}' ", socket_addr),
)
})
}
/// used to retrieve socket from clients by ip address
/// # Example :
/// ```
/// let server = BubbleServer::new(String::from("localhost:25568"));
/// let addr_to_find: SocketAddr = "127.0.0.1:25565"
/// .parse()
/// .expect("Could not parse ip address!");
/// let socket = server.get_sock(&addr_to_find);
/// assert!(socket.is_none());
///
/// ```
///
#[allow(dead_code)]
pub fn get_sock(&self, socket_addr: &SocketAddr) -> Option<TcpStream> {
let clients = self.get_clients().ok()?;
let clients = clients.lock().ok()?;
if let Ok(index) = BubbleServer::get_sock_index(&clients, socket_addr) {
Some(clients[index].try_clone().unwrap())
} else {
None
}
}
///`handle_socket_disconnection` will shutdown the socket, then send a disconnection event to the error_sender.
fn handle_socket_disconnection(&self, socket_: &TcpStream) -> io::Result<()> {
self.remove_socket(&socket_)?;
let sock_addr = socket_.peer_addr().unwrap();
socket_
.shutdown(Shutdown::Both)
.expect("Could not shutdown stream..");
let event = ServerEvent::Disconnection(sock_addr);
let sender = self.error_sender.as_ref();
if let Some(sender) = sender {
sender
.send(event)
.unwrap_or_else(|e| println!("error sending to error_sender {}", e));
}
bubble_print!("Removed client {}", &sock_addr.to_string());
Ok(())
}
/// `get_clients` Returns clients list reference using Arc;
pub fn get_clients(&self) -> io::Result<ClientsListTS> {
Ok(Arc::clone(&self.clients))
}
#[allow(dead_code)]
pub fn set_on_event<F: Fn(ServerEvent) + Send + Sync + 'static>(&self, callback: F) {
let (err_send, err_rec) = channel::<ServerEvent>();
self.error_sender.as_ref().get_or_insert(&err_send);
let call = Arc::new(callback);
std::thread::spawn(move || loop {
let event = err_rec.recv().unwrap_or(ServerEvent::None);
call(event);
});
}
///
///####`remove_socket` Removes given TcpStream from server's `clients`
///
/// * Locks `clients` vec (Vec<TcpStream>)
/// * Removes `socket` from clients list.
fn remove_socket(&self, socket: &TcpStream) -> io::Result<()> {
//Is this the least verbose way to dereference this?
let clients = self.get_clients()?;
let mut clients = clients
.lock()
.map_err(|err| io::Error::new(ErrorKind::Other, err.to_string()))?;
//
let socket_addr = socket.peer_addr()?;
let index = BubbleServer::get_sock_index(&clients, &socket_addr)?;
clients.remove(index);
Ok(())
}
///Runs the user's defined function in a new thread passing in the newly connected socket.
fn handle_client(
&self,
socket: &mut TcpStream,
handle_client: &'static HandleClientCallBack,
) -> io::Result<()> {
let mut socket_ = socket.try_clone()?;
//Again, I hate this cloning
let mut _self = self.clone();
//Need to clone provided function
let handle_client_func = handle_client.clone();
thread::spawn(move || {
handle_client_func((&mut socket_, &mut _self));
_self
.handle_socket_disconnection(&socket_)
.unwrap_or_else(|e| {
println!("Error in handling socket disconnection, Err: '{}' ", e);
});
});
Ok(())
}
///Locks `clients` (Vec<TcpStream>) and adds `socket` to Vec
fn add_client(&self, socket: TcpStream) {
self.clients.lock().unwrap().push(socket);
}
/// Continously accepts incoming connections
/// * Is non blocking.
/// * Locks clients list when user connects
pub fn start(&self, handle_client_cb: &'static HandleClientCallBack) -> io::Result<()> {
let ip = &self.ip;
bubble_print!("server with ip of {} has started...", ip);
let socket_acceptor = TcpListener::bind(ip).expect("Failed to initialize server...");
//I'm not a fan of cloning the server struct, but I didn't want the user to have the potential to deadlock the server
//without explicitily locking mutex's in their own code.
//for example it should be obvious that this will deadlock when a user connects
// ```
// {
// let server = BubbleServer::new(String::from("localhost:25568"));
// let clients = server.get_clients().expect("could not get clients!");
// let clients = clients.lock().unwrap();
// server.start(&|data: HandleClientType| {
// println!("new connection! dropping now!");
// });
// thread::park();
// }
// ```
let self_ = self.clone();
// Accept new incoming connections
thread::spawn(move || loop {
let h_cb_clone = handle_client_cb.clone();
println!("waiting for next client..");
if let Ok((mut socket, _)) = socket_acceptor.accept() {
//Add Client to the clients vector
self_.add_client(
socket
.try_clone()
.expect("Could not clone socket before handling socket.."),
);
//Run user's implementation on how to deal with new client
self_
.handle_client(&mut socket, h_cb_clone)
.unwrap_or_else(|e| {
println!("Error in handle_client : {}", e);
});
}
// Could provide implementation for an erroneous accept, and send that to the error_sender.
// but if many clients were to make erroneous connections this might just clutter the error_sender's receive.
});
Ok(())
}
}
main.rs shows an example of how I create a TextChat server with this module
main.rs:
mod bubble_host;
use crate::bubble_host::{BubbleServer, HandleClientType, ServerEvent};
use std::io;
use std::io::prelude::*;
use std::io::{Error, ErrorKind};
use std::net::TcpStream;
use std::thread;
///Example implementation of BubbleServer
trait TextChatServer {
fn broad_cast_msg(&mut self, msg: &[u8]) -> Result<(), std::io::Error>;
}
impl TextChatServer for BubbleServer {
fn broad_cast_msg(&mut self, msg: &[u8]) -> Result<(), std::io::Error> {
let clients = self.get_clients()?;
let clients = clients
.lock()
.map_err(|e| Error::new(ErrorKind::Other, e.to_string()))?;
for mut client in clients.iter() {
if let Err(e) = client.write_all(&msg) {
println!(
"Error {} for {}",
e,
client.peer_addr().unwrap().to_string()
)
}
}
Ok(())
}
}
fn read_stream(stream: &mut TcpStream, recv_buf: &mut [u8], read_bytes: &mut usize) -> usize {
*read_bytes = match stream.read(recv_buf) {
Ok(read_bytes) => read_bytes,
Err(_) => 0,
};
*read_bytes
}
trait TcpHelper {
fn send_str(&mut self, msg: &str) -> io::Result<()>;
}
impl TcpHelper for TcpStream {
fn send_str(&mut self, msg: &str) -> io::Result<()> {
let data = msg.as_bytes();
self.write_all(&data)?;
Ok(())
}
}
//This Function will be called and ran in a new thread every time a new user connects.
fn handle_new_text_client(client_data: HandleClientType) {
// HandleClientType provides the stream that is being handled and the server it's being called from so the user could access functions from the server
let buffer_size = 1000;
let (stream, server_reference) = client_data;
println!(
"Client connected!! ip {}",
stream.peer_addr().unwrap().to_string()
);
stream
.send_str(&stream.peer_addr().unwrap().to_string())
.unwrap_or_else(|err| {
println!(
"err writing ip to stream, ip {}, err :{}",
&stream.peer_addr().unwrap().to_string(),
err
)
});
let mut recv_buf = vec![0u8; buffer_size];
let mut bytes_read: usize = 0;
while read_stream(stream, &mut recv_buf, &mut bytes_read) > 0 {
if let Err(e) = server_reference.broad_cast_msg(&recv_buf[..bytes_read]) {
println!("Error: {}", e);
};
}
}
fn main() {
let server = BubbleServer::new(String::from("localhost:25568"));
//Receive events from Server
server.set_on_event(|event: ServerEvent| match event {
ServerEvent::Disconnection(socket_addr) => println!("Disconnection from {}", socket_addr),
_ => (),
});
if let Err(e) = server.start(&handle_new_text_client) {
println!("Error starting server: {}", e);
};
//Blocks this thread here as server.start runs in another thread
thread::park();
}
