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

//! An RPKI cache for tests: an RTR server that hands out whatever VRPs and
//! ASPAs the test says there are.

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

use bgp_rpki::Vrp;
use bgp_rpki::rtr::{Pdu, Timers, VERSION_2};
use bytes::BytesMut;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::watch;
use tokio::task::JoinHandle;

const SESSION: u16 = 0x1234;

/// What the cache has to give.
#[derive(Clone, Default)]
struct Data {
    vrps: Vec<Vrp>,
    /// Customer ASes and their providers.
    aspas: Vec<(u32, Vec<u32>)>,
    serial: u32,
}

/// The cache. Stops, closing its connections, when dropped.
pub struct TestCache {
    address: SocketAddr,
    data: Arc<Mutex<Data>>,
    /// The serial, for telling connected routers there is something new.
    serial: watch::Sender<u32>,
    task: JoinHandle<()>,
}

impl TestCache {
    pub async fn start(vrps: Vec<Vrp>) -> Self {
        Self::start_with(vrps, vec![]).await
    }

    /// A cache with ASPAs too: customer ASes and their providers.
    pub async fn start_with(vrps: Vec<Vrp>, aspas: Vec<(u32, Vec<u32>)>) -> Self {
        let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).await.unwrap();
        let address = listener.local_addr().unwrap();
        let data = Arc::new(Mutex::new(Data {
            vrps,
            aspas,
            serial: 1,
        }));
        let (serial, _) = watch::channel(1);
        let task = tokio::spawn({
            let (data, serial) = (Arc::clone(&data), serial.clone());
            async move {
                // Connections end with this task: they are part of it.
                let mut connections = tokio::task::JoinSet::new();
                loop {
                    let Ok((stream, _)) = listener.accept().await else {
                        return;
                    };
                    connections.spawn(serve(stream, Arc::clone(&data), serial.subscribe()));
                }
            }
        });
        Self {
            address,
            data,
            serial,
            task,
        }
    }

    pub fn address(&self) -> SocketAddr {
        self.address
    }

    /// Replaces the VRPs and tells connected routers.
    pub fn set(&self, vrps: Vec<Vrp>) {
        self.change(|data| data.vrps = vrps);
    }

    /// Replaces the ASPAs and tells connected routers.
    pub fn set_aspas(&self, aspas: Vec<(u32, Vec<u32>)>) {
        self.change(|data| data.aspas = aspas);
    }

    fn change(&self, change: impl FnOnce(&mut Data)) {
        let serial = {
            let mut data = self.data.lock().unwrap();
            change(&mut data);
            data.serial += 1;
            data.serial
        };
        self.serial.send_replace(serial);
    }
}

impl Drop for TestCache {
    fn drop(&mut self) {
        self.task.abort();
    }
}

/// One router's connection. Asked for changes, this cache always says to
/// ask for everything: the router works out what changed.
async fn serve(mut stream: TcpStream, data: Arc<Mutex<Data>>, mut serial: watch::Receiver<u32>) {
    let mut buf = BytesMut::new();
    // The version the router speaks, once it has said something.
    let mut spoken = VERSION_2;
    loop {
        let mut out = BytesMut::new();
        tokio::select! {
            read = stream.read_buf(&mut buf) => {
                if !matches!(read, Ok(read) if read > 0) {
                    return;
                }
                while let Ok(Some((version, pdu, used))) = Pdu::decode(&buf) {
                    let _ = buf.split_to(used);
                    spoken = version;
                    match pdu {
                        Pdu::SerialQuery { .. } => Pdu::CacheReset.encode(version, &mut out),
                        Pdu::ResetQuery => {
                            let data = data.lock().unwrap().clone();
                            Pdu::CacheResponse { session_id: SESSION }.encode(version, &mut out);
                            for vrp in data.vrps {
                                Pdu::Prefix { announce: true, vrp }.encode(version, &mut out);
                            }
                            // Only version 2 has ASPAs.
                            for (customer, providers) in data.aspas {
                                if version >= VERSION_2 {
                                    Pdu::Aspa { announce: true, customer, providers }
                                        .encode(version, &mut out);
                                }
                            }
                            Pdu::EndOfData {
                                session_id: SESSION,
                                serial: data.serial,
                                timers: Some(Timers::default()),
                            }
                            .encode(version, &mut out);
                        }
                        _ => return,
                    }
                }
            }
            changed = serial.changed() => {
                if changed.is_err() {
                    return;
                }
                let serial = *serial.borrow_and_update();
                Pdu::SerialNotify { session_id: SESSION, serial }.encode(spoken, &mut out);
            }
        }
        if stream.write_all(&out).await.is_err() {
            return;
        }
    }
}