Implement connection process

This commit is contained in:
Sofia 2026-07-16 01:06:58 +03:00
parent 0e2166fea4
commit d8781b55d0
5 changed files with 507 additions and 5 deletions

View File

@ -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

335
src/connections.rs Normal file
View File

@ -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<std::io::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<SocketAddr, Connection>,
closing_since: Option<Instant>,
}
impl ConnectionManager {
pub fn new(socket: Arc<UdpSocket>) -> 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<PeerMessage> {
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<PeerMessage> {
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<Arc<UdpSocket>>,
}
impl UdpWrapper {
pub fn new(socket: Arc<UdpSocket>) -> 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<ConnectionError>,
}
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,
}

View File

@ -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<ConnectionError>),
}
pub struct Peer {
socket: UdpSocket,
connection_mgr: ConnectionManager,
closing: Arc<AtomicBool>,
receiver: Receiver<ListenerMessage>,
messages: VecDeque<PeerMessage>,
}
impl Peer {
/// Bind given port to listen to incoming messages. Creates a new Peer that
/// can establish new connections.
pub fn listen(port: Option<u16>) -> Result<Peer, PeerError> {
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<Option<PeerMessage>, 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();
}
}

57
src/listener.rs Normal file
View File

@ -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<UdpSocket>,
sender: Sender<ListenerMessage>,
}
impl Listener {
pub fn new(socket: Arc<UdpSocket>, sender: Sender<ListenerMessage>) -> 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::<Package, _>(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),
}

9
src/package.rs Normal file
View File

@ -0,0 +1,9 @@
use serde::{Deserialize, Serialize};
#[derive(Debug, Serialize, Deserialize)]
pub enum Package {
Hello,
Ping,
Pong,
Close,
}