ferrix/crates/server/src/server.rs

use std::io;
use std::net::SocketAddr;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;

use futures::stream::{FuturesUnordered, StreamExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::task::{JoinError, JoinHandle};
use tracing::{error, info, warn};

use bgp_rib::{AdjRibIn, LocalRib, RibStopped};

use crate::advertise::spawn_advertiser;
use crate::collision::{CollisionTable, Direction};
use crate::config::Config;
use crate::connection::{SessionConfig, handle_connection};
use crate::dialer::spawn_dialers;
use crate::error::ServerError;
use crate::family::FamilySet;
use crate::session::{ConfiguredPeer, SessionRegistry};

pub struct BgpServerBuilder {
    config: Config,
}

/// A `BgpServerBuilder` that has bound all its listeners and is ready to accept.
/// Call [`BgpServer::run`] to start the accept loop.
pub struct BgpServer {
    config: Config,
    listeners: Vec<(SocketAddr, TcpListener)>,
    /// Shared RIB handle — created at bind time so callers can clone it before
    /// the accept loop starts (e.g. tests that want to introspect the RIB).
    pub adj_rib_in: AdjRibIn,
    /// Routes this router originates; announced to every established peer.
    pub local_rib: LocalRib,
    /// Live view of established sessions.
    pub registry: SessionRegistry,
}

impl BgpServerBuilder {
    pub fn new(config: Config) -> Self {
        Self { config }
    }

    /// Binds all configured listen endpoints and returns a [`BgpServer`].
    /// Separating bind from run lets callers (e.g. tests) inspect the actual
    /// bound addresses before the accept loop starts.
    pub async fn bind(self) -> Result<BgpServer, ServerError> {
        let mut listeners = Vec::with_capacity(self.config.listen_endpoints.len());
        for &addr in &self.config.listen_endpoints {
            let listener = TcpListener::bind(addr)
                .await
                .map_err(|source| ServerError::Bind { addr, source })?;
            let bound_addr = listener.local_addr()?;
            info!(%bound_addr, "Listening for BGP connections");
            listeners.push((bound_addr, listener));
        }
        // Routes from the configuration go in before any session exists;
        // each session is sent the whole table when it comes up.
        let local_rib = LocalRib::spawn();
        for route in &self.config.originate {
            local_rib.announce(route.clone()).await?;
        }
        Ok(BgpServer {
            listeners,
            adj_rib_in: AdjRibIn::spawn(),
            local_rib,
            registry: SessionRegistry::with_configured(
                self.config
                    .peers
                    .iter()
                    .map(|p| ConfiguredPeer {
                        name: p.name.clone(),
                        endpoint: p.endpoint,
                        remote_asn: p.remote_asn,
                    })
                    .collect(),
            ),
            config: self.config,
        })
    }

    /// Convenience method: bind then run immediately.
    pub async fn run(self) -> Result<(), ServerError> {
        self.bind().await?.run().await
    }
}

impl BgpServer {
    pub fn config(&self) -> &Config {
        &self.config
    }

    /// The addresses actually bound (useful when the config uses port 0).
    pub fn local_addrs(&self) -> Vec<SocketAddr> {
        self.listeners.iter().map(|(addr, _)| *addr).collect()
    }

    pub async fn run(self) -> Result<(), ServerError> {
        crate::metrics::init();
        info!(
            router_id = %self.config.router_id,
            local_asn = self.config.local_asn,
            peers = self.config.peers.len(),
            "BGP server started"
        );

        let session_config = SessionConfig {
            router_id: self.config.router_id,
            local_asn: self.config.local_asn,
            hold_time_secs: self.config.hold_time_secs,
            adj_rib_in: self.adj_rib_in.clone(),
            registry: self.registry.clone(),
            collisions: CollisionTable::new(self.config.collision_detect_established_state),
            direction: Direction::Inbound,
            peer_name: None,
            expected_asn: None,
            families: FamilySet::supported(),
        };

        // The tasks the server cannot do without. None of them ends while
        // the server runs, unless it panics.
        let mut tasks: FuturesUnordered<_> = spawn_dialers(
            &self.config.peers,
            &session_config,
            self.config.connect_retry,
        )
        .into_iter()
        .map(|dialer| named("dialer", dialer))
        .collect();
        tasks.push(named(
            "advertiser",
            spawn_advertiser(
                self.local_rib.clone(),
                self.registry.clone(),
                self.config.local_asn,
            ),
        ));

        let mut accepting: FuturesUnordered<_> = self
            .listeners
            .into_iter()
            .map(|(addr, listener)| accept_one(addr, listener))
            .collect();

        loop {
            let (listen_addr, result, listener) = tokio::select! {
                Some(accepted) = accepting.next() => accepted,
                failure = failed(&self.adj_rib_in, &self.local_rib, &mut tasks) => {
                    return Err(give_up(&self.registry, failure).await);
                }
            };
            match result {
                Ok((stream, peer_addr)) => {
                    // Inbound connections from a configured peer inherit its
                    // name, expected ASN, hold time and address families.
                    let mut config = session_config.clone();
                    if let Some(peer) = self
                        .config
                        .peers
                        .iter()
                        .find(|p| p.endpoint.ip() == peer_addr.ip())
                    {
                        config.peer_name = Some(peer.name.clone());
                        config.expected_asn = Some(peer.remote_asn);
                        config.hold_time_secs = peer.hold_time_secs;
                        config.families = peer.families.clone();
                    }
                    tokio::spawn(handle_connection(stream, peer_addr, Arc::new(config)));
                }
                Err(e) => {
                    error!(%listen_addr, err = %e, "Accept error");
                }
            }
            accepting.push(accept_one(listen_addr, listener));
        }
    }
}

/// Ends the server because a part of it that it cannot do without has
/// failed. Peers are told we are going (Cease / Administrative Shutdown),
/// so that they stop using our routes now and not when their hold timers
/// expire.
async fn give_up(registry: &SessionRegistry, failure: ServerError) -> ServerError {
    error!(err = %failure, "Shutting down: the server cannot continue");
    let left = registry.shutdown(SHUTDOWN_GRACE).await;
    if left > 0 {
        warn!(sessions = left, "Some sessions did not close in time");
    }
    failure
}

/// How long sessions get to send their Cease when the server gives up.
const SHUTDOWN_GRACE: Duration = Duration::from_secs(2);

type NamedTask = Pin<Box<dyn Future<Output = (&'static str, Result<(), JoinError>)> + Send>>;

fn named(name: &'static str, task: JoinHandle<()>) -> NamedTask {
    Box::pin(async move { (name, task.await) })
}

/// Resolves when something the server depends on has stopped: a RIB's
/// actor or one of `tasks`. In a healthy server it never resolves.
///
/// Nothing else would notice. A stopped RIB is only found out by whoever
/// next uses it, and a stopped task by nobody at all: the server would go
/// on accepting sessions while storing or announcing nothing.
async fn failed(
    adj_rib_in: &AdjRibIn,
    local_rib: &LocalRib,
    tasks: &mut FuturesUnordered<NamedTask>,
) -> ServerError {
    tokio::select! {
        () = adj_rib_in.stopped() => RibStopped { rib: "Adj-RIB-In" }.into(),
        () = local_rib.stopped() => RibStopped { rib: "local RIB" }.into(),
        Some((task, result)) = tasks.next() => ServerError::TaskStopped {
            task,
            reason: match result {
                Ok(()) => "it returned".into(),
                Err(e) => e.to_string(),
            },
        },
    }
}

async fn accept_one(
    addr: SocketAddr,
    listener: TcpListener,
) -> (SocketAddr, io::Result<(TcpStream, SocketAddr)>, TcpListener) {
    let result = listener.accept().await;
    (addr, result, listener)
}