//! End-to-end tests of ADD-PATH (RFC 7911): a peer sending several routes
//! for one prefix, told apart by path identifiers.

use std::net::{Ipv4Addr, SocketAddr};
use std::time::Duration;

use bgp_integration_tests::announcement::{Announcement, withdrawal};
use bgp_integration_tests::test_peer::TestPeer;
use bgp_integration_tests::test_server::TestServer;
use bgp_integration_tests::{peer_config, prefix, wait_until};
use bgp_packet::message::PathAttribute;
use bgp_packet::open::{AddPathEntry, Capability, OpenMessage, SendReceive};
use bgp_primitives::constants::AfiSafi;
use bgp_primitives::ip_prefix::Nlri;
use bgp_rib::{LocalRoute, LocalRouteAttrs};
use bgp_server::config::PeerConfig;
use bgp_server::family::SUPPORTED;
use bgp_server::policy::ImportPolicy;
use eyre::Result;
use serde_json::Value;

const LOCAL_ASN: u32 = 65000;
const PEER_ASN: u32 = 65001;

/// A server on which connections from this machine are those of a
/// configured peer, with ADD-PATH offered to it as `add_path` says.
async fn server_offering(add_path: Option<SendReceive>, originate: &[&str]) -> Result<TestServer> {
    let originate: Vec<LocalRoute> = originate
        .iter()
        .map(|p| LocalRoute::new(prefix(p), LocalRouteAttrs::default()))
        .collect();
    TestServer::start_with(LOCAL_ASN, |c| {
        let at = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 1);
        c.peers = vec![PeerConfig {
            add_path,
            import: ImportPolicy::permissive(),
            ..peer_config("route-server-client", at, PEER_ASN)
        }];
        c.originate = originate;
        // Keep the dialer out of the way.
        c.connect_retry = Duration::from_secs(3600);
    })
    .await
}

/// A session from a peer whose OPEN offers ADD-PATH as `offer` says, for
/// IPv4 unicast. Returns the peer and the server's OPEN.
async fn peer_offering(
    server: &TestServer,
    offer: Option<SendReceive>,
) -> Result<(TestPeer, OpenMessage)> {
    let mut capabilities = vec![Capability::FourByteAsn { asn: PEER_ASN }];
    capabilities.extend(SUPPORTED.into_iter().map(Capability::multiprotocol));
    if let Some(send_recv) = offer {
        capabilities.push(Capability::AddPath {
            entries: vec![AddPathEntry {
                afi: AfiSafi::IPV4_UNICAST.afi,
                safi: AfiSafi::IPV4_UNICAST.safi,
                send_recv,
            }],
        });
    }
    let mut peer = TestPeer::connect(server.addr).await?;
    let open = peer
        .handshake_with(PEER_ASN, Ipv4Addr::new(10, 0, 0, 1), 90, capabilities)
        .await?;
    server.wait_for_session().await?;
    Ok((peer, open))
}

fn add_path_offered(open: &OpenMessage) -> Vec<(AfiSafi, SendReceive)> {
    open.capabilities()
        .filter_map(|capability| match capability {
            Capability::AddPath { entries } => Some(entries),
            _ => None,
        })
        .flatten()
        .map(|entry| (entry.family(), entry.send_recv))
        .collect()
}

fn announcement(nlri: Nlri, as_path: &[u32]) -> Announcement {
    Announcement::of([nlri], as_path)
}

/// The routes the server holds, as (path identifier, AS path, best).
fn held(server: &TestServer) -> Vec<(u32, Vec<u32>, bool)> {
    let mut held: Vec<(u32, Vec<u32>, bool)> = server
        .received_routes()
        .into_iter()
        .map(|route| {
            let path = route
                .attributes
                .iter()
                .find_map(|attribute| match attribute {
                    PathAttribute::AsPath(path) => Some(path.segments[0].path.clone()),
                    _ => None,
                })
                .unwrap_or_default();
            (route.path_id, path, route.best)
        })
        .collect();
    held.sort();
    held
}

/// A peer that sends several routes for one prefix: each is kept, the best
/// of them is chosen, and each is withdrawn on its own.
#[tokio::test]
async fn test_several_routes_for_a_prefix_from_one_peer() -> Result<()> {
    let server = server_offering(Some(SendReceive::Receive), &[]).await?;
    let (base, _http) = server.serve_http().await?;
    let (mut peer, open) = peer_offering(&server, Some(SendReceive::Send)).await?;

    // The server offered to receive, for every family it offers.
    assert_eq!(
        add_path_offered(&open),
        SUPPORTED.map(|family| (family, SendReceive::Receive))
    );
    // In use for the one family the peer sends for, and in that direction.
    let session = &server.registry.list()[0];
    assert_eq!(
        session.add_path_receive,
        [AfiSafi::IPV4_UNICAST].into_iter().collect()
    );
    assert!(session.add_path_send.is_empty());
    peer.send_add_path = vec![AfiSafi::IPV4_UNICAST];

    let p = prefix("203.0.113.0/24");
    let path = |path_id| Nlri { prefix: p, path_id };
    peer.announce(announcement(path(1), &[PEER_ASN, 64999, 64998]))
        .await?;
    peer.announce(announcement(path(2), &[PEER_ASN, 64998]))
        .await?;
    server.wait_for_routes(2).await?;
    assert_eq!(
        held(&server),
        [
            (1, vec![PEER_ASN, 64999, 64998], false),
            (2, vec![PEER_ASN, 64998], true),
        ]
    );
    assert_eq!(server.rib.prefix_count(), 1);

    // The API tells them apart too.
    let routes: Value = reqwest::get(format!("{base}/api/routes/received"))
        .await?
        .json()
        .await?;
    let mut ids: Vec<u64> = routes["routes"]
        .as_array()
        .unwrap()
        .iter()
        .map(|route| route["path_id"].as_u64().unwrap())
        .collect();
    ids.sort_unstable();
    assert_eq!(ids, [1, 2]);
    let peers: Value = reqwest::get(format!("{base}/api/peers"))
        .await?
        .json()
        .await?;
    assert_eq!(
        peers[0]["add_path_receive"],
        serde_json::json!(["IPv4 unicast"])
    );

    // Announcing a path again replaces that path only.
    peer.announce(announcement(path(1), &[PEER_ASN])).await?;
    wait_until("the path is replaced", || held(&server)[0].1.len() == 1).await?;
    assert_eq!(
        held(&server),
        [(1, vec![PEER_ASN], true), (2, vec![PEER_ASN, 64998], false),]
    );

    // And withdrawing one leaves the other.
    peer.announce(withdrawal([path(1)])).await?;
    server.wait_for_routes(1).await?;
    assert_eq!(held(&server), [(2, vec![PEER_ASN, 64998], true)]);

    // When the session ends, all of its routes go.
    peer.announce(announcement(path(3), &[PEER_ASN, 1, 2, 3]))
        .await?;
    server.wait_for_routes(2).await?;
    drop(peer);
    server.wait_for_routes(0).await?;
    Ok(())
}

/// ADD-PATH changes how prefixes are written. It must be in use only when
/// both sides said so; a peer that offers it to a server that does not is
/// spoken to, and listened to, in plain BGP.
#[tokio::test]
async fn test_add_path_is_not_used_unless_both_sides_offer_it() -> Result<()> {
    let server = server_offering(None, &["10.100.0.0/16"]).await?;
    let (mut peer, open) = peer_offering(&server, Some(SendReceive::Both)).await?;
    assert!(add_path_offered(&open).is_empty());
    let session = &server.registry.list()[0];
    assert!(session.add_path_receive.is_empty() && session.add_path_send.is_empty());

    // Plain prefixes both ways.
    let update = peer.recv_update().await?;
    assert_eq!(update.announced_routes, vec![prefix("10.100.0.0/16")]);
    peer.announce(announcement(prefix("203.0.113.0/24").into(), &[PEER_ASN]))
        .await?;
    server.wait_for_routes(1).await?;
    assert_eq!(held(&server), [(0, vec![PEER_ASN], true)]);

    // And the other way round: offered by the server, not by the peer.
    let server = server_offering(Some(SendReceive::Both), &["10.100.0.0/16"]).await?;
    let (mut peer, open) = peer_offering(&server, None).await?;
    assert_eq!(add_path_offered(&open).len(), SUPPORTED.len());
    let update = peer.recv_update().await?;
    assert_eq!(update.announced_routes, vec![prefix("10.100.0.0/16")]);
    Ok(())
}

/// Where the peer offered to receive and the server to send, the server's
/// prefixes carry a path identifier, and a peer that reads them so
/// understands them.
#[tokio::test]
async fn test_prefixes_are_sent_with_path_identifiers_when_agreed() -> Result<()> {
    let server = server_offering(Some(SendReceive::Both), &["10.100.0.0/16"]).await?;
    let (mut peer, _) = peer_offering(&server, Some(SendReceive::Receive)).await?;
    let session = &server.registry.list()[0];
    assert_eq!(
        session.add_path_send,
        [AfiSafi::IPV4_UNICAST].into_iter().collect()
    );
    assert!(session.add_path_receive.is_empty());

    // Read as agreed: with identifiers.
    peer.ctx.add_path = vec![AfiSafi::IPV4_UNICAST];
    let update = peer.recv_update().await?;
    assert_eq!(update.announced_routes.len(), 1);
    assert_eq!(update.announced_routes[0].prefix, prefix("10.100.0.0/16"));

    // The peer did not offer to send, so the server reads its prefixes
    // plain.
    peer.announce(announcement(prefix("203.0.113.0/24").into(), &[PEER_ASN]))
        .await?;
    server.wait_for_routes(1).await?;
    assert_eq!(held(&server), [(0, vec![PEER_ASN], true)]);
    Ok(())
}