liminal-server 0.25.4

//! Message-scoped ACK signals and authenticated retransmission receipts.
use super::{
    Datagram, Links, PunchSocket, TAG, UdpCarrier, UdpError, UdpLink, WINDOW, field, send_datagram,
};
use std::net::SocketAddr;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;

pub(super) type Observations = Arc<Mutex<Option<Observation>>>;
type EchoAnswer = Result<(Vec<u8>, SocketAddr), UdpError>;
pub(super) struct Observation {
    nonce: [u8; 16],
    endpoint: SocketAddr,
    reply: std::sync::mpsc::Sender<EchoAnswer>,
}

pub(super) fn route_observation(
    observations: &Observations,
    bytes: &[u8],
    source: SocketAddr,
) -> Result<bool, UdpError> {
    let mut observations = observations.lock().map_err(|_| UdpError::Poisoned)?;
    if !observations.as_ref().is_some_and(|pending| {
        source == pending.endpoint
            && bytes.len() >= 20
            && bytes[..4] == *b"LMOA"
            && bytes[4..20] == pending.nonce
    }) {
        return Ok(false);
    }
    let pending = observations.take().ok_or(UdpError::Stopped)?;
    drop(observations);
    if pending.reply.send(Ok((bytes.to_vec(), source))).is_err() {
        tracing::debug!("broker observation ended before its reply");
    }
    Ok(true)
}

pub(super) fn cancel_observation(observations: &Observations) -> Result<(), UdpError> {
    if let Some(pending) = observations.lock().map_err(|_| UdpError::Poisoned)?.take()
        && pending.reply.send(Err(UdpError::Stopped)).is_err()
    {
        tracing::debug!("broker observation was already closed at stop");
    }
    Ok(())
}

#[derive(Debug, Default)]
pub(super) struct Acknowledgements {
    counters: Mutex<std::collections::BTreeSet<u64>>,
    changed: std::sync::Condvar,
    watched: AtomicBool,
    cancelled: AtomicBool,
    pub(super) route: Mutex<Option<std::sync::mpsc::SyncSender<super::Arrival>>>,
}

impl Acknowledgements {
    pub(super) fn is_cancelled(&self) -> bool {
        self.cancelled.load(Ordering::Acquire)
    }

    pub(super) fn cancel(&self) -> Result<(), UdpError> {
        let mut route = self.route.lock().map_err(|_| UdpError::Poisoned)?;
        self.cancelled.store(true, Ordering::Release);
        route.take();
        drop(route);
        let guard = self.counters.lock().map_err(|_| UdpError::Poisoned)?;
        self.changed.notify_all();
        drop(guard);
        Ok(())
    }
}

pub(super) struct CachedAck {
    pub(super) counter: u64,
    digest: [u8; 32],
    datagram: Datagram,
    pub(super) receipt: Option<Datagram>,
}

impl UdpCarrier {
    /// Observes a new opening on the existing device socket, without another reader.
    /// # Errors
    /// Invalid request, concurrent observation, send failure, or this opening's deadline.
    pub fn observe_raw(
        &self,
        request: &[u8; 20],
        endpoint: SocketAddr,
        bound: Duration,
    ) -> Result<(Vec<u8>, SocketAddr), UdpError> {
        if self.stopping.load(Ordering::Acquire) {
            return Err(UdpError::Stopped);
        }
        if request[..4] != *b"LMOB" || endpoint.port() == 0 || bound.is_zero() {
            return Err(UdpError::InvalidDatagram);
        }
        let nonce = field(&request[4..])?;
        let (reply, incoming) = std::sync::mpsc::channel();
        {
            let mut pending = self.observations.lock().map_err(|_| UdpError::Poisoned)?;
            if pending.is_some() {
                return Err(UdpError::ObservationInProgress);
            }
            *pending = Some(Observation {
                nonce,
                endpoint,
                reply,
            });
        }
        let result = (|| {
            if self.socket.send_to(request, endpoint)? != request.len() {
                return Err(std::io::Error::new(
                    std::io::ErrorKind::WriteZero,
                    "incomplete broker observation",
                )
                .into());
            }
            incoming.recv_timeout(bound).map_err(|error| match error {
                std::sync::mpsc::RecvTimeoutError::Timeout => UdpError::Deadline,
                std::sync::mpsc::RecvTimeoutError::Disconnected => UdpError::Stopped,
            })?
        })();
        let mut pending = self.observations.lock().map_err(|_| UdpError::Poisoned)?;
        if pending
            .as_ref()
            .is_some_and(|pending| pending.nonce == nonce)
        {
            pending.take();
        }
        result
    }
    /// Bounds one opening observation while the receiver stays parked at idle.
    /// # Errors
    /// A receive error, poison, stop, or this opening's deadline.
    pub fn receive_from_until(
        &self,
        deadline: std::time::Instant,
    ) -> Result<(Datagram, SocketAddr), UdpError> {
        let bound = deadline.saturating_duration_since(std::time::Instant::now());
        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_timeout(bound)
            .map_err(|error| match error {
                std::sync::mpsc::RecvTimeoutError::Timeout => UdpError::Deadline,
                std::sync::mpsc::RecvTimeoutError::Disconnected => UdpError::Stopped,
            })?
    }

    /// Routes ACKs to this link's signalled deadline waiter.
    /// # Errors
    /// Unknown index or poisoned state.
    pub fn watch_acknowledgements(&self, index: u32) -> Result<(), UdpError> {
        self.inspect_link(index, |link| {
            link.acknowledgements
                .watched
                .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
                .map(|_| ())
                .map_err(|_| UdpError::AcknowledgementsAlreadyWatched)
        })?
    }

    /// Retires only this owner's keyed slot and interrupts its exact ACK waits.
    /// # Errors
    /// Unknown index or poisoned state.
    pub fn remove_link(&self, index: u32) -> Result<(), UdpError> {
        let link = self
            .links
            .lock()
            .map_err(|_| UdpError::Poisoned)?
            .remove(&index)
            .ok_or(UdpError::UnknownIndex)?;
        link.acknowledgements.cancel()
    }

    /// Waits for this exact authenticated counter with one message deadline.
    /// # Errors
    /// Unknown link, poisoned state, stopped carrier, or overflowing deadline.
    pub fn wait_acknowledged(
        &self,
        index: u32,
        counter: u64,
        bound: Duration,
    ) -> Result<bool, UdpError> {
        let deadline = std::time::Instant::now()
            .checked_add(bound)
            .ok_or(UdpError::Deadline)?;
        self.wait_acknowledged_until(index, counter, deadline)
    }

    /// Waits on an already armed message deadline, including time spent sending.
    /// # Errors
    /// Unknown link, poisoned state, or stopped carrier.
    pub fn wait_acknowledged_until(
        &self,
        index: u32,
        counter: u64,
        deadline: std::time::Instant,
    ) -> Result<bool, UdpError> {
        let bound = deadline.saturating_duration_since(std::time::Instant::now());
        let signal = self.inspect_link(index, |link| Arc::clone(&link.acknowledgements))?;
        let (mut received, _) = signal
            .changed
            .wait_timeout_while(
                signal.counters.lock().map_err(|_| UdpError::Poisoned)?,
                bound,
                |counters| {
                    !counters.contains(&counter) && !signal.cancelled.load(Ordering::Acquire)
                },
            )
            .map_err(|_| UdpError::Poisoned)?;
        let answered = received.remove(&counter);
        drop(received);
        if signal.cancelled.load(Ordering::Acquire) {
            return Err(UdpError::Stopped);
        }
        Ok(answered)
    }

    /// Refreshes a mapping without changing either Noise key or nonce sequence.
    /// # Errors
    /// Unknown link, invalid endpoint, or poisoned state.
    pub fn refresh_endpoint(&self, index: u32, endpoint: SocketAddr) -> Result<(), UdpError> {
        if endpoint.port() == 0 {
            return Err(UdpError::UnexpectedEndpoint);
        }
        let mut links = self.links.lock().map_err(|_| UdpError::Poisoned)?;
        links
            .get_mut(&index)
            .ok_or(UdpError::UnknownIndex)?
            .endpoint = endpoint;
        drop(links);
        Ok(())
    }

    /// Sends retained ciphertext to the link's currently observed endpoint.
    /// # Errors
    /// Unknown link, mismatched index, or send failure.
    pub fn retransmit(&self, index: u32, datagram: &Datagram) -> Result<(), UdpError> {
        let (endpoint, remote) =
            self.inspect_link(index, |link| (link.endpoint, link.remote_index))?;
        if datagram.kind != 3 || datagram.receiver_index != remote {
            return Err(UdpError::InvalidDatagram);
        }
        self.transmit(endpoint, datagram)
    }

    /// Cancels only this keyed link's in-flight observations.
    /// # Errors
    /// Unknown index or poisoned state.
    pub fn cancel_link(&self, index: u32) -> Result<(), UdpError> {
        self.inspect_link(index, |link| link.acknowledgements.cancel())?
    }
    /// Refreshes all keyed links to this authenticated device on its one UDP mapping.
    /// # Errors
    /// Invalid endpoint, missing peer, or poisoned state.
    pub fn refresh_peer(
        &self,
        peer: crate::link_handshake::MachineIdentity,
        endpoint: SocketAddr,
    ) -> Result<(), UdpError> {
        if endpoint.port() == 0 {
            return Err(UdpError::UnexpectedEndpoint);
        }
        let mut links = self.links.lock().map_err(|_| UdpError::Poisoned)?;
        let mut found = false;
        for link in links.values_mut() {
            if link.transport.peer() == &peer {
                link.endpoint = endpoint;
                found = true;
            }
        }
        drop(links);
        if found {
            Ok(())
        } else {
            Err(UdpError::UnknownIndex)
        }
    }
    /// Interrupts only waits owned by this carrier, without a receive timer.
    /// # Errors
    /// Poisoned carrier state.
    pub fn cancel_waits(&self) -> Result<(), UdpError> {
        cancel_links(&self.links)
    }
}

pub(super) fn cancel_links(links: &Links) -> Result<(), UdpError> {
    let links = links.lock().map_err(|_| UdpError::Poisoned)?;
    for link in links.values() {
        link.acknowledgements.cancel()?;
    }
    drop(links);
    Ok(())
}

pub(super) fn watched(links: &Links, index: u32) -> Result<bool, UdpError> {
    let links = links.lock().map_err(|_| UdpError::Poisoned)?;
    Ok(links
        .get(&index)
        .is_some_and(|link| link.acknowledgements.watched.load(Ordering::Acquire)))
}

pub(super) fn open_datagram(
    socket: &dyn PunchSocket,
    links: &Links,
    endpoint: SocketAddr,
    mut datagram: Datagram,
) -> Result<Datagram, UdpError> {
    if matches!(datagram.kind, 1 | 2) {
        return Ok(datagram);
    }
    let mut links = links.lock().map_err(|_| UdpError::Poisoned)?;
    let link = links
        .get_mut(&datagram.receiver_index)
        .ok_or(UdpError::UnknownIndex)?;
    if endpoint != link.endpoint {
        return Err(UdpError::UnexpectedEndpoint);
    }
    let (nonce, sealed) = if datagram.kind == 4 {
        if datagram.payload.len() != 8 + 1 + 8 + TAG
            || datagram.counter >= link.transport.sending_nonce()
        {
            return Err(UdpError::InvalidAck);
        }
        (
            u64::from_be_bytes(field(&datagram.payload[..8])?),
            &datagram.payload[8..],
        )
    } else {
        (datagram.counter, datagram.payload.as_slice())
    };
    if link.window.contains(nonce) {
        return replay_receipt(socket, link, endpoint, &datagram);
    }
    link.transport.set_receiving_nonce(nonce);
    let mut opened = vec![0; sealed.len()];
    let length = link.transport.open(sealed, &mut opened)?;
    opened.truncate(length);
    if datagram.kind == 4 {
        if opened.first() != Some(&4)
            || u64::from_be_bytes(field(&opened[1..])?) != datagram.counter
        {
            return Err(UdpError::InvalidAck);
        }
        link.window.record(nonce)?;
        link.highest_acked = link.highest_acked.max(Some(datagram.counter));
        let mut counters = link
            .acknowledgements
            .counters
            .lock()
            .map_err(|_| UdpError::Poisoned)?;
        counters.insert(datagram.counter);
        while counters.len() > WINDOW {
            counters.pop_first();
        }
        drop(counters);
        link.acknowledgements.changed.notify_all();
    } else {
        if opened.first() == Some(&4) {
            return Err(UdpError::InvalidDatagram);
        }
        if opened.first() == Some(&3) {
            if !matches!(opened.get(1), Some(3 | 4)) {
                return Err(UdpError::InvalidDatagram);
            }
            opened.remove(0);
        }
        link.window.record(nonce)?;
        let mut payload = vec![0; 8 + 1 + 8 + TAG];
        payload[..8].copy_from_slice(&link.transport.sending_nonce().to_be_bytes());
        let mut acknowledged = [4; 9];
        acknowledged[1..].copy_from_slice(&datagram.counter.to_be_bytes());
        link.transport.seal(&acknowledged, &mut payload[8..])?;
        let ack = Datagram {
            kind: 4,
            receiver_index: link.remote_index,
            counter: datagram.counter,
            payload,
        };
        send_datagram(socket, endpoint, &ack)?;
        link.cached_acks.push_back(CachedAck {
            counter: datagram.counter,
            digest: *blake3::hash(&datagram.payload).as_bytes(),
            datagram: ack,
            receipt: None,
        });
        while link.cached_acks.len() > WINDOW {
            link.cached_acks.pop_front();
        }
    }
    drop(links);
    datagram.payload = opened;
    Ok(datagram)
}

fn replay_receipt(
    socket: &dyn PunchSocket,
    link: &mut UdpLink,
    endpoint: SocketAddr,
    datagram: &Datagram,
) -> Result<Datagram, UdpError> {
    link.replayed = link
        .replayed
        .checked_add(1)
        .ok_or(UdpError::ReplayCountExhausted)?;
    if datagram.kind == 3 {
        let digest = *blake3::hash(&datagram.payload).as_bytes();
        if let Some(ack) = link
            .cached_acks
            .iter()
            .find(|ack| ack.counter == datagram.counter && ack.digest == digest)
        {
            send_datagram(socket, endpoint, &ack.datagram)?;
            if let Some(receipt) = &ack.receipt {
                send_datagram(socket, endpoint, receipt)?;
            }
        }
    }
    Err(UdpError::Replayed)
}