ferrix/crates/integration-tests/src/test_server.rs

use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::sync::Arc;

use bgp_fib::Fib;
use bgp_packet::message::{NlriNextHop, PathAttribute};
use bgp_primitives::ip_prefix::IpPrefix;
use bgp_rib::LocalRib;
use bgp_rib::table::{Event, PeerIndex, Reader, RibTable};
use bgp_rpki::Rpki;
use bgp_server::config::Config;
use bgp_server::error::ServerError;
use bgp_server::refusals::RefusalLog;
use bgp_server::restart::Restarts;
use bgp_server::server::BgpServer;
use bgp_server::session::SessionRegistry;
use bgp_web::{AppState, RouterInfo};
use eyre::Result;
use tokio::net::TcpListener;
use tokio::task::JoinHandle;

/// A BGP server started on an ephemeral port for testing.
/// Aborts the server task when dropped.
pub struct TestServer {
    handle: JoinHandle<Result<(), ServerError>>,
    /// The single IPv4 loopback address the server is bound to.
    pub addr: SocketAddr,
    /// The server's routing table.
    pub rib: Arc<RibTable>,
    /// Handle to the server's locally originated routes.
    pub local_rib: LocalRib,
    /// Live view of the server's established sessions.
    pub registry: SessionRegistry,
    /// What peers announced that was not taken.
    pub refusals: RefusalLog,
    /// The data for route origin validation, if the server validates.
    pub rpki: Option<Rpki>,
    /// What keeps the best routes in the kernel, if the server does.
    pub kernel: Option<Fib>,
    /// The routes kept of restarting peers (graceful restart).
    pub restarts: Arc<Restarts>,
    pub local_asn: u32,
}

impl TestServer {
    /// Start a server with `local_asn` and the default 90 s hold time.
    pub async fn start(local_asn: u32) -> Result<Self> {
        Self::start_with_hold_time(local_asn, 90).await
    }

    /// Start a server with an explicit hold time (useful for timer tests).
    pub async fn start_with_hold_time(local_asn: u32, hold_time_secs: u16) -> Result<Self> {
        Self::start_with(local_asn, |c| c.hold_time_secs = hold_time_secs).await
    }

    /// Start a server after letting `customize` edit the default test config
    /// (e.g. to add configured peers).
    pub async fn start_with(local_asn: u32, customize: impl FnOnce(&mut Config)) -> Result<Self> {
        let mut config = Config {
            router_id: Ipv4Addr::new(127, 0, 0, 1),
            cluster_id: Ipv4Addr::new(127, 0, 0, 1),
            local_asn,
            hold_time_secs: 90,
            listen_endpoints: vec![SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0)],
            peers: vec![],
            connect_retry: std::time::Duration::from_secs(5),
            collision_detect_established_state: false,
            // Most tests play a peer that connects from wherever the
            // system puts it; the ones about who may connect turn this off.
            accept_unconfigured_peers: true,
            originate: vec![],
            http_endpoint: None,
            rpki: None,
            kernel: None,
            export: Default::default(),
            control: None,
            graceful_restart: Default::default(),
        };
        customize(&mut config);

        let bound = BgpServer::bind(config).await?;
        let addr = bound
            .local_addrs()
            .into_iter()
            .next()
            .expect("at least one listener");
        let rib = Arc::clone(&bound.rib);
        let local_rib = bound.local_rib.clone();
        let registry = bound.registry.clone();
        let refusals = bound.refusals.clone();
        let rpki = bound.rpki.clone();
        let kernel = bound.kernel.clone();
        let restarts = Arc::clone(&bound.restarts);

        let handle = tokio::spawn(bound.run());
        Ok(Self {
            handle,
            addr,
            rib,
            local_rib,
            registry,
            refusals,
            rpki,
            kernel,
            restarts,
            local_asn,
        })
    }
}

/// A route in the server's table, as a test wants to look at it.
#[derive(Debug, Clone)]
pub struct ReceivedRoute {
    pub prefix: IpPrefix,
    pub peer: PeerIndex,
    /// The peer's identifier for the route; zero without ADD-PATH.
    pub path_id: u32,
    /// The route's attributes, without MP_REACH_NLRI.
    pub attributes: Vec<PathAttribute>,
    /// The next hop, if the route came in MP_REACH_NLRI.
    pub next_hop: Option<NlriNextHop>,
    /// Whether it is the best route for its prefix.
    pub best: bool,
}

/// A change of a prefix's best route.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RouteChange {
    /// The prefix has a best route, or another one than before.
    Announced { prefix: IpPrefix, peer: PeerIndex },
    /// The prefix has no route any more.
    Withdrawn { prefix: IpPrefix },
}

/// The changes of best route in a server's table from the moment this was
/// made, read from the table's journal.
pub struct RouteChanges {
    reader: Reader,
    pending: std::collections::VecDeque<RouteChange>,
}

impl RouteChanges {
    /// The next change, waiting for one if need be. Changes to one prefix
    /// that happen before they are read arrive as one.
    pub async fn recv(&mut self) -> Result<RouteChange> {
        loop {
            if let Some(change) = self.pending.pop_front() {
                return Ok(change);
            }
            self.reader.changed().await;
            self.read();
        }
    }

    fn read(&mut self) {
        let pending = &mut self.pending;
        let mut read = |event: Event<'_>| match event {
            Event::Best { prefix, peer, .. } => {
                pending.push_back(RouteChange::Announced { prefix, peer })
            }
            Event::Withdrawn { prefix } => pending.push_back(RouteChange::Withdrawn { prefix }),
            // Not given to a follower of the best routes.
            Event::Routes { .. } => {}
            Event::Reset | Event::Synced => {}
        };
        while !self.reader.poll(1024, &mut read).caught_up {}
    }
}

impl TestServer {
    /// Every route in the server's table.
    pub fn received_routes(&self) -> Vec<ReceivedRoute> {
        let mut routes = Vec::new();
        self.rib.for_each_route(|route| {
            let (attributes, next_hop) = route.attrs.attributes().expect("stored attributes parse");
            routes.push(ReceivedRoute {
                prefix: route.prefix,
                peer: route.peer,
                path_id: route.path_id,
                attributes,
                next_hop,
                best: route.best,
            });
        });
        routes
    }

    /// Follows the changes of best route from now on.
    pub fn route_changes(&self) -> RouteChanges {
        let mut changes = RouteChanges {
            reader: self.rib.reader(),
            pending: Default::default(),
        };
        // A new reader is first given the table as it is. That is not a
        // change.
        changes.read();
        changes.pending.clear();
        changes
    }

    /// Waits for the server to stop by itself and returns why. A healthy
    /// server never does.
    pub async fn stopped(&mut self) -> Result<(), ServerError> {
        (&mut self.handle)
            .await
            .expect("the server task does not panic")
    }

    /// Starts the HTTP API on an ephemeral loopback port and returns its
    /// base URL (e.g. `http://127.0.0.1:12345`). The API task is aborted
    /// when the returned guard drops.
    pub async fn serve_http(&self) -> Result<(String, HttpGuard)> {
        let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).await?;
        let base_url = format!("http://{}", listener.local_addr()?);
        let state = AppState {
            info: std::sync::Arc::new(
                RouterInfo::new(Ipv4Addr::new(127, 0, 0, 1), self.local_asn, "test")
                    .with_config("test.textpb", format!("local_asn: {}\n", self.local_asn)),
            ),
            registry: self.registry.clone(),
            local_rib: self.local_rib.clone(),
            rib: Arc::clone(&self.rib),
            refusals: self.refusals.clone(),
            rpki: self.rpki.clone(),
            kernel: self.kernel.clone(),
            restarts: Arc::clone(&self.restarts),
        };
        let handle = tokio::spawn(bgp_web::serve_on(state, listener));
        Ok((base_url, HttpGuard { handle }))
    }
}

/// Aborts the HTTP API task when dropped.
pub struct HttpGuard {
    handle: JoinHandle<std::io::Result<()>>,
}

impl Drop for HttpGuard {
    fn drop(&mut self) {
        self.handle.abort();
    }
}

impl Drop for TestServer {
    fn drop(&mut self) {
        self.handle.abort();
    }
}