liminal-server 0.25.4

//! Indexed Noise messages on one blocking UDP receiver per device.
//!
//! Integers are big endian. ACK headers name the acknowledged data counter.
//! ACK bodies carry an eight-byte nonce and sealed kind and counter in data's sending sequence.
//! Sealed ACKs start with 4; escaping data starting with 3 or 4 keeps kinds disjoint.
use crate::link_handshake::{LinkHandshakeError, LinkTransport};
use std::collections::HashMap;
use std::io;
use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, UdpSocket};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, mpsc};
use std::thread::{self, JoinHandle};
use std::time::Duration;
/// Datagram I/O shared by address observation and the direct carrier.
/// A gateway adapter must preserve the actual remote endpoint on receive.
pub trait PunchSocket: Send + Sync {
    /// # Errors
    /// Propagates socket or gateway failures.
    fn send_to(&self, bytes: &[u8], target: SocketAddr) -> io::Result<usize>;
    /// # Errors
    /// Propagates socket or gateway failures.
    fn recv_from(&self, bytes: &mut [u8]) -> io::Result<(usize, SocketAddr)>;
    /// # Errors
    /// Propagates address inspection failure.
    fn local_addr(&self) -> io::Result<SocketAddr>;
    /// Sets an observation bound, or parks indefinitely on delivery.
    /// # Errors
    /// Refuses connected sockets and propagates socket setup failures.
    fn prepare(&self, bound: Option<Duration>) -> io::Result<()>;
    /// Interrupts a parked receiver through its owned socket.
    /// # Errors
    /// Propagates address inspection and wake delivery failures.
    fn wake(&self) -> io::Result<()> {
        let target = match self.local_addr()? {
            SocketAddr::V4(address) if address.ip().is_unspecified() => {
                SocketAddr::from((Ipv4Addr::LOCALHOST, address.port()))
            }
            SocketAddr::V6(address) if address.ip().is_unspecified() => {
                SocketAddr::from((Ipv6Addr::LOCALHOST, address.port()))
            }
            address => address,
        };
        self.send_to(&[], target)?;
        Ok(())
    }
}
impl PunchSocket for UdpSocket {
    fn send_to(&self, bytes: &[u8], target: SocketAddr) -> io::Result<usize> {
        Self::send_to(self, bytes, target)
    }
    fn recv_from(&self, bytes: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
        Self::recv_from(self, bytes)
    }
    fn local_addr(&self) -> io::Result<SocketAddr> {
        Self::local_addr(self)
    }
    fn prepare(&self, bound: Option<Duration>) -> io::Result<()> {
        match self.peer_addr() {
            Ok(_) => return Err(io::Error::other("udp_socket_must_be_unconnected")),
            Err(error) if error.kind() == io::ErrorKind::NotConnected => {}
            Err(error) => return Err(error),
        }
        self.set_nonblocking(false)?;
        self.set_read_timeout(bound)
    }
}
pub const MAX_DATAGRAM: usize = 1200;
const HEADER: usize = 13;
const TAG: usize = 16;
const WINDOW: usize = 2048;
/// A wire datagram; `payload` contains sealed bytes for data and ACKs.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Datagram {
    /// First handshake (1), second handshake (2), data (3), or ACK (4).
    pub kind: u8,
    /// The destination's locally assigned link index.
    pub receiver_index: u32,
    /// Data's Noise nonce, or the data counter acknowledged by an ACK.
    pub counter: u64,
    pub payload: Vec<u8>,
}
impl Datagram {
    /// # Errors
    /// Refuses unknown kinds and frames larger than the datagram bound.
    pub fn encode(&self) -> Result<Vec<u8>, UdpError> {
        if self.payload.len() > MAX_DATAGRAM - HEADER {
            return Err(UdpError::TooLarge);
        }
        if !(1..=4).contains(&self.kind) {
            return Err(UdpError::InvalidDatagram);
        }
        let mut bytes = Vec::with_capacity(HEADER + self.payload.len());
        bytes.push(self.kind);
        bytes.extend_from_slice(&self.receiver_index.to_be_bytes());
        bytes.extend_from_slice(&self.counter.to_be_bytes());
        bytes.extend_from_slice(&self.payload);
        Ok(bytes)
    }
    /// # Errors
    /// Refuses short, unknown-kind, or oversized datagrams.
    pub fn decode(bytes: &[u8]) -> Result<Self, UdpError> {
        if bytes.len() > MAX_DATAGRAM {
            return Err(UdpError::TooLarge);
        }
        if bytes.len() < HEADER || !(1..=4).contains(&bytes[0]) {
            return Err(UdpError::InvalidDatagram);
        }
        Ok(Self {
            kind: bytes[0],
            receiver_index: u32::from_be_bytes(field(&bytes[1..5])?),
            counter: u64::from_be_bytes(field(&bytes[5..HEADER])?),
            payload: bytes[HEADER..].to_vec(),
        })
    }
}
/// A named carrier refusal or a propagated transport failure.
#[derive(Debug, thiserror::Error)]
pub enum UdpError {
    #[error("udp_receive_deadline")]
    Deadline,
    #[error("udp_observation_in_progress")]
    ObservationInProgress,
    #[error("udp_acknowledgements_already_watched")]
    AcknowledgementsAlreadyWatched,
    #[error("too_large_for_datagram")]
    TooLarge,
    #[error("invalid_datagram")]
    InvalidDatagram,
    #[error("replayed_datagram")]
    Replayed,
    #[error("unknown_receiver_index")]
    UnknownIndex,
    #[error("receiver_index_already_registered")]
    DuplicateIndex,
    #[error("unexpected_link_endpoint")]
    UnexpectedEndpoint,
    #[error("invalid_acknowledgement")]
    InvalidAck,
    #[error("udp_state_poisoned")]
    Poisoned,
    #[error("udp_receiver_stopped")]
    Stopped,
    #[error("udp_receiver_panicked")]
    WorkerPanicked,
    #[error("replay_count_exhausted")]
    ReplayCountExhausted,
    #[error("udp_socket_must_be_unconnected")]
    ConnectedSocket,
    #[error("udp_io: {0}")]
    Io(#[from] std::io::Error),
    #[error("udp_noise: {0}")]
    Noise(#[from] LinkHandshakeError),
}
fn field<const N: usize>(bytes: &[u8]) -> Result<[u8; N], UdpError> {
    bytes.try_into().map_err(|_| UdpError::InvalidDatagram)
}
#[derive(Default)]
struct ReplayWindow {
    highest: Option<u64>,
    bits: [u64; WINDOW / 64],
}
impl ReplayWindow {
    fn contains(&self, counter: u64) -> bool {
        self.highest.is_some_and(|highest| {
            let Some(age) = highest.checked_sub(counter) else {
                return false;
            };
            let Ok(age) = usize::try_from(age) else {
                return true;
            };
            age >= WINDOW || self.bits[age / 64] & (1 << (age % 64)) != 0
        })
    }
    fn record(&mut self, counter: u64) -> Result<(), UdpError> {
        let highest = *self.highest.get_or_insert(counter);
        if counter > highest {
            if counter - highest >= 2048 {
                self.bits.fill(0);
            } else {
                let distance =
                    usize::try_from(counter - highest).map_err(|_| UdpError::InvalidDatagram)?;
                let words = distance / 64;
                let shift = distance % 64;
                self.bits.rotate_right(words);
                self.bits[..words].fill(0);
                for destination in (0..self.bits.len()).rev() {
                    let carry = if shift != 0 && destination > 0 {
                        self.bits[destination - 1] >> (64 - shift)
                    } else {
                        0
                    };
                    self.bits[destination] = (self.bits[destination] << shift) | carry;
                }
            }
        }
        let highest = highest.max(counter);
        let age = usize::try_from(highest - counter).map_err(|_| UdpError::Replayed)?;
        if age >= WINDOW {
            return Err(UdpError::Replayed);
        }
        self.bits[age / 64] |= 1 << (age % 64);
        self.highest = Some(highest);
        Ok(())
    }
}
struct UdpLink {
    remote_index: u32,
    endpoint: SocketAddr,
    transport: LinkTransport,
    window: ReplayWindow,
    replayed: u64,
    highest_acked: Option<u64>,
    acknowledgements: Arc<Acknowledgements>,
    cached_acks: std::collections::VecDeque<CachedAck>,
}
mod cancellation;
pub(crate) use cancellation::LinkCancellation;
mod recovery;
pub mod streams;
use recovery::{Acknowledgements, CachedAck, Observations, open_datagram};

type Links = Arc<Mutex<HashMap<u32, UdpLink>>>;
type Arrival = Result<(Datagram, SocketAddr), UdpError>;
struct Incoming {
    receiver: mpsc::Receiver<Arrival>,
    parked: std::collections::VecDeque<Arrival>,
}
/// One device socket, one parked receive thread, and independently keyed links.
/// Received data payloads are opened; handshake payloads remain untouched.
pub struct UdpCarrier {
    socket: Arc<dyn PunchSocket>,
    links: Links,
    incoming: Option<Mutex<Incoming>>,
    stopping: Arc<AtomicBool>,
    worker: Option<JoinHandle<Result<(), UdpError>>>,
    observations: Observations,
}
impl UdpCarrier {
    /// Starts a blocking receiver with a bounded delivery queue.
    /// # Errors
    /// Propagates socket setup and thread creation errors; refuses connected sockets.
    pub fn new(socket: UdpSocket) -> Result<Self, UdpError> {
        match socket.peer_addr() {
            Ok(_) => return Err(UdpError::ConnectedSocket),
            Err(error) if error.kind() == std::io::ErrorKind::NotConnected => {}
            Err(error) => return Err(error.into()),
        }
        Self::from_socket(Arc::new(socket))
    }
    /// Starts the carrier on the same socket used for broker observation.
    /// # Errors
    /// Propagates socket setup and thread creation failures.
    pub fn from_socket(socket: Arc<dyn PunchSocket>) -> Result<Self, UdpError> {
        socket.prepare(None)?;
        let reader = Arc::clone(&socket);
        let links = Arc::new(Mutex::new(HashMap::new()));
        let stopping = Arc::new(AtomicBool::new(false));
        let (outgoing, incoming) = mpsc::sync_channel(64);
        let worker_links = Arc::clone(&links);
        let worker_stopping = Arc::clone(&stopping);
        let observations = Arc::new(Mutex::new(None));
        let worker_observations = Arc::clone(&observations);
        let worker = thread::Builder::new()
            .name("link-udp".to_owned())
            .spawn(move || {
                receive_loop(
                    reader.as_ref(),
                    &worker_links,
                    &worker_stopping,
                    &outgoing,
                    &worker_observations,
                )
            })?;
        Ok(Self {
            socket,
            links,
            incoming: Some(Mutex::new(Incoming {
                receiver: incoming,
                parked: std::collections::VecDeque::new(),
            })),
            stopping,
            worker: Some(worker),
            observations,
        })
    }
    /// # Errors
    /// Refuses replacement of an existing index and poisoned state.
    pub fn add_link(
        &self,
        local_index: u32,
        remote_index: u32,
        endpoint: SocketAddr,
        transport: LinkTransport,
    ) -> Result<(), UdpError> {
        let mut links = self.links.lock().map_err(|_| UdpError::Poisoned)?;
        if links.contains_key(&local_index) {
            return Err(UdpError::DuplicateIndex);
        }
        links.insert(
            local_index,
            UdpLink {
                remote_index,
                endpoint,
                transport,
                window: ReplayWindow::default(),
                replayed: 0,
                highest_acked: None,
                acknowledgements: Arc::default(),
                cached_acks: std::collections::VecDeque::new(),
            },
        );
        drop(links);
        Ok(())
    }
    /// # Errors
    /// Refuses oversize payloads before consuming a nonce, unknown links, and Noise errors.
    pub fn seal(&self, index: u32, payload: &[u8]) -> Result<Datagram, UdpError> {
        let escaped = matches!(payload.first(), Some(3 | 4));
        if payload.len() > MAX_DATAGRAM - HEADER - TAG - usize::from(escaped) {
            return Err(UdpError::TooLarge);
        }
        let mut plaintext = payload.to_vec();
        if escaped {
            plaintext.insert(0, 3);
        }
        let mut links = self.links.lock().map_err(|_| UdpError::Poisoned)?;
        let link = links.get_mut(&index).ok_or(UdpError::UnknownIndex)?;
        let counter = link.transport.sending_nonce();
        let receiver_index = link.remote_index;
        let mut sealed = vec![0; plaintext.len() + TAG];
        let length = link.transport.seal(&plaintext, &mut sealed)?;
        drop(links);
        sealed.truncate(length);
        Ok(Datagram {
            kind: 3,
            receiver_index,
            counter,
            payload: sealed,
        })
    }
    /// # Errors
    /// Propagates sealing and send failures without changing carriers.
    pub fn send(&self, index: u32, payload: &[u8]) -> Result<u64, UdpError> {
        let datagram = self.seal(index, payload)?;
        let endpoint = self.inspect_link(index, |link| link.endpoint)?;
        send_datagram(self.socket.as_ref(), endpoint, &datagram)?;
        Ok(datagram.counter)
    }
    /// Sends a handshake or an already sealed frame without deriving a key.
    /// # Errors
    /// Refuses invalid or oversized datagrams; propagates send errors.
    pub fn transmit(&self, endpoint: SocketAddr, datagram: &Datagram) -> Result<(), UdpError> {
        send_datagram(self.socket.as_ref(), endpoint, datagram)
    }
    /// Receives an opened message or a named refusal, waiting only on delivery.
    /// # Errors
    /// Propagates receive, authentication, framing, and replay refusals.
    pub fn receive(&self) -> Result<Datagram, UdpError> {
        self.receive_from().map(|(datagram, _)| datagram)
    }
    /// # Errors
    /// Propagates the same refusals as `receive` and a stopped receiver.
    pub fn receive_from(&self) -> Arrival {
        let mut incoming = self
            .incoming
            .as_ref()
            .ok_or(UdpError::Stopped)?
            .lock()
            .map_err(|_| UdpError::Poisoned)?;
        if let Some(arrival) = incoming.parked.pop_front() {
            return arrival;
        }
        incoming.receiver.recv().map_err(|_| UdpError::Stopped)?
    }
    /// # Errors
    /// Refuses unknown links or poisoned state.
    pub fn highest_acked(&self, index: u32) -> Result<Option<u64>, UdpError> {
        self.inspect_link(index, |link| link.highest_acked)
    }
    /// # Errors
    /// Refuses unknown links or poisoned state.
    pub fn replayed(&self, index: u32) -> Result<u64, UdpError> {
        self.inspect_link(index, |link| link.replayed)
    }
    fn inspect_link<T>(&self, index: u32, read: impl FnOnce(&UdpLink) -> T) -> Result<T, UdpError> {
        let links = self.links.lock().map_err(|_| UdpError::Poisoned)?;
        Ok(read(links.get(&index).ok_or(UdpError::UnknownIndex)?))
    }
    /// Signals the parked receiver and joins it, including when its queue is full.
    /// # Errors
    /// Propagates interruption, worker, or join failures.
    pub fn shutdown(mut self) -> Result<(), UdpError> {
        self.stop_worker()
    }
    fn stop_worker(&mut self) -> Result<(), UdpError> {
        let Some(worker) = self.worker.as_ref() else {
            return Ok(());
        };
        self.stopping.store(true, Ordering::SeqCst);
        self.cancel_waits()?;
        recovery::cancel_observation(&self.observations)?;
        self.incoming.take();
        if !worker.is_finished() {
            self.socket.wake()?;
        }
        let worker = self.worker.take().ok_or(UdpError::Stopped)?;
        worker.join().map_err(|_| UdpError::WorkerPanicked)?
    }
}
impl Drop for UdpCarrier {
    fn drop(&mut self) {
        if let Err(error) = self.stop_worker() {
            tracing::error!(%error, "UDP receiver shutdown failed");
        }
    }
}
fn send_datagram(
    socket: &dyn PunchSocket,
    endpoint: SocketAddr,
    datagram: &Datagram,
) -> Result<(), UdpError> {
    let bytes = datagram.encode()?;
    if socket.send_to(&bytes, endpoint)? != bytes.len() {
        return Err(
            std::io::Error::new(std::io::ErrorKind::WriteZero, "partial UDP datagram").into(),
        );
    }
    Ok(())
}
fn receive_loop(
    socket: &dyn PunchSocket,
    links: &Links,
    stopping: &AtomicBool,
    outgoing: &mpsc::SyncSender<Arrival>,
    observations: &Observations,
) -> Result<(), UdpError> {
    let mut bytes = [0; MAX_DATAGRAM + 1];
    while !stopping.load(Ordering::SeqCst) {
        let (length, endpoint) = match socket.recv_from(&mut bytes) {
            Ok(received) => received,
            Err(error) if error.kind() == std::io::ErrorKind::Interrupted => continue,
            Err(error) => {
                tracing::warn!(%error, "UDP receiver failed");
                stopping.store(true, Ordering::Release);
                recovery::cancel_links(links)?;
                recovery::cancel_observation(observations)?;
                outgoing
                    .send(Err(error.into()))
                    .map_err(|_| UdpError::Stopped)?;
                return Ok(());
            }
        };
        if stopping.load(Ordering::SeqCst) {
            break;
        }
        if recovery::route_observation(observations, &bytes[..length], endpoint)? {
            continue;
        }
        let index = Datagram::decode(&bytes[..length])
            .ok()
            .map(|datagram| datagram.receiver_index);
        let result = Datagram::decode(&bytes[..length])
            .and_then(|datagram| open_datagram(socket, links, endpoint, datagram))
            .map(|datagram| (datagram, endpoint));
        if let Ok((datagram, _)) = &result
            && datagram.kind == 4
            && recovery::watched(links, datagram.receiver_index)?
        {
            continue;
        }
        if !streams::route_arrival(links, index, result, outgoing)? {
            break;
        }
    }
    Ok(())
}