Implement connection process
This commit is contained in:
parent
0e2166fea4
commit
d8781b55d0
@ -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
335
src/connections.rs
Normal 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,
|
||||
}
|
||||
109
src/lib.rs
109
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<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
57
src/listener.rs
Normal 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
9
src/package.rs
Normal file
@ -0,0 +1,9 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub enum Package {
|
||||
Hello,
|
||||
Ping,
|
||||
Pong,
|
||||
Close,
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user