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

use std::collections::BTreeMap;
use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr};
use std::sync::Arc;
use std::time::Duration;

use bgp_integration_tests::announcement::{Announcement, withdrawal};
use bgp_integration_tests::test_peer::{TestPeer, collector_capabilities};
use bgp_integration_tests::test_server::TestServer;
use bgp_integration_tests::{peer_config, prefix, wait_until};
use bgp_packet::message::{
    AsPathAttribute, Message, MpReachNlriPathAttribute, NlriNextHop, OriginPathAttribute,
    PathAttribute, UpdateMessage,
};
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::{ExportPolicy, 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(())
}

const COLLECTOR_ASN: u32 = 65100;

/// What a collector holds: for each prefix and path identifier, the
/// route's AS path.
type Held = BTreeMap<(String, u32), Vec<u32>>;

/// Reads what the collector is told, into `held`, until `done` says so.
async fn told(
    collector: &mut TestPeer,
    held: &mut Held,
    done: impl Fn(&Held) -> bool,
) -> Result<()> {
    while !done(held) {
        let update = collector.recv_update().await?;
        let path = update
            .path_attributes
            .iter()
            .find_map(|attribute| match attribute {
                PathAttribute::AsPath(path) => Some(path.segments[0].path.clone()),
                _ => None,
            })
            .unwrap_or_default();
        for nlri in &update.withdrawn_routes {
            held.remove(&(nlri.prefix.to_string(), nlri.path_id));
        }
        for nlri in &update.announced_routes {
            held.insert((nlri.prefix.to_string(), nlri.path_id), path.clone());
        }
        // And what came and went in the multiprotocol attributes.
        for attribute in &update.path_attributes {
            match attribute {
                PathAttribute::MpReachNlri(reach) => {
                    for nlri in &reach.prefixes {
                        held.insert((nlri.prefix.to_string(), nlri.path_id), path.clone());
                    }
                }
                PathAttribute::MpUnreachNlri(unreach) => {
                    for nlri in &unreach.prefixes {
                        held.remove(&(nlri.prefix.to_string(), nlri.path_id));
                    }
                }
                _ => {}
            }
        }
    }
    Ok(())
}

/// The routes held for 203.0.113.0/24, by path identifier.
fn paths(held: &Held) -> Vec<(u32, Vec<u32>)> {
    held.iter()
        .map(|((prefix, id), path)| {
            assert_eq!(prefix, "203.0.113.0/24");
            (*id, path.clone())
        })
        .collect()
}

/// A collector that agreed to be sent several routes for a prefix, and
/// whose export policy says to send it every route, is told of every
/// peer's route for a prefix, each with a path identifier of its own that
/// stays with the route; a route that goes is withdrawn by its identifier.
#[tokio::test]
async fn test_every_route_goes_to_a_collector_told_apart_by_path_identifiers() -> Result<()> {
    let server = TestServer::start_with(LOCAL_ASN, |c| {
        // The collector is the configured peer at the IPv4 loopback
        // address; the feeders come over IPv6 loopback as peers of no
        // account.
        c.listen_endpoints
            .push(SocketAddr::new(Ipv6Addr::LOCALHOST.into(), 0));
        c.peers = vec![PeerConfig {
            add_path: Some(SendReceive::Send),
            import: ImportPolicy::permissive(),
            export: Arc::new(ExportPolicy {
                learned_routes: true,
                all_paths: true,
                ..ExportPolicy::default()
            }),
            ..peer_config(
                "collector",
                SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 1),
                COLLECTOR_ASN,
            )
        }];
        c.connect_retry = Duration::from_secs(3600);
    })
    .await?;
    let route = |path: &[u32]| Announcement::of([prefix("203.0.113.0/24").into()], path);
    let feeder = async |n: u32| -> Result<TestPeer> {
        let mut peer = TestPeer::connect(server.addrs[1]).await?;
        peer.handshake(PEER_ASN + n, Ipv4Addr::new(10, 0, 0, n as u8 + 1), 90)
            .await?;
        Ok(peer)
    };
    let mut a = feeder(0).await?;
    let mut b = feeder(1).await?;
    server.wait_for_sessions(2).await?;
    a.announce(route(&[PEER_ASN, 64510])).await?;
    b.announce(route(&[PEER_ASN + 1, 64510])).await?;
    server.wait_for_routes(2).await?;

    let mut collector = TestPeer::connect(server.addr).await?;
    collector
        .handshake_with(
            COLLECTOR_ASN,
            Ipv4Addr::new(10, 0, 0, 100),
            90,
            collector_capabilities(COLLECTOR_ASN, &SUPPORTED, false),
        )
        .await?;
    collector.ctx.add_path = SUPPORTED.to_vec();

    // Both routes, told apart.
    let mut held = Held::new();
    told(&mut collector, &mut held, |held| held.len() == 2).await?;
    assert_eq!(
        paths(&held),
        [
            (1, vec![LOCAL_ASN, PEER_ASN, 64510]),
            (2, vec![LOCAL_ASN, PEER_ASN + 1, 64510]),
        ]
    );

    // The second peer's route goes: withdrawn by its identifier, the
    // first's untouched.
    b.announce(withdrawal([prefix("203.0.113.0/24").into()]))
        .await?;
    told(&mut collector, &mut held, |held| held.len() == 1).await?;
    assert_eq!(paths(&held), [(1, vec![LOCAL_ASN, PEER_ASN, 64510])]);

    // The first peer's route changes: it goes out again with the same
    // identifier. A third peer's route gets the lowest free one.
    a.announce(route(&[PEER_ASN, 64511, 64510])).await?;
    let mut c = feeder(2).await?;
    server.wait_for_sessions(4).await?;
    c.announce(route(&[PEER_ASN + 2, 64510])).await?;
    let changed = |held: &Held| {
        held.len() == 2
            && held.get(&("203.0.113.0/24".to_string(), 1))
                == Some(&vec![LOCAL_ASN, PEER_ASN, 64511, 64510])
    };
    told(&mut collector, &mut held, changed).await?;
    assert_eq!(
        paths(&held),
        [
            (1, vec![LOCAL_ASN, PEER_ASN, 64511, 64510]),
            (2, vec![LOCAL_ASN, PEER_ASN + 2, 64510]),
        ]
    );
    Ok(())
}

/// A collector that is connected before the table arrives is told the
/// whole table, however large: the walk of the table that a new follower
/// starts with takes many batches with a table of some size, and goes on
/// while the feeders keep sending.
#[tokio::test]
async fn test_a_large_table_reaches_a_collector_that_was_there_first() -> Result<()> {
    const PREFIXES: u32 = 40_000;
    let server = TestServer::start_with(LOCAL_ASN, |c| {
        c.listen_endpoints
            .push(SocketAddr::new(Ipv6Addr::LOCALHOST.into(), 0));
        c.peers = vec![PeerConfig {
            add_path: Some(SendReceive::Send),
            import: ImportPolicy::permissive(),
            export: Arc::new(ExportPolicy {
                learned_routes: true,
                all_paths: true,
                ..ExportPolicy::default()
            }),
            ..peer_config(
                "collector",
                SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 1),
                COLLECTOR_ASN,
            )
        }];
        c.connect_retry = Duration::from_secs(3600);
    })
    .await?;
    let mut collector = TestPeer::connect(server.addr).await?;
    collector
        .handshake_with(
            COLLECTOR_ASN,
            Ipv4Addr::new(10, 0, 0, 100),
            90,
            collector_capabilities(COLLECTOR_ASN, &SUPPORTED, false),
        )
        .await?;
    collector.ctx.add_path = SUPPORTED.to_vec();
    server.wait_for_sessions(1).await?;

    // Two feeders, each sending the same prefixes a hundred to an UPDATE.
    let nth = |n: u32| {
        let [_, b, c, _] = (0x0a00_0000u32 + n * 256).to_be_bytes();
        prefix(&format!("10.{b}.{c}.0/24"))
    };
    let mut feeders = Vec::new();
    for n in 0..2u32 {
        let mut peer = TestPeer::connect(server.addrs[1]).await?;
        peer.handshake(PEER_ASN + n, Ipv4Addr::new(10, 0, 0, n as u8 + 1), 90)
            .await?;
        feeders.push(peer);
    }
    server.wait_for_sessions(3).await?;
    for start in (0..PREFIXES).step_by(100) {
        for (n, feeder) in feeders.iter_mut().enumerate() {
            feeder
                .announce(Announcement::of(
                    (start..start + 100).map(|i| nth(i).into()),
                    &[PEER_ASN + n as u32, 64510],
                ))
                .await?;
        }
    }
    let mut held = Held::new();
    told(&mut collector, &mut held, |held| {
        held.len() == 2 * PREFIXES as usize
    })
    .await?;
    let ids: std::collections::BTreeSet<u32> = held.keys().map(|(_, id)| *id).collect();
    assert_eq!(ids.into_iter().collect::<Vec<_>>(), [1, 2]);
    Ok(())
}

/// As a router at an exchange feeds a collector: one IPv6 session that
/// carries both families, with ADD-PATH for both and IPv4 routes with
/// IPv6 next hops (RFC 8950), and a table of some size from two feeders.
#[tokio::test]
async fn test_both_families_reach_a_collector_over_one_ipv6_session() -> Result<()> {
    const V4: u32 = 20_000;
    const V6: u32 = 1_000;
    let server = TestServer::start_with(LOCAL_ASN, |c| {
        c.listen_endpoints
            .push(SocketAddr::new(Ipv6Addr::LOCALHOST.into(), 0));
        c.peers = vec![PeerConfig {
            add_path: Some(SendReceive::Send),
            import: ImportPolicy::permissive(),
            export: Arc::new(ExportPolicy {
                learned_routes: true,
                all_paths: true,
                ..ExportPolicy::default()
            }),
            ..peer_config(
                "collector",
                SocketAddr::new(Ipv6Addr::LOCALHOST.into(), 1),
                COLLECTOR_ASN,
            )
        }];
        c.connect_retry = Duration::from_secs(3600);
    })
    .await?;
    let mut collector = TestPeer::connect(server.addrs[1]).await?;
    collector
        .handshake_with(
            COLLECTOR_ASN,
            Ipv4Addr::new(10, 0, 0, 100),
            90,
            collector_capabilities(COLLECTOR_ASN, &SUPPORTED, true),
        )
        .await?;
    collector.ctx.add_path = SUPPORTED.to_vec();
    server.wait_for_sessions(1).await?;
    let session = &server.registry.list()[0];
    assert!(session.extended_next_hop, "RFC 8950 agreed");
    assert_eq!(session.add_path_send.iter().count(), 2);

    let v4 = |n: u32| {
        let [_, b, c, _] = (0x0a00_0000u32 + n * 256).to_be_bytes();
        prefix(&format!("10.{b}.{c}.0/24"))
    };
    let v6 = |n: u32| prefix(&format!("2001:db8:{:x}::/48", 0x1000 + n));
    let mut feeders = Vec::new();
    for n in 0..2u32 {
        let mut peer = TestPeer::connect(server.addr).await?;
        peer.handshake(PEER_ASN + n, Ipv4Addr::new(10, 0, 0, n as u8 + 1), 90)
            .await?;
        feeders.push(peer);
    }
    server.wait_for_sessions(3).await?;
    for start in (0..V4).step_by(100) {
        for (n, feeder) in feeders.iter_mut().enumerate() {
            feeder
                .announce(Announcement::of(
                    (start..start + 100).map(|i| v4(i).into()),
                    &[PEER_ASN + n as u32, 64510],
                ))
                .await?;
        }
    }
    for start in (0..V6).step_by(100) {
        for (n, feeder) in feeders.iter_mut().enumerate() {
            let update = UpdateMessage {
                withdrawn_routes: vec![],
                path_attributes: vec![
                    PathAttribute::Origin(OriginPathAttribute::IGP),
                    AsPathAttribute::from_asns(vec![PEER_ASN + n as u32, 64510]),
                    PathAttribute::MpReachNlri(MpReachNlriPathAttribute {
                        family: AfiSafi::IPV6_UNICAST,
                        next_hop: NlriNextHop::Ipv6("2001:db8::9".parse().unwrap()),
                        prefixes: (start..start + 100).map(|i| v6(i).into()).collect(),
                    }),
                ],
                announced_routes: vec![],
            };
            feeder.send(&Message::Update(update)).await?;
        }
    }
    let mut held = Held::new();
    told(&mut collector, &mut held, |held| {
        held.len() == 2 * (V4 + V6) as usize
    })
    .await?;
    let v6_routes = held
        .keys()
        .filter(|(prefix, _)| prefix.contains(':'))
        .count();
    assert_eq!(v6_routes, 2 * V6 as usize);
    Ok(())
}