1use std::io::ErrorKind;
4use std::net::SocketAddr;
5use std::time::Instant;
6
7use msrt::endpoint::{EndpointPoll, EngineConfig, PeerState, ServerEndpoint};
8use tokio::net::{ToSocketAddrs, UdpSocket};
9
10use crate::error::{Error, Result};
11use crate::event::UdpServerEvent;
12
13const RX_BYTES: usize = 2048;
14const TX_BYTES: usize = 256;
15
16#[derive(Debug)]
18pub struct UdpServer<const N: usize> {
19 socket: UdpSocket,
20 endpoint: ServerEndpoint<SocketAddr, N>,
21 start: Instant,
22 rx_buf: [u8; RX_BYTES],
23 tx_buf: [u8; TX_BYTES],
24}
25
26impl<const N: usize> UdpServer<N> {
27 pub async fn bind<A>(local: A) -> Result<Self>
29 where
30 A: ToSocketAddrs,
31 {
32 Self::bind_with_config(local, EngineConfig::default()).await
33 }
34
35 pub async fn bind_with_config<A>(local: A, config: EngineConfig) -> Result<Self>
37 where
38 A: ToSocketAddrs,
39 {
40 let socket = UdpSocket::bind(local).await?;
41 Ok(Self::from_socket(socket, config))
42 }
43
44 pub fn from_socket(socket: UdpSocket, config: EngineConfig) -> Self {
46 Self {
47 socket,
48 endpoint: ServerEndpoint::new(config),
49 start: Instant::now(),
50 rx_buf: [0; RX_BYTES],
51 tx_buf: [0; TX_BYTES],
52 }
53 }
54
55 pub fn local_addr(&self) -> Result<SocketAddr> {
57 Ok(self.socket.local_addr()?)
58 }
59
60 pub const fn socket(&self) -> &UdpSocket {
62 &self.socket
63 }
64
65 pub fn socket_mut(&mut self) -> &mut UdpSocket {
67 &mut self.socket
68 }
69
70 pub fn into_socket(self) -> UdpSocket {
72 self.socket
73 }
74
75 pub fn peer_state(&mut self, peer: SocketAddr) -> Option<PeerState> {
77 self.endpoint.peer_mut(peer).map(|slot| slot.state())
78 }
79
80 pub fn peers(&self) -> impl Iterator<Item = SocketAddr> + '_ {
82 self.endpoint.peers().map(|peer| *peer.peer_id())
83 }
84
85 pub fn send_to(&mut self, peer: SocketAddr, message: &[u8]) -> Result<bool> {
87 match self.endpoint.send(peer, message) {
88 Some(Ok(Some(_))) => Ok(true),
89 Some(Ok(None)) | None => Ok(false),
90 Some(Err(error)) => Err(error.into()),
91 }
92 }
93
94 pub fn disconnect(&mut self, peer: SocketAddr) -> bool {
96 self.endpoint.disconnect(peer)
97 }
98
99 pub fn disconnect_idle(&mut self, timeout_ms: u64) -> usize {
101 self.endpoint.disconnect_idle(self.now_ms(), timeout_ms)
102 }
103
104 pub fn receive_available(&mut self) -> Result<usize> {
109 let mut datagrams = 0;
110 loop {
111 match self.socket.try_recv_from(&mut self.rx_buf) {
112 Ok((n, peer)) => {
113 datagrams += 1;
114 let now_ms = self.now_ms();
115 if self.endpoint.peer_mut(peer).is_none() {
116 self.endpoint.accept(peer, now_ms)?;
117 }
118 let _ = self.endpoint.receive(peer, now_ms, &self.rx_buf[..n]);
119 }
120 Err(error) if error.kind() == ErrorKind::WouldBlock => return Ok(datagrams),
121 Err(error) if error.kind() == ErrorKind::Interrupted => continue,
122 Err(error) => return Err(error.into()),
123 }
124 }
125 }
126
127 pub async fn poll(&mut self) -> Result<UdpServerEvent> {
129 let now_ms = self.now_ms();
130 let peers: Vec<SocketAddr> = self.peers().collect();
131
132 for peer in peers {
133 loop {
134 let poll = self
135 .endpoint
136 .poll(peer, now_ms, &mut self.tx_buf)
137 .ok_or_else(|| Error::Io(ErrorKind::NotFound.into()))??;
138
139 match poll {
140 EndpointPoll::Transmit { bytes, .. } => {
141 let _ = self.socket.send_to(bytes, peer).await?;
142 }
143 EndpointPoll::Message(message) => {
144 return Ok(UdpServerEvent::Message { peer, message });
145 }
146 EndpointPoll::SendFailed(failed) => {
147 return Ok(UdpServerEvent::SendFailed { peer, failed });
148 }
149 EndpointPoll::Idle => break,
150 }
151 }
152 }
153
154 Ok(UdpServerEvent::Idle)
155 }
156
157 pub async fn tick(&mut self) -> Result<UdpServerEvent> {
159 let _ = self.receive_available()?;
160 self.poll().await
161 }
162
163 fn now_ms(&self) -> u64 {
164 self.start
165 .elapsed()
166 .as_millis()
167 .try_into()
168 .unwrap_or(u64::MAX)
169 }
170}