//! Graceful restart (RFC 4724, RFC 8538): a peer that restarts keeps its
//! routes with us while it does, and we keep ours with it while we do.

use std::net::Ipv4Addr;
use std::time::Duration;

use bgp_integration_tests::announcement::{Announcement, announcement};
use bgp_integration_tests::test_peer::TestPeer;
use bgp_integration_tests::test_server::TestServer;
use bgp_integration_tests::{TIMEOUT, wait_until as until};
use bgp_packet::errors::{BgpError, CeaseSubcode};
use bgp_packet::message::{Message, NotificationMessage, UpdateMessage};
use bgp_packet::open::{Capability, GracefulRestartEntry};
use bgp_primitives::constants::{Afi, AfiSafi, Safi};
use bgp_server::restart::Leave;
use eyre::{Result, bail};

const PEER_ASN: u32 = 65001;
const PEER_ID: Ipv4Addr = Ipv4Addr::new(10, 0, 0, 2);
const PREFIXES: [&str; 3] = ["10.1.0.0/16", "10.2.0.0/16", "10.3.0.0/16"];

/// The graceful restart capability of a peer that speaks IPv4 and IPv6.
fn restart(restarted: bool, notification: bool, forwarding: bool, time: u16) -> Capability {
    Capability::GracefulRestart {
        restart_state: restarted,
        restart_time: time,
        notification,
        entries: [Afi::IPV4, Afi::IPV6]
            .into_iter()
            .map(|afi| GracefulRestartEntry {
                afi,
                safi: Safi::UNICAST,
                preserve_forwarding: forwarding,
            })
            .collect(),
    }
}

/// What a dual-stack peer offers, with `more` besides.
fn capabilities(more: impl IntoIterator<Item = Capability>) -> Vec<Capability> {
    let mut capabilities = vec![
        Capability::FourByteAsn { asn: PEER_ASN },
        Capability::MultiProtocol {
            afi: Afi::IPV4,
            safi: Safi::UNICAST,
        },
        Capability::MultiProtocol {
            afi: Afi::IPV6,
            safi: Safi::UNICAST,
        },
        Capability::RouteRefresh,
    ];
    capabilities.extend(more);
    capabilities
}

/// The peer's own routes.
fn own(prefixes: &[&str]) -> Announcement {
    announcement(prefixes, &[PEER_ASN])
}

fn end_of_rib(family: AfiSafi) -> UpdateMessage {
    UpdateMessage::end_of_rib(family)
}

/// The prefixes in the server's table, and how many of its routes are
/// stale.
fn table(server: &TestServer) -> (Vec<String>, usize) {
    let mut prefixes = Vec::new();
    let mut stale = 0;
    server.rib.for_each_route(|route| {
        prefixes.push(route.prefix.to_string());
        stale += usize::from(route.stale);
    });
    prefixes.sort();
    (prefixes, stale)
}

/// A peer that has announced `PREFIXES` with graceful restart as `offer`
/// says, and is Established.
async fn peer_with_routes(server: &TestServer, offer: Capability) -> Result<TestPeer> {
    let mut peer = TestPeer::connect(server.addr).await?;
    peer.handshake_with(PEER_ASN, PEER_ID, 90, capabilities([offer]))
        .await?;
    peer.announce(own(&PREFIXES)).await?;
    until("the routes are in", || server.rib.len() == 3).await?;
    Ok(peer)
}

/// The peer's connection fails; its routes are kept, stale, and used as
/// before. It comes back, sends some of them again, and its End-of-RIB:
/// the others go.
#[tokio::test]
async fn test_a_restarting_peer_keeps_its_routes_until_it_has_sent_them_again() -> Result<()> {
    let server = TestServer::start(65000).await?;
    let peer = peer_with_routes(&server, restart(false, true, true, 90)).await?;
    drop(peer);

    until("the session is down", || server.registry.list().is_empty()).await?;
    until("the routes are stale", || table(&server).1 == 3).await?;
    let kept = server.restarts.kept();
    assert_eq!(kept.len(), 1);
    assert_eq!(kept[0].1.routes, 3);
    // Still chosen: the RFC has stale routes used like any others.
    assert!(server.received_routes().iter().all(|route| route.best));
    // And shown, as the peer's, stale.
    let (base, _http) = server.serve_http().await?;
    let received: serde_json::Value = reqwest::get(format!("{base}/api/routes/received"))
        .await?
        .json()
        .await?;
    let routes = received["routes"].as_array().unwrap();
    assert_eq!(routes.len(), 3, "{received}");
    for route in routes {
        assert_eq!(route["stale"], true, "{route}");
        assert_eq!(route["peer_asn"], PEER_ASN, "{route}");
    }

    let mut peer = TestPeer::connect(server.addr).await?;
    peer.handshake_with(
        PEER_ASN,
        PEER_ID,
        90,
        capabilities([restart(true, true, true, 90)]),
    )
    .await?;
    until("the routes are taken over", || {
        server.registry.list().len() == 1 && server.restarts.kept().is_empty()
    })
    .await?;
    peer.announce(own(&PREFIXES[..2])).await?;
    until("two are fresh", || table(&server).1 == 1).await?;
    // The peer has them all, as far as its limit goes.
    assert_eq!(server.registry.list()[0].counters.routes(), 3);

    peer.announce(end_of_rib(AfiSafi::IPV4_UNICAST)).await?;
    until("the one not sent again is gone", || {
        table(&server) == (vec!["10.1.0.0/16".into(), "10.2.0.0/16".into()], 0)
    })
    .await?;
    until("it is counted out", || {
        server
            .registry
            .list()
            .first()
            .is_some_and(|session| session.counters.routes() == 2)
    })
    .await?;
    Ok(())
}

/// What ends a session decides whether the peer is restarting: a Hard
/// Reset never is, and a NOTIFICATION is only where both sides said so
/// (RFC 8538). A peer that does not offer graceful restart never is.
#[tokio::test]
async fn test_only_a_restart_keeps_the_routes() -> Result<()> {
    let cease = |subcode| NotificationMessage {
        error: BgpError::Cease(subcode),
        data: vec![],
    };
    let hard_reset = NotificationMessage {
        error: BgpError::Cease(CeaseSubcode::HardReset),
        data: vec![6, 2],
    };
    for (offer, goodbye, kept) in [
        (
            Some(restart(false, true, true, 90)),
            Some(cease(CeaseSubcode::AdminReset)),
            true,
        ),
        (
            Some(restart(false, true, true, 90)),
            Some(hard_reset.clone()),
            false,
        ),
        (
            Some(restart(false, false, true, 90)),
            Some(cease(CeaseSubcode::AdminReset)),
            false,
        ),
        (Some(restart(false, false, true, 90)), None, true),
        (Some(restart(false, true, true, 0)), None, false),
        (None, None, false),
    ] {
        let server = TestServer::start(65000).await?;
        let mut peer = TestPeer::connect(server.addr).await?;
        peer.handshake_with(PEER_ASN, PEER_ID, 90, capabilities(offer))
            .await?;
        peer.announce(own(&PREFIXES)).await?;
        until("the routes are in", || server.rib.len() == 3).await?;
        if let Some(goodbye) = &goodbye {
            peer.send(&Message::Notification(goodbye.clone())).await?;
        }
        drop(peer);
        until("the session is down", || server.registry.list().is_empty()).await?;
        if kept {
            until("the routes are stale", || table(&server).1 == 3).await?;
        } else {
            until("the routes are gone", || server.rib.is_empty()).await?;
            assert!(server.restarts.kept().is_empty());
        }
    }
    Ok(())
}

/// A peer that does not come back within its restart time has its routes
/// removed then.
#[tokio::test]
async fn test_the_routes_go_when_the_peer_does_not_come_back_in_time() -> Result<()> {
    let server = TestServer::start(65000).await?;
    let peer = peer_with_routes(&server, restart(false, true, true, 1)).await?;
    drop(peer);
    until("the routes are stale", || table(&server).1 == 3).await?;
    until("the routes are gone", || server.rib.is_empty()).await?;
    assert!(server.restarts.kept().is_empty());
    Ok(())
}

/// A peer that comes back without having kept its forwarding (F bit), or
/// without graceful restart at all, has its stale routes removed at once.
#[tokio::test]
async fn test_routes_whose_forwarding_was_lost_go_when_the_peer_is_back() -> Result<()> {
    for back_with in [Some(restart(true, true, false, 90)), None] {
        let server = TestServer::start(65000).await?;
        let peer = peer_with_routes(&server, restart(false, true, true, 90)).await?;
        drop(peer);
        until("the routes are stale", || table(&server).1 == 3).await?;

        let mut peer = TestPeer::connect(server.addr).await?;
        peer.handshake_with(PEER_ASN, PEER_ID, 90, capabilities(back_with.clone()))
            .await?;
        until("the stale routes are gone", || server.rib.is_empty()).await?;
        until("they are counted out", || {
            server
                .registry
                .list()
                .first()
                .is_some_and(|session| session.counters.routes() == 0)
        })
        .await?;
        drop(peer);
    }
    Ok(())
}

/// Our OPEN offers graceful restart, as just restarted, and a peer that
/// agreed is sent End-of-RIB markers once it has our routes.
#[tokio::test]
async fn test_we_offer_graceful_restart_and_say_when_our_routes_are_all_sent() -> Result<()> {
    let server = TestServer::start(65000).await?;
    let mut peer = TestPeer::connect(server.addr).await?;
    let open = peer
        .handshake_with(
            PEER_ASN,
            PEER_ID,
            90,
            capabilities([restart(false, true, true, 90)]),
        )
        .await?;
    let offered = open.capabilities().find_map(|capability| match capability {
        Capability::GracefulRestart {
            restart_state,
            restart_time,
            notification,
            entries,
        } => Some((
            *restart_state,
            *restart_time,
            *notification,
            entries.clone(),
        )),
        _ => None,
    });
    let Some((restarted, time, notification, entries)) = offered else {
        bail!("no graceful restart in {open:?}");
    };
    assert!(restarted, "started just now");
    assert_eq!(time, 120);
    assert!(notification);
    // Without kernel routes, our forwarding is not ours to lose.
    assert!(entries.iter().all(|entry| entry.preserve_forwarding));

    // Nothing is originated: the markers are all there is.
    let mut markers = Vec::new();
    while markers.len() < 2 {
        let update = peer.recv_update().await?;
        match update.end_of_rib_family() {
            Some(family) => markers.push(family),
            None => bail!("expected End-of-RIB, got {update:?}"),
        }
    }
    markers.sort();
    assert_eq!(markers, [AfiSafi::IPV4_UNICAST, AfiSafi::IPV6_UNICAST]);
    Ok(())
}

/// Leaving to restart closes the connection without a NOTIFICATION where
/// graceful restart was agreed, so that the peer keeps our routes; leaving
/// for good is a Hard Reset there, and a Cease elsewhere.
#[tokio::test]
async fn test_leaving_tells_each_peer_as_graceful_restart_has_it() -> Result<()> {
    let hard_reset = BgpError::Cease(CeaseSubcode::HardReset);
    let shutdown = BgpError::Cease(CeaseSubcode::AdminShutdown);
    for (leave, offer, told) in [
        (Leave::Restart, Some(restart(false, true, true, 90)), None),
        (Leave::Restart, None, Some(shutdown.clone())),
        (
            Leave::Stop,
            Some(restart(false, true, true, 90)),
            Some(hard_reset),
        ),
        (Leave::Stop, None, Some(shutdown)),
    ] {
        let server = TestServer::start(65000).await?;
        let mut peer = TestPeer::connect(server.addr).await?;
        peer.handshake_with(PEER_ASN, PEER_ID, 90, capabilities(offer.clone()))
            .await?;
        until("the session is up", || server.registry.list().len() == 1).await?;
        server.registry.leave(leave, Duration::from_secs(2)).await;
        let mut notification = None;
        // Whatever came before (End-of-RIB markers) is passed over.
        while let Ok(message) = peer.recv().await {
            if let Message::Notification(n) = message {
                notification = Some(n.error);
                break;
            }
        }
        assert_eq!(notification, told, "{leave:?} with {offer:?}");
        assert!(peer.closed_within(TIMEOUT).await);
    }
    Ok(())
}