diff --git a/Cargo.toml b/Cargo.toml index 66570da..503eaae 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,5 +4,7 @@ version = "0.1.0" edition = "2024" [dependencies] +ciborium = "0.2.2" +serde = "1.0.228" thiserror = "2.0.18" ## Make it easier to generate errors diff --git a/src/connections.rs b/src/connections.rs new file mode 100644 index 0000000..94c7748 --- /dev/null +++ b/src/connections.rs @@ -0,0 +1,335 @@ +use std::{ + collections::HashMap, + io::Cursor, + net::{SocketAddr, UdpSocket}, + sync::Arc, + time::{Duration, Instant}, +}; + +use thiserror::Error; + +use crate::{ + PeerMessage, + listener::{DATAGRAM_SIZE, Datagram, ListenerError}, + package::Package, +}; + +#[derive(Debug, Error)] +pub enum SendError { + #[error("Serialization Error: {0}")] + SerializationError(ciborium::ser::Error), + #[error("Datagram is too large: {0}")] + DatagramTooLarge(usize), + #[error("Socket disconnected")] + SocketDisconnected, + #[error("Error while sending datagram: {0}")] + SendError(std::io::Error), +} + +#[derive(Debug, Error)] +pub enum ConnectionError { + #[error(transparent)] + SendError(#[from] SendError), + #[error(transparent)] + ListenerError(#[from] ListenerError), + #[error("Connection timed out")] + TimedOut, +} + +pub struct ConnectionManager { + udp: UdpWrapper, + connections: HashMap, + closing_since: Option, +} + +impl ConnectionManager { + pub fn new(socket: Arc) -> ConnectionManager { + ConnectionManager { + udp: UdpWrapper::new(socket), + connections: HashMap::new(), + closing_since: None, + } + } + + pub fn close(&mut self) { + self.closing_since = Some(Instant::now()); + for (addr, _) in self.connections.clone() { + self.close_connection(&addr); + } + } + + pub fn connect_to(&mut self, addr: SocketAddr) { + self.connections + .insert(addr, Connection::from(addr, ConnectionState::Connecting)); + } + + pub fn send_pings(&mut self) { + let now = Instant::now(); + + if let Some(closing_since) = self.closing_since { + if self.connections.len() == 0 { + self.udp.close(); + } else if now - closing_since > Duration::from_millis(1000) { + self.udp.close(); + } + } + + if let Some(_) = &self.udp.socket { + for (addr, conn) in &mut self.connections { + let duration = now - conn.last_sent_ping; + if duration > Duration::from_millis(100) { + match conn.state { + ConnectionState::ReceivingConnection | ConnectionState::Connecting => { + match self.udp.send_to(*addr, Package::Hello) { + Ok(_) => { + conn.last_sent_ping = now; + } + Err(err) => { + conn.closing_since = Instant::now(); + conn.state = ConnectionState::Error; + conn.error = Some(ConnectionError::SendError(err)); + } + } + } + ConnectionState::Connected | ConnectionState::ConnectingNearlyReady => { + match self.udp.send_to(*addr, Package::Ping) { + Ok(_) => { + conn.last_sent_ping = now; + } + Err(err) => { + conn.closing_since = Instant::now(); + conn.state = ConnectionState::Error; + conn.error = Some(ConnectionError::SendError(err)); + } + } + } + ConnectionState::Closing => match self.udp.send_to(*addr, Package::Close) { + Ok(_) => { + conn.last_sent_ping = now; + } + Err(err) => { + conn.closing_since = Instant::now(); + conn.state = ConnectionState::Error; + conn.error = Some(ConnectionError::SendError(err)); + } + }, + _ => {} + } + } + } + } + } + + pub fn close_connection(&mut self, addr: &SocketAddr) { + if let Some(conn) = self.connections.get_mut(addr) { + if conn.state != ConnectionState::Closing { + conn.closing_since = Instant::now(); + conn.state = ConnectionState::Closing; + } + } + } + + pub fn error_connection(&mut self, addr: &SocketAddr, error: ConnectionError) { + if let Some(conn) = self.connections.get_mut(addr) { + if conn.state != ConnectionState::Error { + conn.closing_since = Instant::now(); + conn.state = ConnectionState::Error; + conn.error = Some(error); + } + } + } + + #[must_use] + pub fn handle_package(&mut self, package: Package, addr: &SocketAddr) -> Vec { + let mut messages = Vec::new(); + + match package { + Package::Hello => { + if let Some(conn) = self.connections.get_mut(addr) { + if conn.state == ConnectionState::Connecting { + conn.state = ConnectionState::ConnectingNearlyReady; + } + } else { + self.connections.insert( + *addr, + Connection::from(*addr, ConnectionState::ReceivingConnection), + ); + } + } + Package::Ping => { + if let Some(conn) = self.connections.get_mut(addr) { + if conn.state == ConnectionState::ConnectingNearlyReady + || conn.state == ConnectionState::ReceivingConnection + { + conn.state = ConnectionState::Connected; + messages.push(PeerMessage::NewConnection(conn.clone())); + } + + if conn.state == ConnectionState::Connected { + conn.last_recv_ping = Instant::now(); + match self.udp.send_to(*addr, Package::Pong) { + Ok(_) => {} + Err(err) => { + conn.closing_since = Instant::now(); + conn.state = ConnectionState::Error; + conn.error = Some(ConnectionError::SendError(err)); + } + } + } + } + } + Package::Pong => { + if let Some(conn) = self.connections.get_mut(addr) { + if conn.state == ConnectionState::Connected { + conn.ping = Instant::now() - conn.last_sent_ping; + } + } + } + Package::Close => { + let remove = if let Some(conn) = self.connections.get_mut(addr) { + match conn.state { + ConnectionState::Closing => true, + ConnectionState::Error => false, + ConnectionState::ReceivingClosing => { + conn.last_recv_close = Instant::now(); + false + } + _ => { + conn.state = ConnectionState::ReceivingClosing; + false + } + } + } else { + false + }; + if remove { + self.connections.remove(addr); + } + } + } + + messages + } + + pub fn close_silent_connections(&mut self) { + let now = Instant::now(); + for (addr, conn) in self.connections.clone() { + if (now - conn.last_recv_ping) > Duration::from_millis(2000) { + self.error_connection(&addr, ConnectionError::TimedOut); + } + } + } + + pub fn purge_old_connections(&mut self) -> Vec { + let mut messages = Vec::new(); + + let now = Instant::now(); + for (addr, conn) in self.connections.clone() { + match conn.state { + ConnectionState::ReceivingClosing => { + if (now - conn.last_recv_close) > Duration::from_millis(500) { + self.connections.remove(&addr); + messages.push(PeerMessage::Disconnected(conn, None)); + } + } + ConnectionState::Error => { + self.connections.remove(&addr); + messages.push(PeerMessage::Disconnected(conn.clone(), conn.error)); + } + _ => {} + } + } + + messages + } +} + +pub struct UdpWrapper { + socket: Option>, +} + +impl UdpWrapper { + pub fn new(socket: Arc) -> UdpWrapper { + UdpWrapper { + socket: Some(socket), + } + } + + pub fn close(&mut self) { + self.socket = None; + } + + fn send_to(&self, addr: SocketAddr, package: Package) -> Result<(), SendError> { + if let Some(socket) = &self.socket { + let mut buf = Vec::new(); + let cursor = Cursor::new(&mut buf); + ciborium::into_writer(&package, cursor) + .map_err(|err| SendError::SerializationError(err))?; + if buf.len() > DATAGRAM_SIZE { + return Err(SendError::DatagramTooLarge(buf.len())); + } + socket + .send_to(&buf, addr) + .map_err(|e| SendError::SendError(e)) + .map(|_| ()) + } else { + Err(SendError::SocketDisconnected) + } + } +} + +pub struct Connection { + pub address: SocketAddr, + pub ping: Duration, + + last_sent_ping: Instant, + last_recv_ping: Instant, + last_recv_close: Instant, + closing_since: Instant, + + pub state: ConnectionState, + error: Option, +} + +impl Clone for Connection { + fn clone(&self) -> Self { + Self { + address: self.address.clone(), + ping: self.ping.clone(), + last_sent_ping: self.last_sent_ping.clone(), + last_recv_ping: self.last_recv_ping.clone(), + last_recv_close: self.last_recv_close.clone(), + closing_since: self.closing_since.clone(), + state: self.state.clone(), + error: None, + } + } +} + +impl Connection { + pub fn from(address: SocketAddr, state: ConnectionState) -> Connection { + Connection { + address, + ping: Duration::default(), + last_sent_ping: Instant::now() - Duration::from_hours(1), + last_recv_ping: Instant::now(), + last_recv_close: Instant::now(), + closing_since: Instant::now(), + state, + error: None, + } + } +} + +#[derive(PartialEq, Eq, PartialOrd, Ord, Clone, Copy)] +pub enum ConnectionState { + Connecting, + ConnectingNearlyReady, + ReceivingConnection, + Connected, + + Closing, + ReceivingClosing, + + Error, +} diff --git a/src/lib.rs b/src/lib.rs index d7f1e62..5264177 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,23 +1,122 @@ -use std::net::{SocketAddr, UdpSocket}; +use std::{ + collections::VecDeque, + net::{SocketAddr, UdpSocket}, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + mpsc::{Receiver, channel}, + }, + thread::spawn, +}; use thiserror::*; +use crate::{ + connections::{Connection, ConnectionError, ConnectionManager}, + listener::{Listener, ListenerMessage}, +}; + +pub mod connections; +pub(crate) mod listener; +pub(crate) mod package; + #[derive(Error, Debug)] pub enum PeerError { #[error("Error binding to port: {0}")] BindError(std::io::Error), + #[error("Error listening to messages: {0}")] + ListenerError(std::io::Error), +} + +pub enum PeerMessage { + NewConnection(Connection), + Disconnected(Connection, Option), } pub struct Peer { - socket: UdpSocket, + connection_mgr: ConnectionManager, + closing: Arc, + receiver: Receiver, + messages: VecDeque, } impl Peer { /// Bind given port to listen to incoming messages. Creates a new Peer that /// can establish new connections. pub fn listen(port: Option) -> Result { - let socket = UdpSocket::bind(SocketAddr::from(([0, 0, 0, 0], port.unwrap_or(0)))) - .map_err(|e| PeerError::BindError(e))?; - Ok(Peer { socket }) + let socket = Arc::new( + UdpSocket::bind(SocketAddr::from(([0, 0, 0, 0], port.unwrap_or(0)))) + .map_err(|e| PeerError::BindError(e))?, + ); + let closing = Arc::new(AtomicBool::new(false)); + + let (sender, receiver) = channel(); + + spawn({ + let socket = socket.clone(); + let closing = closing.clone(); + move || { + let mut listener = Listener::new(socket, sender); + + while !closing.load(Ordering::Relaxed) { + listener.poll(); + } + } + }); + + Ok(Peer { + connection_mgr: ConnectionManager::new(socket), + closing, + receiver, + messages: VecDeque::new(), + }) + } + + pub fn connect_to(&mut self, addr: SocketAddr) { + self.connection_mgr.connect_to(addr); + } + + #[must_use] + pub fn poll(&mut self) -> Result, PeerError> { + self.connection_mgr.send_pings(); + + let mut msg; + while { + msg = match self.receiver.try_recv() { + Ok(msg) => Some(msg), + Err(_) => None, + }; + + msg.is_some() + } { + match msg.unwrap() { + ListenerMessage::Package(package, socket_addr) => { + self.messages + .extend(self.connection_mgr.handle_package(package, &socket_addr)); + } + ListenerMessage::PackageError(listener_error, socket_addr) => { + self.connection_mgr.error_connection( + &socket_addr, + ConnectionError::ListenerError(listener_error), + ); + } + ListenerMessage::Error(error) => { + self.close(); + return Err(PeerError::ListenerError(error)); + } + } + } + + self.connection_mgr.close_silent_connections(); + + self.messages + .extend(self.connection_mgr.purge_old_connections()); + + Ok(self.messages.pop_front()) + } + + pub fn close(&mut self) { + self.closing.store(true, Ordering::Relaxed); + self.connection_mgr.close(); } } diff --git a/src/listener.rs b/src/listener.rs new file mode 100644 index 0000000..8f73023 --- /dev/null +++ b/src/listener.rs @@ -0,0 +1,57 @@ +use std::{ + io::Cursor, + net::{SocketAddr, UdpSocket}, + sync::{Arc, mpsc::Sender}, +}; + +use thiserror::Error; + +use crate::package::Package; + +pub const DATAGRAM_SIZE: usize = 102400; +pub type Datagram = [u8; DATAGRAM_SIZE]; + +pub struct Listener { + socket: Arc, + sender: Sender, +} + +impl Listener { + pub fn new(socket: Arc, sender: Sender) -> Listener { + Listener { socket, sender } + } + + pub fn poll(&mut self) { + let mut buffer: Datagram = [0; _]; + match self.socket.recv_from(&mut buffer) { + Ok((num_bytes, from_addr)) => { + let mut ciborium_buf: Datagram = [0; _]; + let bytes = Cursor::new(&mut buffer[..num_bytes]); + let res = ciborium::from_reader_with_buffer::(bytes, &mut ciborium_buf); + match res { + Ok(pkg) => self.sender.send(ListenerMessage::Package(pkg, from_addr)), + Err(_) => self.sender.send(ListenerMessage::PackageError( + ListenerError::ParseError, + from_addr, + )), + } + .ok(); + } + Err(err) => { + self.sender.send(ListenerMessage::Error(err)).ok(); + } + } + } +} + +#[derive(Debug, Error)] +pub enum ListenerError { + #[error("Error parsing datagram")] + ParseError, +} + +pub enum ListenerMessage { + Package(Package, SocketAddr), + PackageError(ListenerError, SocketAddr), + Error(std::io::Error), +} diff --git a/src/package.rs b/src/package.rs new file mode 100644 index 0000000..258e82d --- /dev/null +++ b/src/package.rs @@ -0,0 +1,9 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Serialize, Deserialize)] +pub enum Package { + Hello, + Ping, + Pong, + Close, +}