1use crate::link_handshake::{LinkHandshakeError, LinkTransport};
7use std::collections::HashMap;
8use std::io;
9use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, UdpSocket};
10use std::sync::atomic::{AtomicBool, Ordering};
11use std::sync::{Arc, Mutex, mpsc};
12use std::thread::{self, JoinHandle};
13use std::time::Duration;
14pub trait PunchSocket: Send + Sync {
17 fn send_to(&self, bytes: &[u8], target: SocketAddr) -> io::Result<usize>;
20 fn recv_from(&self, bytes: &mut [u8]) -> io::Result<(usize, SocketAddr)>;
23 fn local_addr(&self) -> io::Result<SocketAddr>;
26 fn prepare(&self, bound: Option<Duration>) -> io::Result<()>;
30 fn wake(&self) -> io::Result<()> {
34 let target = match self.local_addr()? {
35 SocketAddr::V4(address) if address.ip().is_unspecified() => {
36 SocketAddr::from((Ipv4Addr::LOCALHOST, address.port()))
37 }
38 SocketAddr::V6(address) if address.ip().is_unspecified() => {
39 SocketAddr::from((Ipv6Addr::LOCALHOST, address.port()))
40 }
41 address => address,
42 };
43 self.send_to(&[], target)?;
44 Ok(())
45 }
46}
47impl PunchSocket for UdpSocket {
48 fn send_to(&self, bytes: &[u8], target: SocketAddr) -> io::Result<usize> {
49 Self::send_to(self, bytes, target)
50 }
51 fn recv_from(&self, bytes: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
52 Self::recv_from(self, bytes)
53 }
54 fn local_addr(&self) -> io::Result<SocketAddr> {
55 Self::local_addr(self)
56 }
57 fn prepare(&self, bound: Option<Duration>) -> io::Result<()> {
58 match self.peer_addr() {
59 Ok(_) => return Err(io::Error::other("udp_socket_must_be_unconnected")),
60 Err(error) if error.kind() == io::ErrorKind::NotConnected => {}
61 Err(error) => return Err(error),
62 }
63 self.set_nonblocking(false)?;
64 self.set_read_timeout(bound)
65 }
66}
67pub const MAX_DATAGRAM: usize = 1200;
68const HEADER: usize = 13;
69const TAG: usize = 16;
70const WINDOW: usize = 2048;
71#[derive(Debug, Clone, PartialEq, Eq)]
73pub struct Datagram {
74 pub kind: u8,
76 pub receiver_index: u32,
78 pub counter: u64,
80 pub payload: Vec<u8>,
81}
82impl Datagram {
83 pub fn encode(&self) -> Result<Vec<u8>, UdpError> {
86 if self.payload.len() > MAX_DATAGRAM - HEADER {
87 return Err(UdpError::TooLarge);
88 }
89 if !(1..=4).contains(&self.kind) {
90 return Err(UdpError::InvalidDatagram);
91 }
92 let mut bytes = Vec::with_capacity(HEADER + self.payload.len());
93 bytes.push(self.kind);
94 bytes.extend_from_slice(&self.receiver_index.to_be_bytes());
95 bytes.extend_from_slice(&self.counter.to_be_bytes());
96 bytes.extend_from_slice(&self.payload);
97 Ok(bytes)
98 }
99 pub fn decode(bytes: &[u8]) -> Result<Self, UdpError> {
102 if bytes.len() > MAX_DATAGRAM {
103 return Err(UdpError::TooLarge);
104 }
105 if bytes.len() < HEADER || !(1..=4).contains(&bytes[0]) {
106 return Err(UdpError::InvalidDatagram);
107 }
108 Ok(Self {
109 kind: bytes[0],
110 receiver_index: u32::from_be_bytes(field(&bytes[1..5])?),
111 counter: u64::from_be_bytes(field(&bytes[5..HEADER])?),
112 payload: bytes[HEADER..].to_vec(),
113 })
114 }
115}
116#[derive(Debug, thiserror::Error)]
118pub enum UdpError {
119 #[error("udp_receive_deadline")]
120 Deadline,
121 #[error("udp_observation_in_progress")]
122 ObservationInProgress,
123 #[error("udp_acknowledgements_already_watched")]
124 AcknowledgementsAlreadyWatched,
125 #[error("too_large_for_datagram")]
126 TooLarge,
127 #[error("invalid_datagram")]
128 InvalidDatagram,
129 #[error("replayed_datagram")]
130 Replayed,
131 #[error("unknown_receiver_index")]
132 UnknownIndex,
133 #[error("receiver_index_already_registered")]
134 DuplicateIndex,
135 #[error("unexpected_link_endpoint")]
136 UnexpectedEndpoint,
137 #[error("invalid_acknowledgement")]
138 InvalidAck,
139 #[error("udp_state_poisoned")]
140 Poisoned,
141 #[error("udp_receiver_stopped")]
142 Stopped,
143 #[error("udp_receiver_panicked")]
144 WorkerPanicked,
145 #[error("replay_count_exhausted")]
146 ReplayCountExhausted,
147 #[error("udp_socket_must_be_unconnected")]
148 ConnectedSocket,
149 #[error("udp_io: {0}")]
150 Io(#[from] std::io::Error),
151 #[error("udp_noise: {0}")]
152 Noise(#[from] LinkHandshakeError),
153}
154fn field<const N: usize>(bytes: &[u8]) -> Result<[u8; N], UdpError> {
155 bytes.try_into().map_err(|_| UdpError::InvalidDatagram)
156}
157#[derive(Default)]
158struct ReplayWindow {
159 highest: Option<u64>,
160 bits: [u64; WINDOW / 64],
161}
162impl ReplayWindow {
163 fn contains(&self, counter: u64) -> bool {
164 self.highest.is_some_and(|highest| {
165 let Some(age) = highest.checked_sub(counter) else {
166 return false;
167 };
168 let Ok(age) = usize::try_from(age) else {
169 return true;
170 };
171 age >= WINDOW || self.bits[age / 64] & (1 << (age % 64)) != 0
172 })
173 }
174 fn record(&mut self, counter: u64) -> Result<(), UdpError> {
175 let highest = *self.highest.get_or_insert(counter);
176 if counter > highest {
177 if counter - highest >= 2048 {
178 self.bits.fill(0);
179 } else {
180 let distance =
181 usize::try_from(counter - highest).map_err(|_| UdpError::InvalidDatagram)?;
182 let words = distance / 64;
183 let shift = distance % 64;
184 self.bits.rotate_right(words);
185 self.bits[..words].fill(0);
186 for destination in (0..self.bits.len()).rev() {
187 let carry = if shift != 0 && destination > 0 {
188 self.bits[destination - 1] >> (64 - shift)
189 } else {
190 0
191 };
192 self.bits[destination] = (self.bits[destination] << shift) | carry;
193 }
194 }
195 }
196 let highest = highest.max(counter);
197 let age = usize::try_from(highest - counter).map_err(|_| UdpError::Replayed)?;
198 if age >= WINDOW {
199 return Err(UdpError::Replayed);
200 }
201 self.bits[age / 64] |= 1 << (age % 64);
202 self.highest = Some(highest);
203 Ok(())
204 }
205}
206struct UdpLink {
207 remote_index: u32,
208 endpoint: SocketAddr,
209 transport: LinkTransport,
210 window: ReplayWindow,
211 replayed: u64,
212 highest_acked: Option<u64>,
213 acknowledgements: Arc<Acknowledgements>,
214 cached_acks: std::collections::VecDeque<CachedAck>,
215}
216mod cancellation;
217pub(crate) use cancellation::LinkCancellation;
218mod recovery;
219pub mod streams;
220use recovery::{Acknowledgements, CachedAck, Observations, open_datagram};
221
222type Links = Arc<Mutex<HashMap<u32, UdpLink>>>;
223type Arrival = Result<(Datagram, SocketAddr), UdpError>;
224struct Incoming {
225 receiver: mpsc::Receiver<Arrival>,
226 parked: std::collections::VecDeque<Arrival>,
227}
228pub struct UdpCarrier {
231 socket: Arc<dyn PunchSocket>,
232 links: Links,
233 incoming: Option<Mutex<Incoming>>,
234 stopping: Arc<AtomicBool>,
235 worker: Option<JoinHandle<Result<(), UdpError>>>,
236 observations: Observations,
237}
238impl UdpCarrier {
239 pub fn new(socket: UdpSocket) -> Result<Self, UdpError> {
243 match socket.peer_addr() {
244 Ok(_) => return Err(UdpError::ConnectedSocket),
245 Err(error) if error.kind() == std::io::ErrorKind::NotConnected => {}
246 Err(error) => return Err(error.into()),
247 }
248 Self::from_socket(Arc::new(socket))
249 }
250 pub fn from_socket(socket: Arc<dyn PunchSocket>) -> Result<Self, UdpError> {
254 socket.prepare(None)?;
255 let reader = Arc::clone(&socket);
256 let links = Arc::new(Mutex::new(HashMap::new()));
257 let stopping = Arc::new(AtomicBool::new(false));
258 let (outgoing, incoming) = mpsc::sync_channel(64);
259 let worker_links = Arc::clone(&links);
260 let worker_stopping = Arc::clone(&stopping);
261 let observations = Arc::new(Mutex::new(None));
262 let worker_observations = Arc::clone(&observations);
263 let worker = thread::Builder::new()
264 .name("link-udp".to_owned())
265 .spawn(move || {
266 receive_loop(
267 reader.as_ref(),
268 &worker_links,
269 &worker_stopping,
270 &outgoing,
271 &worker_observations,
272 )
273 })?;
274 Ok(Self {
275 socket,
276 links,
277 incoming: Some(Mutex::new(Incoming {
278 receiver: incoming,
279 parked: std::collections::VecDeque::new(),
280 })),
281 stopping,
282 worker: Some(worker),
283 observations,
284 })
285 }
286 pub fn add_link(
289 &self,
290 local_index: u32,
291 remote_index: u32,
292 endpoint: SocketAddr,
293 transport: LinkTransport,
294 ) -> Result<(), UdpError> {
295 let mut links = self.links.lock().map_err(|_| UdpError::Poisoned)?;
296 if links.contains_key(&local_index) {
297 return Err(UdpError::DuplicateIndex);
298 }
299 links.insert(
300 local_index,
301 UdpLink {
302 remote_index,
303 endpoint,
304 transport,
305 window: ReplayWindow::default(),
306 replayed: 0,
307 highest_acked: None,
308 acknowledgements: Arc::default(),
309 cached_acks: std::collections::VecDeque::new(),
310 },
311 );
312 drop(links);
313 Ok(())
314 }
315 pub fn seal(&self, index: u32, payload: &[u8]) -> Result<Datagram, UdpError> {
318 let escaped = matches!(payload.first(), Some(3 | 4));
319 if payload.len() > MAX_DATAGRAM - HEADER - TAG - usize::from(escaped) {
320 return Err(UdpError::TooLarge);
321 }
322 let mut plaintext = payload.to_vec();
323 if escaped {
324 plaintext.insert(0, 3);
325 }
326 let mut links = self.links.lock().map_err(|_| UdpError::Poisoned)?;
327 let link = links.get_mut(&index).ok_or(UdpError::UnknownIndex)?;
328 let counter = link.transport.sending_nonce();
329 let receiver_index = link.remote_index;
330 let mut sealed = vec![0; plaintext.len() + TAG];
331 let length = link.transport.seal(&plaintext, &mut sealed)?;
332 drop(links);
333 sealed.truncate(length);
334 Ok(Datagram {
335 kind: 3,
336 receiver_index,
337 counter,
338 payload: sealed,
339 })
340 }
341 pub fn send(&self, index: u32, payload: &[u8]) -> Result<u64, UdpError> {
344 let datagram = self.seal(index, payload)?;
345 let endpoint = self.inspect_link(index, |link| link.endpoint)?;
346 send_datagram(self.socket.as_ref(), endpoint, &datagram)?;
347 Ok(datagram.counter)
348 }
349 pub fn transmit(&self, endpoint: SocketAddr, datagram: &Datagram) -> Result<(), UdpError> {
353 send_datagram(self.socket.as_ref(), endpoint, datagram)
354 }
355 pub fn receive(&self) -> Result<Datagram, UdpError> {
359 self.receive_from().map(|(datagram, _)| datagram)
360 }
361 pub fn receive_from(&self) -> Arrival {
364 let mut incoming = self
365 .incoming
366 .as_ref()
367 .ok_or(UdpError::Stopped)?
368 .lock()
369 .map_err(|_| UdpError::Poisoned)?;
370 if let Some(arrival) = incoming.parked.pop_front() {
371 return arrival;
372 }
373 incoming.receiver.recv().map_err(|_| UdpError::Stopped)?
374 }
375 pub fn highest_acked(&self, index: u32) -> Result<Option<u64>, UdpError> {
378 self.inspect_link(index, |link| link.highest_acked)
379 }
380 pub fn replayed(&self, index: u32) -> Result<u64, UdpError> {
383 self.inspect_link(index, |link| link.replayed)
384 }
385 fn inspect_link<T>(&self, index: u32, read: impl FnOnce(&UdpLink) -> T) -> Result<T, UdpError> {
386 let links = self.links.lock().map_err(|_| UdpError::Poisoned)?;
387 Ok(read(links.get(&index).ok_or(UdpError::UnknownIndex)?))
388 }
389 pub fn shutdown(mut self) -> Result<(), UdpError> {
393 self.stop_worker()
394 }
395 fn stop_worker(&mut self) -> Result<(), UdpError> {
396 let Some(worker) = self.worker.as_ref() else {
397 return Ok(());
398 };
399 self.stopping.store(true, Ordering::SeqCst);
400 self.cancel_waits()?;
401 recovery::cancel_observation(&self.observations)?;
402 self.incoming.take();
403 if !worker.is_finished() {
404 self.socket.wake()?;
405 }
406 let worker = self.worker.take().ok_or(UdpError::Stopped)?;
407 worker.join().map_err(|_| UdpError::WorkerPanicked)?
408 }
409}
410impl Drop for UdpCarrier {
411 fn drop(&mut self) {
412 if let Err(error) = self.stop_worker() {
413 tracing::error!(%error, "UDP receiver shutdown failed");
414 }
415 }
416}
417fn send_datagram(
418 socket: &dyn PunchSocket,
419 endpoint: SocketAddr,
420 datagram: &Datagram,
421) -> Result<(), UdpError> {
422 let bytes = datagram.encode()?;
423 if socket.send_to(&bytes, endpoint)? != bytes.len() {
424 return Err(
425 std::io::Error::new(std::io::ErrorKind::WriteZero, "partial UDP datagram").into(),
426 );
427 }
428 Ok(())
429}
430fn receive_loop(
431 socket: &dyn PunchSocket,
432 links: &Links,
433 stopping: &AtomicBool,
434 outgoing: &mpsc::SyncSender<Arrival>,
435 observations: &Observations,
436) -> Result<(), UdpError> {
437 let mut bytes = [0; MAX_DATAGRAM + 1];
438 while !stopping.load(Ordering::SeqCst) {
439 let (length, endpoint) = match socket.recv_from(&mut bytes) {
440 Ok(received) => received,
441 Err(error) if error.kind() == std::io::ErrorKind::Interrupted => continue,
442 Err(error) => {
443 tracing::warn!(%error, "UDP receiver failed");
444 stopping.store(true, Ordering::Release);
445 recovery::cancel_links(links)?;
446 recovery::cancel_observation(observations)?;
447 outgoing
448 .send(Err(error.into()))
449 .map_err(|_| UdpError::Stopped)?;
450 return Ok(());
451 }
452 };
453 if stopping.load(Ordering::SeqCst) {
454 break;
455 }
456 if recovery::route_observation(observations, &bytes[..length], endpoint)? {
457 continue;
458 }
459 let index = Datagram::decode(&bytes[..length])
460 .ok()
461 .map(|datagram| datagram.receiver_index);
462 let result = Datagram::decode(&bytes[..length])
463 .and_then(|datagram| open_datagram(socket, links, endpoint, datagram))
464 .map(|datagram| (datagram, endpoint));
465 if let Ok((datagram, _)) = &result
466 && datagram.kind == 4
467 && recovery::watched(links, datagram.receiver_index)?
468 {
469 continue;
470 }
471 if !streams::route_arrival(links, index, result, outgoing)? {
472 break;
473 }
474 }
475 Ok(())
476}