Skip to main content

liminal_server/
link_udp.rs

1//! Indexed Noise messages on one blocking UDP receiver per device.
2//!
3//! Integers are big endian. ACK headers name the acknowledged data counter.
4//! ACK bodies carry an eight-byte nonce and sealed kind and counter in data's sending sequence.
5//! Sealed ACKs start with 4; escaping data starting with 3 or 4 keeps kinds disjoint.
6use 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;
14/// Datagram I/O shared by address observation and the direct carrier.
15/// A gateway adapter must preserve the actual remote endpoint on receive.
16pub trait PunchSocket: Send + Sync {
17    /// # Errors
18    /// Propagates socket or gateway failures.
19    fn send_to(&self, bytes: &[u8], target: SocketAddr) -> io::Result<usize>;
20    /// # Errors
21    /// Propagates socket or gateway failures.
22    fn recv_from(&self, bytes: &mut [u8]) -> io::Result<(usize, SocketAddr)>;
23    /// # Errors
24    /// Propagates address inspection failure.
25    fn local_addr(&self) -> io::Result<SocketAddr>;
26    /// Sets an observation bound, or parks indefinitely on delivery.
27    /// # Errors
28    /// Refuses connected sockets and propagates socket setup failures.
29    fn prepare(&self, bound: Option<Duration>) -> io::Result<()>;
30    /// Interrupts a parked receiver through its owned socket.
31    /// # Errors
32    /// Propagates address inspection and wake delivery failures.
33    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/// A wire datagram; `payload` contains sealed bytes for data and ACKs.
72#[derive(Debug, Clone, PartialEq, Eq)]
73pub struct Datagram {
74    /// First handshake (1), second handshake (2), data (3), or ACK (4).
75    pub kind: u8,
76    /// The destination's locally assigned link index.
77    pub receiver_index: u32,
78    /// Data's Noise nonce, or the data counter acknowledged by an ACK.
79    pub counter: u64,
80    pub payload: Vec<u8>,
81}
82impl Datagram {
83    /// # Errors
84    /// Refuses unknown kinds and frames larger than the datagram bound.
85    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    /// # Errors
100    /// Refuses short, unknown-kind, or oversized datagrams.
101    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/// A named carrier refusal or a propagated transport failure.
117#[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}
228/// One device socket, one parked receive thread, and independently keyed links.
229/// Received data payloads are opened; handshake payloads remain untouched.
230pub 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    /// Starts a blocking receiver with a bounded delivery queue.
240    /// # Errors
241    /// Propagates socket setup and thread creation errors; refuses connected sockets.
242    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    /// Starts the carrier on the same socket used for broker observation.
251    /// # Errors
252    /// Propagates socket setup and thread creation failures.
253    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    /// # Errors
287    /// Refuses replacement of an existing index and poisoned state.
288    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    /// # Errors
316    /// Refuses oversize payloads before consuming a nonce, unknown links, and Noise errors.
317    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    /// # Errors
342    /// Propagates sealing and send failures without changing carriers.
343    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    /// Sends a handshake or an already sealed frame without deriving a key.
350    /// # Errors
351    /// Refuses invalid or oversized datagrams; propagates send errors.
352    pub fn transmit(&self, endpoint: SocketAddr, datagram: &Datagram) -> Result<(), UdpError> {
353        send_datagram(self.socket.as_ref(), endpoint, datagram)
354    }
355    /// Receives an opened message or a named refusal, waiting only on delivery.
356    /// # Errors
357    /// Propagates receive, authentication, framing, and replay refusals.
358    pub fn receive(&self) -> Result<Datagram, UdpError> {
359        self.receive_from().map(|(datagram, _)| datagram)
360    }
361    /// # Errors
362    /// Propagates the same refusals as `receive` and a stopped receiver.
363    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    /// # Errors
376    /// Refuses unknown links or poisoned state.
377    pub fn highest_acked(&self, index: u32) -> Result<Option<u64>, UdpError> {
378        self.inspect_link(index, |link| link.highest_acked)
379    }
380    /// # Errors
381    /// Refuses unknown links or poisoned state.
382    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    /// Signals the parked receiver and joins it, including when its queue is full.
390    /// # Errors
391    /// Propagates interruption, worker, or join failures.
392    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}