liminal-server 0.25.4

//! Bounded opaque relay handoffs, woken and written on the owning connection.

use std::collections::{BTreeSet, VecDeque};
use std::sync::{Arc, Mutex, MutexGuard};

use liminal::protocol::{Frame, LinkOpenOutcome, codec};

use super::super::outbound::{OutboundError, OutboundWriter};
use super::super::state::{ConnectionProcessState, FrameAction};
use super::super::wake::ReadyWaker;
use super::ConnectionRuntime;
use crate::ServerError;
use crate::console::TopologyChange;
use crate::console::records::{Carrier, Link, LinkCounters, LinkState};
use crate::server::signals::ServerSignals;
use crate::services::relay::{RelayConnection, RelayRefusal, RelaySink};

// A burst cannot hold more than sixty-four frames or one connection's byte allowance.
const FRAME_BOUND: usize = 64;
const REFUSED_CODE: u16 = 0x0302;
const SUBJECT_PREFIX: &str = "net.relay.";

#[derive(Debug, Default)]
struct Pending {
    frames: VecDeque<Frame>,
    bytes: usize,
    closed: bool,
}

#[derive(Debug)]
struct Queue {
    pending: Mutex<Pending>,
    byte_bound: usize,
    wake: ReadyWaker,
}

impl Queue {
    fn lock(&self) -> Result<MutexGuard<'_, Pending>, RelayRefusal> {
        self.pending
            .lock()
            .map_err(|_| RelayRefusal::Named("relay_queue_poisoned"))
    }

    fn close(&self) -> Result<(), RelayRefusal> {
        let mut pending = self.lock()?;
        pending.closed = true;
        pending.frames.clear();
        pending.bytes = 0;
        drop(pending);
        Ok(())
    }

    fn pump(&self, outbound: &mut OutboundWriter) -> Result<(), PumpError> {
        let mut pending = self.lock()?;
        while let Some(frame) = pending.frames.front() {
            let bytes = codec::encoded_len(frame).map_err(OutboundError::Encode)?;
            if bytes <= outbound.capacity() && !outbound.has_room(bytes) {
                break;
            }
            outbound.enqueue_frame(frame)?;
            pending.frames.pop_front();
            pending.bytes -= bytes;
        }
        drop(pending);
        Ok(())
    }
}

impl RelaySink for Queue {
    fn send(&self, frame: Frame) -> Result<(), RelayRefusal> {
        let bytes =
            codec::encoded_len(&frame).map_err(|_| RelayRefusal::Named("invalid_relay_frame"))?;
        let mut pending = self.lock()?;
        if pending.closed {
            return Err(RelayRefusal::Named("peer_not_connected"));
        }
        let total = pending
            .bytes
            .checked_add(bytes)
            .ok_or(RelayRefusal::Named("relay_outbound_full"))?;
        if pending.frames.len() >= FRAME_BOUND || total > self.byte_bound {
            return Err(RelayRefusal::Named("relay_outbound_full"));
        }
        pending.frames.push_back(frame);
        if !self.wake.fire() {
            pending.frames.pop_back();
            return Err(RelayRefusal::Named("peer_not_connected"));
        }
        pending.bytes = total;
        drop(pending);
        Ok(())
    }
}

/// The registration is strongly owned only by the connection process.
#[derive(Debug)]
pub(in crate::server::connection) struct LiveRelay {
    registration: RelayConnection,
    queue: Arc<Queue>,
    streams: BTreeSet<u32>,
    signals: Option<ServerSignals>,
    local: String,
    relay_device: String,
    device: String,
}

impl LiveRelay {
    fn connect(
        pid: u64,
        runtime: &ConnectionRuntime,
        state: &ConnectionProcessState,
    ) -> Result<Self, RelayRefusal> {
        let bindings = runtime
            .service_bindings()
            .ok_or(RelayRefusal::Named("relay_not_configured"))?;
        let relay = bindings
            .relay()
            .ok_or(RelayRefusal::Named("relay_not_configured"))?;
        let local = bindings
            .relay_topology_local()
            .ok_or(RelayRefusal::Named("relay_topology_not_configured"))?;
        let session = state
            .lys_session
            .as_ref()
            .ok_or(RelayRefusal::Named("relay_session_not_admitted"))?;
        let host = runtime
            .link_host()
            .ok_or(RelayRefusal::Named("relay_machine_not_configured"))?;
        let device = host
            .recorded_machine(session.holder())
            .ok_or(RelayRefusal::Named("machine_not_recorded"))?;
        let queue = Arc::new(Queue {
            pending: Mutex::new(Pending::default()),
            byte_bound: runtime
                .limits()
                .max_connection_outbound_bytes
                .min(crate::link_carrier::WIRE_FRAME_BOUND),
            wake: runtime
                .ready_waker(pid)
                .ok_or(RelayRefusal::Named("peer_not_connected"))?,
        });
        let sink = Arc::clone(&queue);
        Ok(Self {
            registration: relay.connect(Arc::clone(session), sink)?,
            queue,
            streams: BTreeSet::new(),
            signals: runtime.signals.get().cloned(),
            local: local.to_owned(),
            relay_device: hex::encode(host.local().machine().ed25519_public()),
            device: hex::encode(device.ed25519_public()),
        })
    }

    fn observe(&self, state: LinkState) {
        if let Some(signals) = &self.signals {
            signals.topology(TopologyChange::Link {
                link_id: Link::to_peer(&self.local, &self.device).id,
                peer: self.device.clone(),
                carrier: Some(Carrier::Relay(self.relay_device.clone())),
                state,
                reason: (state == LinkState::Down).then(|| "peer_not_connected".to_owned()),
                counters: LinkCounters::default(),
            });
        }
    }
}

impl Drop for LiveRelay {
    fn drop(&mut self) {
        if let Err(error) = self.queue.close() {
            tracing::error!(%error, "relay queue close refused");
        }
        if let Err(error) = self.registration.disconnect() {
            tracing::error!(%error, "relay registration cleanup refused");
        }
        if !self.streams.is_empty() {
            self.observe(LinkState::Down);
        }
    }
}

pub(in crate::server::connection) fn apply_frame(
    pid: u64,
    runtime: &ConnectionRuntime,
    state: &mut ConnectionProcessState,
    frame: Frame,
) -> FrameAction {
    if !super::super::apply::rejects_unauthenticated(runtime, state, &frame)
        && handles(state, &frame)
    {
        return dispatch(pid, runtime, state, frame);
    }
    super::super::apply::apply_frame(pid, runtime, state, frame)
}

fn handles(state: &ConnectionProcessState, frame: &Frame) -> bool {
    match frame {
        Frame::ConversationOpen {
            subject, stream_id, ..
        } => {
            subject.starts_with(SUBJECT_PREFIX)
                || state
                    .relay
                    .as_ref()
                    .is_some_and(|relay| relay.streams.contains(stream_id))
        }
        Frame::LinkOpen { stream_id, .. }
        | Frame::LinkMessage { stream_id, .. }
        | Frame::PublishAck { stream_id, .. }
        | Frame::PublishError { stream_id, .. } => state
            .relay
            .as_ref()
            .is_some_and(|relay| relay.streams.contains(stream_id)),
        _ => false,
    }
}

fn dispatch(
    pid: u64,
    runtime: &ConnectionRuntime,
    state: &mut ConnectionProcessState,
    frame: Frame,
) -> FrameAction {
    let stream_id = frame.stream_id();
    let result = match frame {
        Frame::ConversationOpen {
            stream_id,
            subject,
            conversation_id,
            ..
        } => {
            let marked = mark(pid, runtime, state, stream_id, &subject);
            return match marked {
                Ok(()) => FrameAction::NoResponse,
                Err(error) => FrameAction::Respond(Frame::ConversationError {
                    flags: 0,
                    stream_id,
                    conversation_id,
                    reason_code: REFUSED_CODE,
                    message: Some(error.to_string()),
                }),
            };
        }
        Frame::LinkOpen { .. } => {
            return FrameAction::Respond(Frame::LinkOpenAck {
                flags: 0,
                stream_id,
                outcome: LinkOpenOutcome::Refused {
                    reason: "relay_stream_in_use".to_owned(),
                },
            });
        }
        frame => state
            .relay
            .as_ref()
            .ok_or(RelayRefusal::Named("relay_stream_not_marked"))
            .and_then(|relay| {
                relay
                    .registration
                    .receive(frame, super::super::apply::unix_seconds_now())
            }),
    };
    match result {
        Ok(()) => FrameAction::NoResponse,
        Err(error) => FrameAction::Respond(Frame::PublishError {
            flags: 0,
            stream_id,
            reason_code: REFUSED_CODE,
            message: Some(error.to_string()),
        }),
    }
}

fn mark(
    pid: u64,
    runtime: &ConnectionRuntime,
    state: &mut ConnectionProcessState,
    stream: u32,
    subject: &str,
) -> Result<(), RelayRefusal> {
    let peer = subject
        .strip_prefix(SUBJECT_PREFIX)
        .ok_or(RelayRefusal::Named("relay_stream_in_use"))?;
    if state
        .links
        .contains_key(stream)
        .map_err(|_| RelayRefusal::Named("link_state_poisoned"))?
    {
        return Err(RelayRefusal::Named("relay_stream_in_use"));
    }
    let newly_connected = state.relay.is_none();
    if newly_connected {
        state.relay = Some(LiveRelay::connect(pid, runtime, state)?);
    }
    let relay = state
        .relay
        .as_mut()
        .ok_or(RelayRefusal::Named("peer_not_connected"))?;
    if let Err(error) =
        relay
            .registration
            .mark_stream(stream, peer, super::super::apply::unix_seconds_now())
    {
        if newly_connected {
            state.relay.take();
        }
        return Err(error);
    }
    if relay.streams.insert(stream) && relay.streams.len() == 1 {
        relay.observe(LinkState::Up);
    }
    Ok(())
}

#[derive(Debug, thiserror::Error)]
pub(in crate::server::connection) enum PumpError {
    #[error(transparent)]
    Queue(#[from] RelayRefusal),
    #[error(transparent)]
    Outbound(#[from] OutboundError),
}

pub(in crate::server::connection) fn pump(
    state: &ConnectionProcessState,
    outbound: &mut OutboundWriter,
) -> Result<(), PumpError> {
    state
        .relay
        .as_ref()
        .map_or(Ok(()), |relay| relay.queue.pump(outbound))
}

pub(in crate::server::connection) fn has_ready(
    state: &ConnectionProcessState,
    outbound: &OutboundWriter,
) -> Result<bool, ServerError> {
    let Some(relay) = &state.relay else {
        return Ok(false);
    };
    let pending = relay
        .queue
        .lock()
        .map_err(|error| ServerError::ListenerAccept {
            message: error.to_string(),
        })?;
    pending.frames.front().map_or(Ok(false), |frame| {
        let bytes = codec::encoded_len(frame).map_err(|error| ServerError::ListenerAccept {
            message: error.to_string(),
        })?;
        Ok(bytes > outbound.capacity() || outbound.has_room(bytes))
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicU64, Ordering};

    type ResultOf = Result<(), Box<dyn std::error::Error>>;

    fn queue(byte_bound: usize) -> (Queue, Arc<AtomicU64>) {
        let fired = Arc::new(AtomicU64::new(0));
        (
            Queue {
                pending: Mutex::new(Pending::default()),
                byte_bound,
                wake: ReadyWaker::for_test(Arc::clone(&fired)),
            },
            fired,
        )
    }

    fn frame(message_id: u64) -> Frame {
        Frame::PublishAck {
            flags: 0,
            stream_id: 7,
            message_id,
            durable_position: None,
        }
    }

    #[test]
    fn a_bounded_relay_queue_wakes_and_refuses_without_dropping() -> ResultOf {
        let (queue, fired) = queue(4096);
        for id in 0..FRAME_BOUND {
            queue.send(frame(u64::try_from(id)?))?;
        }
        assert_eq!(fired.load(Ordering::SeqCst), u64::try_from(FRAME_BOUND)?);
        assert_eq!(
            queue
                .send(frame(99))
                .err()
                .ok_or("full queue accepted")?
                .to_string(),
            "relay_outbound_full"
        );
        let mut outbound = OutboundWriter::with_capacity(4096);
        queue.pump(&mut outbound)?;
        let mut bytes = Vec::new();
        outbound.drain(&mut bytes, None)?;
        for id in 0..FRAME_BOUND {
            let (actual, used) = codec::decode(&bytes)?;
            assert_eq!(actual, frame(u64::try_from(id)?));
            bytes.drain(..used);
        }
        assert!(bytes.is_empty());
        queue.send(frame(99))?;
        assert_eq!(
            fired.load(Ordering::SeqCst),
            u64::try_from(FRAME_BOUND + 1)?
        );
        Ok(())
    }

    #[test]
    fn relay_byte_pressure_holds_the_head_until_the_writer_drains() -> ResultOf {
        let size = codec::encoded_len(&frame(1))?;
        let (queue, fired) = queue(2 * size);
        queue.send(frame(1))?;
        queue.send(frame(2))?;
        assert_eq!(
            queue
                .send(frame(3))
                .err()
                .ok_or("byte bound ignored")?
                .to_string(),
            "relay_outbound_full"
        );
        let mut outbound = OutboundWriter::with_capacity(size);
        outbound.enqueue_frame(&frame(0))?;
        queue.pump(&mut outbound)?;
        assert_eq!(queue.lock()?.frames.len(), 2);
        let mut bytes = Vec::new();
        for id in 0..3 {
            outbound.drain(&mut bytes, None)?;
            assert_eq!(codec::decode(&bytes)?.0, frame(id));
            bytes.clear();
            queue.pump(&mut outbound)?;
        }
        assert_eq!(queue.lock()?.bytes, 0);
        assert_eq!(fired.load(Ordering::SeqCst), 2);
        Ok(())
    }

    #[test]
    fn a_closed_relay_queue_refuses_without_waking() -> ResultOf {
        let (queue, fired) = queue(4096);
        queue.send(frame(1))?;
        queue.close()?;
        assert_eq!(
            queue
                .send(frame(2))
                .err()
                .ok_or("closed queue accepted")?
                .to_string(),
            "peer_not_connected"
        );
        assert!(queue.lock()?.frames.is_empty());
        assert_eq!(fired.load(Ordering::SeqCst), 1);
        Ok(())
    }
}