ferrix/crates/server/src/session.rs

//! Registry of established BGP sessions.
//!
//! A session registers itself when it reaches Established and is removed when
//! its [`SessionGuard`] drops, so a session that errors out is always
//! deregistered. Other parts of the server (the advertiser, the HTTP API) use
//! the registry to enumerate live sessions and queue outbound messages.

use std::collections::HashMap;
use std::net::{IpAddr, SocketAddr};
use std::sync::{Arc, RwLock};
use std::time::SystemTime;

use bgp_packet::message::Message;
use bgp_rib::PeerId;
use tokio::sync::{broadcast, mpsc};

use crate::metrics::SESSIONS_ESTABLISHED;

const EVENT_CHANNEL_CAPACITY: usize = 256;

/// Where the dialer is in bringing up a configured peer.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DialState {
    /// Not dialling: the peer connected to us, or the dialer has not started.
    Idle,
    Connecting,
    /// The last attempt failed or the session ended; waiting to retry.
    Backoff { until: SystemTime, last_error: String },
}

/// A peer from the configuration file, whether or not it is currently up.
#[derive(Debug, Clone)]
pub struct ConfiguredPeer {
    pub name: String,
    pub endpoint: SocketAddr,
    pub remote_asn: u32,
}

/// Handle to one established session. Cloning is cheap.
#[derive(Debug, Clone)]
pub struct SessionHandle {
    pub peer: PeerId,
    /// Name from the configuration, if this peer is configured.
    pub name: Option<String>,
    /// Our side of the TCP connection; used as the next hop for routes we
    /// originate to this peer.
    pub local_addr: SocketAddr,
    pub established_at: SystemTime,
    /// Whether the peer negotiated RFC 6793 four-octet ASNs. Determines how
    /// AS_PATH is encoded on the wire for this session.
    pub four_octet_asn: bool,
    outbound: mpsc::Sender<Message>,
}

#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
#[error("session with {0} has closed")]
pub struct SessionClosed(pub PeerId);

impl SessionHandle {
    pub fn new(
        peer: PeerId,
        name: Option<String>,
        local_addr: SocketAddr,
        four_octet_asn: bool,
        outbound: mpsc::Sender<Message>,
    ) -> Self {
        Self { peer, name, local_addr, established_at: SystemTime::now(), four_octet_asn, outbound }
    }

    /// Queues a message for transmission on this session. Waits for queue
    /// space if the session's writer is behind.
    pub async fn send(&self, msg: Message) -> Result<(), SessionClosed> {
        self.outbound.send(msg).await.map_err(|_| SessionClosed(self.peer))
    }
}

#[derive(Debug, Clone)]
pub enum SessionEvent {
    Up(SessionHandle),
    Down(PeerId),
}

#[derive(Clone)]
pub struct SessionRegistry {
    inner: Arc<RwLock<HashMap<PeerId, SessionHandle>>>,
    configured: Arc<Vec<ConfiguredPeer>>,
    dial_states: Arc<RwLock<HashMap<String, DialState>>>,
    events: broadcast::Sender<SessionEvent>,
}

impl Default for SessionRegistry {
    fn default() -> Self {
        Self::new()
    }
}

impl SessionRegistry {
    pub fn new() -> Self {
        Self::with_configured(vec![])
    }

    pub fn with_configured(configured: Vec<ConfiguredPeer>) -> Self {
        let (events, _) = broadcast::channel(EVENT_CHANNEL_CAPACITY);
        let dial_states = configured.iter().map(|p| (p.name.clone(), DialState::Idle)).collect();
        Self {
            inner: Arc::new(RwLock::new(HashMap::new())),
            configured: Arc::new(configured),
            dial_states: Arc::new(RwLock::new(dial_states)),
            events,
        }
    }

    /// Peers from the configuration file, each with the dialer's current
    /// state for it.
    pub fn configured_peers(&self) -> Vec<(ConfiguredPeer, DialState)> {
        let states = self.dial_states.read().expect("dial states poisoned");
        self.configured
            .iter()
            .map(|p| (p.clone(), states.get(&p.name).cloned().unwrap_or(DialState::Idle)))
            .collect()
    }

    /// The configured peer whose endpoint IP is `ip`, if any.
    pub fn configured_by_ip(&self, ip: IpAddr) -> Option<&ConfiguredPeer> {
        self.configured.iter().find(|p| p.endpoint.ip() == ip)
    }

    pub fn set_dial_state(&self, name: &str, state: DialState) {
        self.dial_states.write().expect("dial states poisoned").insert(name.to_string(), state);
    }

    /// An established session with this remote IP and ASN, regardless of
    /// which side opened the TCP connection.
    pub fn find(&self, ip: IpAddr, asn: u32) -> Option<SessionHandle> {
        self.inner
            .read()
            .expect("session registry poisoned")
            .values()
            .find(|h| h.peer.addr.ip() == ip && h.peer.asn == asn)
            .cloned()
    }

    /// Adds a session and emits [`SessionEvent::Up`]. The returned guard
    /// removes the session again when dropped.
    pub fn register(&self, handle: SessionHandle) -> SessionGuard {
        let peer = handle.peer;
        self.inner.write().expect("session registry poisoned").insert(peer, handle.clone());
        SESSIONS_ESTABLISHED.inc();
        let _ = self.events.send(SessionEvent::Up(handle));
        SessionGuard { registry: self.clone(), peer }
    }

    pub fn get(&self, peer: &PeerId) -> Option<SessionHandle> {
        self.inner.read().expect("session registry poisoned").get(peer).cloned()
    }

    /// All currently established sessions, in no particular order.
    pub fn list(&self) -> Vec<SessionHandle> {
        self.inner.read().expect("session registry poisoned").values().cloned().collect()
    }

    pub fn subscribe(&self) -> broadcast::Receiver<SessionEvent> {
        self.events.subscribe()
    }

    fn remove(&self, peer: &PeerId) {
        let removed = self.inner.write().expect("session registry poisoned").remove(peer);
        if removed.is_some() {
            SESSIONS_ESTABLISHED.dec();
            let _ = self.events.send(SessionEvent::Down(*peer));
        }
    }
}

/// Removes the session from its registry on drop.
pub struct SessionGuard {
    registry: SessionRegistry,
    peer: PeerId,
}

impl Drop for SessionGuard {
    fn drop(&mut self) {
        self.registry.remove(&self.peer);
    }
}