1use std::io::ErrorKind;
4use std::net::SocketAddr;
5use std::time::Instant;
6
7use msrt::endpoint::{ClientEndpoint, EndpointPoll, EngineConfig, PeerState};
8use tokio::net::{ToSocketAddrs, UdpSocket};
9
10use crate::error::Result;
11use crate::event::UdpClientEvent;
12
13const RX_BYTES: usize = 2048;
14const TX_BYTES: usize = 256;
15
16#[derive(Debug)]
18pub struct UdpClient {
19 socket: UdpSocket,
20 endpoint: ClientEndpoint,
21 start: Instant,
22 rx_buf: [u8; RX_BYTES],
23 tx_buf: [u8; TX_BYTES],
24}
25
26impl UdpClient {
27 pub async fn bind<A, R>(local: A, remote: R) -> Result<Self>
29 where
30 A: ToSocketAddrs,
31 R: ToSocketAddrs,
32 {
33 Self::bind_with_config(local, remote, EngineConfig::default()).await
34 }
35
36 pub async fn bind_with_config<A, R>(local: A, remote: R, config: EngineConfig) -> Result<Self>
38 where
39 A: ToSocketAddrs,
40 R: ToSocketAddrs,
41 {
42 let socket = UdpSocket::bind(local).await?;
43 socket.connect(remote).await?;
44 Ok(Self::from_socket(socket, config))
45 }
46
47 pub fn from_socket(socket: UdpSocket, config: EngineConfig) -> Self {
49 Self {
50 socket,
51 endpoint: ClientEndpoint::new(config),
52 start: Instant::now(),
53 rx_buf: [0; RX_BYTES],
54 tx_buf: [0; TX_BYTES],
55 }
56 }
57
58 pub fn local_addr(&self) -> Result<SocketAddr> {
60 Ok(self.socket.local_addr()?)
61 }
62
63 pub fn peer_addr(&self) -> Result<SocketAddr> {
65 Ok(self.socket.peer_addr()?)
66 }
67
68 pub fn peer_state(&self) -> PeerState {
70 self.endpoint.peer().state()
71 }
72
73 pub const fn socket(&self) -> &UdpSocket {
75 &self.socket
76 }
77
78 pub fn socket_mut(&mut self) -> &mut UdpSocket {
80 &mut self.socket
81 }
82
83 pub fn into_socket(self) -> UdpSocket {
85 self.socket
86 }
87
88 pub fn connect(&mut self) -> Result<()> {
90 self.endpoint.connect(self.now_ms())?;
91 Ok(())
92 }
93
94 pub fn disconnect(&mut self) {
96 self.endpoint.disconnect();
97 }
98
99 pub fn send(&mut self, message: &[u8]) -> Result<bool> {
101 Ok(self.endpoint.send(message)?.is_some())
102 }
103
104 pub fn receive_available(&mut self) -> Result<usize> {
106 let mut datagrams = 0;
107 loop {
108 match self.socket.try_recv(&mut self.rx_buf) {
109 Ok(n) => {
110 datagrams += 1;
111 let now_ms = self.now_ms();
112 let _ = self.endpoint.receive(now_ms, &self.rx_buf[..n]);
113 }
114 Err(error) if error.kind() == ErrorKind::WouldBlock => return Ok(datagrams),
115 Err(error) if error.kind() == ErrorKind::Interrupted => continue,
116 Err(error) => return Err(error.into()),
117 }
118 }
119 }
120
121 pub async fn poll(&mut self) -> Result<UdpClientEvent> {
123 let now_ms = self.now_ms();
124 loop {
125 match self.endpoint.poll(now_ms, &mut self.tx_buf)? {
126 EndpointPoll::Transmit { bytes, .. } => {
127 let _ = self.socket.send(bytes).await?;
128 }
129 EndpointPoll::Message(message) => return Ok(UdpClientEvent::Message(message)),
130 EndpointPoll::SendFailed(failed) => return Ok(UdpClientEvent::SendFailed(failed)),
131 EndpointPoll::Idle => return Ok(UdpClientEvent::Idle),
132 }
133 }
134 }
135
136 pub async fn tick(&mut self) -> Result<UdpClientEvent> {
138 if let Err(error) = self.receive_available() {
139 if let Some(kind) = recoverable_transport_error(&error) {
140 return Ok(UdpClientEvent::TransportUnavailable { kind });
141 }
142 return Err(error);
143 }
144
145 match self.poll().await {
146 Ok(event) => Ok(event),
147 Err(error) => {
148 if let Some(kind) = recoverable_transport_error(&error) {
149 return Ok(UdpClientEvent::TransportUnavailable { kind });
150 }
151 Err(error)
152 }
153 }
154 }
155
156 fn now_ms(&self) -> u64 {
157 self.start
158 .elapsed()
159 .as_millis()
160 .try_into()
161 .unwrap_or(u64::MAX)
162 }
163}
164
165fn recoverable_transport_error(error: &crate::Error) -> Option<ErrorKind> {
166 let crate::Error::Io(error) = error else {
167 return None;
168 };
169
170 match error.kind() {
171 ErrorKind::ConnectionRefused
172 | ErrorKind::ConnectionReset
173 | ErrorKind::ConnectionAborted
174 | ErrorKind::NotConnected
175 | ErrorKind::TimedOut => Some(error.kind()),
176 _ => None,
177 }
178}