liminal-server 0.25.4

//! Blind relay streams resolved from admitted holders to recorded machines.

use crate::auth_lys::session::LysSession;
use crate::auth_lys::{AccessDecisionError, HotAction, SubjectBinding};
use crate::link_handshake::MachineIdentity;
use liminal::protocol::{Frame, FrameType, codec, validate_stream};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::sync::{Arc, Mutex, mpsc};

type Device = [u8; 32];
type Pair = (Device, Device);
const STREAM_BOUND: usize = 256;

/// `[auth.lys.relay]`: explicitly enables opaque forwarding.
#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct RelayConfig {
    /// Whether this liminal serves relay streams.
    pub enabled: bool,
}

/// A relay boundary refused an operation without opening its payload.
#[derive(Debug, thiserror::Error)]
pub enum RelayRefusal {
    /// The admitted pass did not grant the relay call.
    #[error(transparent)]
    Decision(#[from] AccessDecisionError),
    /// A named protocol, connection or capacity refusal.
    #[error("{0}")]
    Named(&'static str),
}

/// A live connection's bounded, nonblocking outbound frame queue.
pub trait RelaySink: Send + Sync + std::fmt::Debug {
    /// Enqueues one frame without waiting for the peer.
    ///
    /// # Errors
    /// A full or disconnected queue must refuse by name.
    fn send(&self, frame: Frame) -> Result<(), RelayRefusal>;
}

impl RelaySink for mpsc::SyncSender<Frame> {
    fn send(&self, frame: Frame) -> Result<(), RelayRefusal> {
        self.try_send(frame).map_err(|error| match error {
            mpsc::TrySendError::Full(_) => RelayRefusal::Named("relay_outbound_full"),
            mpsc::TrySendError::Disconnected(_) => RelayRefusal::Named("peer_not_connected"),
        })
    }
}

/// Frames and encoded wire bytes accepted by the peer's outbound queue.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize)]
pub struct RelayCounters {
    /// Forwarded `LinkMessage` frames, including handshake messages.
    pub frames_relayed: u64,
    /// Encoded bytes of those frames, including the wire header.
    pub bytes_relayed: u64,
    /// Local refusals and peer `PublishError` replies for this ordered pair.
    pub refused: u64,
}

#[derive(Debug)]
struct Endpoint {
    generation: u64,
    sink: Arc<dyn RelaySink>,
    streams: BTreeMap<Device, u32>,
}

#[derive(Debug, Default)]
struct State {
    generation: u64,
    endpoints: BTreeMap<Device, Endpoint>,
    pending: BTreeMap<Pair, u32>,
    counters: BTreeMap<Pair, RelayCounters>,
}

/// The relay owns queues and routes; it has no Noise transport or private key.
#[derive(Debug)]
pub struct Relay {
    network: String,
    recorded: BTreeMap<String, MachineIdentity>,
    state: Mutex<State>,
}

impl Relay {
    /// Uses the holder-to-machine directory shared by the link host.
    #[must_use]
    pub fn new(network: String, recorded: BTreeMap<String, MachineIdentity>) -> Arc<Self> {
        Arc::new(Self {
            network,
            recorded,
            state: Mutex::new(State::default()),
        })
    }

    /// Registers an already admitted session's live connection.
    ///
    /// # Errors
    /// An unknown holder, a duplicate live machine, poisoned state or exhausted ids.
    pub fn connect(
        self: &Arc<Self>,
        session: Arc<dyn LysSession>,
        sink: Arc<dyn RelaySink>,
    ) -> Result<RelayConnection, RelayRefusal> {
        let device = *self
            .recorded
            .get(session.holder())
            .ok_or(RelayRefusal::Named("machine_not_recorded"))?
            .ed25519_public();
        let mut state = self.lock()?;
        if state.endpoints.contains_key(&device) {
            return Err(RelayRefusal::Named("device_already_connected"));
        }
        state.generation = state
            .generation
            .checked_add(1)
            .ok_or(RelayRefusal::Named("relay_ids_exhausted"))?;
        let generation = state.generation;
        state.endpoints.insert(
            device,
            Endpoint {
                generation,
                sink,
                streams: BTreeMap::new(),
            },
        );
        drop(state);
        Ok(RelayConnection {
            relay: Arc::clone(self),
            device,
            generation,
            session,
        })
    }

    /// Reads counters for one direction, retaining counts across reconnects.
    ///
    /// # Errors
    /// The relay state was poisoned.
    pub fn counters(
        &self,
        from: MachineIdentity,
        to: MachineIdentity,
    ) -> Result<RelayCounters, RelayRefusal> {
        Ok(self
            .lock()?
            .counters
            .get(&(*from.ed25519_public(), *to.ed25519_public()))
            .copied()
            .unwrap_or_default())
    }

    fn lock(&self) -> Result<std::sync::MutexGuard<'_, State>, RelayRefusal> {
        self.state
            .lock()
            .map_err(|_| RelayRefusal::Named("relay_state_poisoned"))
    }
}

/// A generation-bound connection; dropping it removes its routes and pending work.
pub struct RelayConnection {
    relay: Arc<Relay>,
    device: Device,
    generation: u64,
    session: Arc<dyn LysSession>,
}

impl std::fmt::Debug for RelayConnection {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("RelayConnection")
            .field("generation", &self.generation)
            .finish_non_exhaustive()
    }
}

impl RelayConnection {
    /// Marks a non-control stream after deciding call on net.relay.<peer>.
    ///
    /// # Errors
    /// Invalid peer or stream, missing grant, ended connection, duplicate stream or capacity.
    pub fn mark_stream(&self, stream: u32, peer: &str, now: u64) -> Result<(), RelayRefusal> {
        validate_stream(FrameType::LinkMessage, stream)
            .map_err(|_| RelayRefusal::Named("invalid_relay_stream"))?;
        let mut device = [0; 32];
        hex::decode_to_slice(peer, &mut device)
            .map_err(|_| RelayRefusal::Named("invalid_peer_device"))?;
        if hex::encode(device) != peer || device == self.device {
            return Err(RelayRefusal::Named("invalid_peer_device"));
        }
        self.decide(device, now)?;
        let mut state = self.relay.lock()?;
        let endpoint = self.endpoint(&state)?;
        if endpoint
            .streams
            .iter()
            .any(|(to, id)| (*id == stream && *to != device) || (*to == device && *id != stream))
        {
            return Err(RelayRefusal::Named("relay_stream_in_use"));
        }
        if !endpoint.streams.contains_key(&device) && endpoint.streams.len() >= STREAM_BOUND {
            return Err(RelayRefusal::Named("relay_stream_limit"));
        }
        state
            .endpoints
            .get_mut(&self.device)
            .ok_or(RelayRefusal::Named("peer_not_connected"))?
            .streams
            .insert(device, stream);
        drop(state);
        Ok(())
    }

    /// Dispatches a marked stream before the keyed-link route can open it.
    ///
    /// # Errors
    /// An unmarked stream, unsolicited reply, ended connection, queue failure or poisoned state.
    /// A refused `LinkMessage` is answered with `PublishError` on the sender's stream.
    pub fn receive(&self, mut frame: Frame, now: u64) -> Result<(), RelayRefusal> {
        let stream = frame.stream_id();
        let mut state = self.relay.lock()?;
        let endpoint = self.endpoint(&state)?;
        let peer = *endpoint
            .streams
            .iter()
            .find(|(_, id)| **id == stream)
            .map(|(peer, _)| peer)
            .ok_or(RelayRefusal::Named("relay_stream_not_marked"))?;
        let pair = (self.device, peer);
        if matches!(frame, Frame::LinkMessage { .. }) {
            if let Err(error) = self.decide(peer, now) {
                return refuse(&mut state, pair, stream, &error.to_string());
            }
            let Some(target) = state.endpoints.get(&peer) else {
                return refuse(&mut state, pair, stream, "peer_not_connected");
            };
            let Some(&target_stream) = target.streams.get(&self.device) else {
                return refuse(&mut state, pair, stream, "peer_relay_stream_not_marked");
            };
            if state.pending.contains_key(&pair) {
                return refuse(&mut state, pair, stream, "relay_in_flight");
            }
            remap(&mut frame, target_stream)?;
            let bytes = codec::encoded_len(&frame)
                .map_err(|_| RelayRefusal::Named("invalid_relay_frame"))?;
            if bytes > crate::link_carrier::WIRE_FRAME_BOUND {
                return refuse(&mut state, pair, stream, "relay_frame_too_large");
            }
            let previous = state.counters.get(&pair).copied().unwrap_or_default();
            let frames_relayed = increment(previous.frames_relayed)?;
            let bytes_relayed = previous
                .bytes_relayed
                .checked_add(
                    u64::try_from(bytes)
                        .map_err(|_| RelayRefusal::Named("relay_counter_overflow"))?,
                )
                .ok_or(RelayRefusal::Named("relay_counter_overflow"))?;
            if let Err(error) = target.sink.send(frame) {
                return refuse(&mut state, pair, stream, &error.to_string());
            }
            state.counters.insert(
                pair,
                RelayCounters {
                    frames_relayed,
                    bytes_relayed,
                    ..previous
                },
            );
            state.pending.insert(pair, stream);
            drop(state);
            return Ok(());
        }
        let original = (peer, self.device);
        let sender_stream = *state
            .pending
            .get(&original)
            .ok_or(RelayRefusal::Named("unexpected_relay_reply"))?;
        let refused = matches!(frame, Frame::PublishError { .. });
        remap(&mut frame, sender_stream)?;
        let count = if refused {
            Some(increment(
                state.counters.entry(original).or_default().refused,
            )?)
        } else {
            None
        };
        state
            .endpoints
            .get(&peer)
            .ok_or(RelayRefusal::Named("peer_not_connected"))?
            .sink
            .send(frame)?;
        state.pending.remove(&original);
        if let Some(count) = count {
            state.counters.entry(original).or_default().refused = count;
        }
        drop(state);
        Ok(())
    }

    /// Removes this exact connection and refuses unsettled frames without buffering.
    ///
    /// # Errors
    /// A remaining sender's queue refused the disconnect notification, or state was poisoned.
    pub fn disconnect(&self) -> Result<(), RelayRefusal> {
        let mut state = self.relay.lock()?;
        if state
            .endpoints
            .get(&self.device)
            .is_none_or(|endpoint| endpoint.generation != self.generation)
        {
            return Ok(());
        }
        state.endpoints.remove(&self.device);
        let affected: Vec<_> = state
            .pending
            .iter()
            .filter(|((from, to), _)| *from == self.device || *to == self.device)
            .map(|(pair, stream)| (*pair, *stream))
            .collect();
        let mut failed = None;
        for (pair, stream) in affected {
            state.pending.remove(&pair);
            if pair.1 == self.device
                && let Err(error) = refuse(&mut state, pair, stream, "peer_not_connected")
            {
                tracing::error!(%error, "relay disconnect notification refused");
                failed = Some(error);
            }
        }
        drop(state);
        failed.map_or(Ok(()), Err)
    }

    fn endpoint<'a>(&self, state: &'a State) -> Result<&'a Endpoint, RelayRefusal> {
        state
            .endpoints
            .get(&self.device)
            .filter(|endpoint| endpoint.generation == self.generation)
            .ok_or(RelayRefusal::Named("peer_not_connected"))
    }

    fn decide(&self, peer: Device, now: u64) -> Result<(), RelayRefusal> {
        let subject = format!("net.relay.{}", hex::encode(peer));
        let binding = SubjectBinding::new(self.relay.network.clone(), subject.clone(), subject)
            .map_err(|_| RelayRefusal::Named("invalid_relay_subject"))?;
        self.session.decide(&binding, HotAction::Call, now)?;
        Ok(())
    }
}

impl Drop for RelayConnection {
    fn drop(&mut self) {
        if let Err(error) = self.disconnect() {
            tracing::error!(%error, "relay disconnect refused");
        }
    }
}

fn increment(value: u64) -> Result<u64, RelayRefusal> {
    value
        .checked_add(1)
        .ok_or(RelayRefusal::Named("relay_counter_overflow"))
}

fn refuse(state: &mut State, pair: Pair, stream: u32, reason: &str) -> Result<(), RelayRefusal> {
    let counters = state.counters.entry(pair).or_default();
    counters.refused = increment(counters.refused)?;
    state
        .endpoints
        .get(&pair.0)
        .ok_or(RelayRefusal::Named("peer_not_connected"))?
        .sink
        .send(Frame::PublishError {
            flags: 0,
            stream_id: stream,
            reason_code: 0x0302,
            message: Some(reason.to_owned()),
        })
}

const fn remap(frame: &mut Frame, stream: u32) -> Result<(), RelayRefusal> {
    match frame {
        Frame::LinkMessage { stream_id, .. }
        | Frame::PublishAck { stream_id, .. }
        | Frame::PublishError { stream_id, .. } => {
            *stream_id = stream;
            Ok(())
        }
        _ => Err(RelayRefusal::Named("invalid_relay_frame")),
    }
}