ferrix/crates/server/src/server.rs

use std::io;
use std::net::SocketAddr;
use std::sync::Arc;

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

use bgp_rib::{AdjRibIn, LocalRib};

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::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
                .map_err(|_| ServerError::LocalRibStopped)?;
        }
        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,
        };

        let _advertiser = spawn_advertiser(
            self.local_rib.clone(),
            self.registry.clone(),
            self.config.local_asn,
        );
        let _dialers = spawn_dialers(
            &self.config.peers,
            &session_config,
            self.config.connect_retry,
        );

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

        while let Some((listen_addr, result, listener)) = accepting.next().await {
            match result {
                Ok((stream, peer_addr)) => {
                    // Inbound connections from a configured peer inherit its
                    // name, expected ASN and hold time.
                    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;
                    }
                    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));
        }

        Ok(())
    }
}

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