//! What is taken from a peer: the checks every route goes through, and the
//! limit on how many a peer may have.

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

use bgp_integration_tests::default_policy;
use bgp_integration_tests::test_peer::TestPeer;
use bgp_integration_tests::test_server::TestServer;
use bgp_packet::errors::{BgpError, CeaseSubcode};
use bgp_packet::message::{
    AsPathAttribute, CommunitiesPathAttribute, LargeCommunitiesPathAttribute, LargeCommunity,
    Message, NextHopPathAttribute, OriginPathAttribute, PathAttribute, UpdateMessage,
};
use bgp_primitives::ip_prefix::IpPrefix;
use bgp_server::config::PeerConfig;
use bgp_server::family::FamilySet;
use bgp_server::marks::Scrub;
use bgp_server::policy::{ImportPolicy, PrefixSet};
use bgp_server::session::DialState;
use eyre::Result;
use tokio::net::TcpListener;

/// AS numbers as they are on the Internet: the checks refuse the private
/// ones the other tests use.
const LOCAL_ASN: u32 = 1299;
const PEER_ASN: u32 = 3333;
const TIMEOUT: Duration = Duration::from_secs(5);

fn prefix(s: &str) -> IpPrefix {
    IpPrefix::try_from(s).unwrap()
}

/// A server with one configured peer, at the address the test connects
/// from, whose routes are judged by `import`. Returns the listener that
/// keeps the server's own dialling occupied.
async fn server_with(import: ImportPolicy) -> Result<(TestServer, TcpListener)> {
    let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).await?;
    let endpoint = listener.local_addr()?;
    let server = TestServer::start_with(LOCAL_ASN, |c| {
        c.accept_unconfigured_peers = false;
        c.connect_retry = Duration::from_millis(50);
        c.peers = vec![PeerConfig {
            graceful_restart: Default::default(),
            groups: Vec::new(),
            name: "r1".into(),
            endpoint: endpoint.into(),
            remote_asn: PEER_ASN,
            hold_time_secs: 90,
            local_address: None,
            families: FamilySet::supported(),
            add_path: None,
            import,
            role: None,
            require_role: false,
            export: Default::default(),
        }];
    })
    .await?;
    Ok((server, listener))
}

async fn established(server: &TestServer) -> Result<TestPeer> {
    let mut peer = TestPeer::connect(server.addr).await?;
    peer.handshake(PEER_ASN, Ipv4Addr::new(10, 0, 0, 2), 90)
        .await?;
    tokio::time::timeout(TIMEOUT, async {
        while server.registry.list().is_empty() {
            tokio::time::sleep(Duration::from_millis(5)).await;
        }
    })
    .await?;
    Ok(peer)
}

/// Announces IPv4 `prefixes` with the AS path `path`.
async fn announce(peer: &mut TestPeer, prefixes: &[&str], path: &[u32]) -> Result<()> {
    peer.send(&Message::Update(UpdateMessage {
        withdrawn_routes: vec![],
        path_attributes: vec![
            PathAttribute::Origin(OriginPathAttribute::IGP),
            AsPathAttribute::from_asns(path.to_vec()),
            PathAttribute::NextHop(NextHopPathAttribute(Ipv4Addr::new(192, 0, 2, 1))),
        ],
        announced_routes: prefixes.iter().map(|p| prefix(p).into()).collect(),
    }))
    .await
}

/// Waits until the server has dealt with `count` announced routes of the
/// session's, taken or refused, and returns the prefixes in its table.
async fn table_after(server: &TestServer, taken: u64, refused: u64) -> Result<BTreeSet<String>> {
    let counters = server.registry.list()[0].counters.clone();
    tokio::time::timeout(TIMEOUT, async {
        while (counters.routes(), counters.rejected()) != (taken, refused) {
            tokio::time::sleep(Duration::from_millis(5)).await;
        }
    })
    .await
    .map_err(|_| {
        eyre::eyre!(
            "the server took {} routes and refused {}; expected {taken} and {refused}",
            counters.routes(),
            counters.rejected()
        )
    })?;
    Ok(server
        .received_routes()
        .iter()
        .map(|route| route.prefix.to_string())
        .collect())
}

fn set(prefixes: &[&str]) -> BTreeSet<String> {
    prefixes.iter().map(|p| p.to_string()).collect()
}

/// With nothing configured for a peer but its address and AS, its routes
/// still go through the checks that hold for anyone: only space that is
/// routed on the Internet, nothing too specific, a path that begins with
/// the peer and has no reserved AS in it, and never one that has been
/// through us. Everything a peer sent used to be taken.
#[tokio::test]
async fn test_routes_that_no_peer_should_announce_are_refused() -> Result<()> {
    let (server, _listener) = server_with(default_policy()).await?;
    let mut peer = established(&server).await?;

    // One UPDATE with good and bad prefixes: the good ones are taken.
    announce(
        &mut peer,
        &[
            "193.0.0.0/21",
            "10.0.0.0/8",    // private
            "193.0.10.0/25", // more specific than is routed
            "0.0.0.0/0",     // default
            "193.0.16.0/24",
        ],
        &[PEER_ASN],
    )
    .await?;
    assert_eq!(
        table_after(&server, 2, 3).await?,
        set(&["193.0.0.0/21", "193.0.16.0/24"])
    );

    // Paths that are refused whatever the prefix: not the peer's own,
    // with a private AS in it, and one that has been through us.
    announce(&mut peer, &["194.0.1.0/24"], &[174, PEER_ASN]).await?;
    announce(&mut peer, &["194.0.2.0/24"], &[PEER_ASN, 64512]).await?;
    announce(&mut peer, &["194.0.3.0/24"], &[PEER_ASN, LOCAL_ASN, 174]).await?;
    assert_eq!(
        table_after(&server, 2, 6).await?,
        set(&["193.0.0.0/21", "193.0.16.0/24"])
    );
    Ok(())
}

/// A route that is refused takes the place of the one the peer had
/// announced before, which goes: the peer no longer has the old route,
/// and the new one is not acceptable.
#[tokio::test]
async fn test_a_refused_route_replaces_the_one_that_was_taken() -> Result<()> {
    let (server, _listener) = server_with(default_policy()).await?;
    let mut peer = established(&server).await?;

    announce(&mut peer, &["193.0.0.0/21", "193.0.16.0/24"], &[PEER_ASN]).await?;
    assert_eq!(table_after(&server, 2, 0).await?.len(), 2);

    announce(&mut peer, &["193.0.0.0/21"], &[PEER_ASN, 64512]).await?;
    assert_eq!(table_after(&server, 1, 1).await?, set(&["193.0.16.0/24"]));
    Ok(())
}

/// A peer with a list of what it may announce is held to it: its own
/// prefixes, originated by the AS numbers given.
#[tokio::test]
async fn test_a_peer_is_held_to_its_prefix_list_and_origins() -> Result<()> {
    let mut prefixes = PrefixSet::default();
    prefixes.allow(prefix("193.0.0.0/21"), 24);
    let (server, _listener) = server_with(ImportPolicy {
        prefixes: Some(prefixes),
        origin_asns: Some([PEER_ASN, 2121].into()),
        ..default_policy()
    })
    .await?;
    let mut peer = established(&server).await?;

    // Its own, and a more specific that the list allows.
    announce(&mut peer, &["193.0.0.0/21", "193.0.4.0/24"], &[PEER_ASN]).await?;
    // A customer of its that is listed as an origin.
    announce(&mut peer, &["193.0.6.0/24"], &[PEER_ASN, 2121]).await?;
    // Someone else's prefix, as a hijack or a leak would be.
    announce(&mut peer, &["8.8.8.0/24"], &[PEER_ASN]).await?;
    // Its own prefix, but passed on from an origin that is not listed.
    announce(&mut peer, &["193.0.7.0/24"], &[PEER_ASN, 15169]).await?;

    assert_eq!(
        table_after(&server, 3, 2).await?,
        set(&["193.0.0.0/21", "193.0.4.0/24", "193.0.6.0/24"])
    );
    Ok(())
}

/// A peer that sends more routes than it is allowed is cut off and told
/// why, and kept down for a while: neither its connections are taken nor
/// is it dialled. A session that came straight back would only fill up
/// again.
#[tokio::test]
async fn test_a_peer_over_its_limit_is_cut_off_and_kept_down() -> Result<()> {
    let (server, listener) = server_with(ImportPolicy {
        max_prefixes: Some(3),
        hold_down: Duration::from_millis(1500),
        ..default_policy()
    })
    .await?;
    // The server dials the peer as well; that connection is left unanswered
    // until the end.
    let dialled = tokio::time::timeout(TIMEOUT, listener.accept()).await??;

    let mut peer = established(&server).await?;
    // Refused routes do not count towards the limit, and replacing a
    // route does not either.
    announce(
        &mut peer,
        &["193.0.0.0/24", "193.0.1.0/24", "10.0.0.0/8"],
        &[PEER_ASN],
    )
    .await?;
    announce(&mut peer, &["193.0.1.0/24", "193.0.2.0/24"], &[PEER_ASN]).await?;
    assert_eq!(table_after(&server, 3, 1).await?.len(), 3);

    announce(&mut peer, &["193.0.3.0/24"], &[PEER_ASN]).await?;
    let notification = tokio::time::timeout(TIMEOUT, async {
        loop {
            if let Message::Notification(notification) = peer.recv().await? {
                return Ok::<_, eyre::Report>(notification);
            }
        }
    })
    .await??;
    assert_eq!(
        notification.error,
        BgpError::Cease(CeaseSubcode::MaxPrefixes)
    );
    drop(peer);
    tokio::time::timeout(TIMEOUT, async {
        while !server.received_routes().is_empty() {
            tokio::time::sleep(Duration::from_millis(5)).await;
        }
    })
    .await?;

    // Held down: a new connection gets nothing.
    let mut early = TestPeer::connect(server.addr).await?;
    assert!(
        matches!(
            tokio::time::timeout(TIMEOUT, early.recv()).await,
            Ok(Err(_))
        ),
        "a peer that is held down is not answered"
    );
    // Nor is it dialled: once the connection the server had open to it is
    // gone, the dialer waits out the hold-down and says so.
    drop(dialled);
    let reason = tokio::time::timeout(TIMEOUT, async {
        loop {
            if let Some((_, DialState::Backoff { last_error, .. })) =
                server.registry.configured_peers().into_iter().next()
                && last_error.contains("held down")
            {
                return last_error;
            }
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
    })
    .await?;
    assert!(reason.contains("limit of routes"), "{reason}");

    // When the time is up the peer is dialled again, and taken back.
    tokio::time::timeout(TIMEOUT, listener.accept()).await??;
    let mut later = TestPeer::connect(server.addr).await?;
    later
        .handshake(PEER_ASN, Ipv4Addr::new(10, 0, 0, 2), 90)
        .await?;
    Ok(())
}

/// Routes from a peer are marked with where they were learned and given
/// a preference, as the peer's `import` says; and marks that mean
/// something to us are not taken on the peer's word.
#[tokio::test]
async fn test_a_peers_routes_are_marked_and_ranked_as_configured() -> Result<()> {
    use bgp_packet::message::{CommunitiesPathAttribute, LocalPrefPathAttribute};
    use bgp_server::marks::Marks;
    let (server, _listener) = server_with(ImportPolicy {
        marks: Marks::parse(
            &["1299:*".to_string()],
            &["1299:1001".to_string()],
            &["1299:9:7".to_string()],
        )
        .unwrap(),
        local_pref: Some(250),
        ..default_policy()
    })
    .await?;
    let mut peer = established(&server).await?;
    peer.send(&Message::Update(UpdateMessage {
        withdrawn_routes: vec![],
        path_attributes: vec![
            PathAttribute::Origin(OriginPathAttribute::IGP),
            AsPathAttribute::from_asns(vec![PEER_ASN]),
            PathAttribute::NextHop(NextHopPathAttribute(Ipv4Addr::new(192, 0, 2, 1))),
            PathAttribute::Communities(CommunitiesPathAttribute(vec![(3333, 5), (1299, 2002)])),
        ],
        announced_routes: vec![prefix("193.0.0.0/21").into(), prefix("193.0.8.0/22").into()],
    }))
    .await?;
    table_after(&server, 2, 0).await?;
    let routes = server.received_routes();
    for route in &routes {
        assert!(
            route
                .attributes
                .contains(&PathAttribute::Communities(CommunitiesPathAttribute(vec![
                    (3333, 5),
                    (1299, 1001)
                ]))),
            "{route:?}"
        );
        assert!(
            route
                .attributes
                .contains(&PathAttribute::LocalPref(LocalPrefPathAttribute(250))),
            "{route:?}"
        );
        assert!(
            route.attributes.iter().any(|attribute| matches!(
                attribute,
                PathAttribute::LargeCommunities(large) if large.communities.len() == 1
            )),
            "{route:?}"
        );
    }
    assert_eq!(server.rib.attrs().len(), 1, "they share their attributes");
    Ok(())
}

/// The communities of our own AS are taken off what another network
/// sends, before the route is kept, and nothing else is; from a peer that
/// is to steer its routes with them (a customer), they are kept.
#[tokio::test]
async fn test_our_communities_are_not_taken_from_other_networks() -> Result<()> {
    let ours_and_theirs = || {
        vec![
            PathAttribute::Communities(CommunitiesPathAttribute(vec![
                (0, 13335),
                (3356, 3),
                (1299, 5),
            ])),
            PathAttribute::LargeCommunities(LargeCommunitiesPathAttribute {
                communities: vec![large(1299, 0, 0), large(1299, 31, 2), large(3356, 1, 2)],
            }),
        ]
    };
    let kept = |server: &TestServer| -> Vec<PathAttribute> {
        let routes = server.received_routes();
        routes[0]
            .attributes
            .iter()
            .filter(|attribute| {
                matches!(
                    attribute,
                    PathAttribute::Communities(_) | PathAttribute::LargeCommunities(_)
                )
            })
            .cloned()
            .collect()
    };

    let scrubbed = ImportPolicy {
        scrub: Some(Scrub {
            local_asn: LOCAL_ASN,
            zero: true,
        }),
        ..default_policy()
    };
    let (server, _listener) = server_with(scrubbed).await?;
    let mut peer = established(&server).await?;
    announce_with(&mut peer, &["193.0.0.0/21"], ours_and_theirs()).await?;
    table_after(&server, 1, 0).await?;
    assert_eq!(
        kept(&server),
        [
            PathAttribute::Communities(CommunitiesPathAttribute(vec![(3356, 3)])),
            PathAttribute::LargeCommunities(LargeCommunitiesPathAttribute {
                communities: vec![large(3356, 1, 2)],
            }),
        ]
    );

    let (server, _listener) = server_with(default_policy()).await?;
    let mut peer = established(&server).await?;
    announce_with(&mut peer, &["193.0.0.0/21"], ours_and_theirs()).await?;
    table_after(&server, 1, 0).await?;
    assert_eq!(kept(&server), ours_and_theirs());
    Ok(())
}

fn large(global_admin: u32, data_1: u32, data_2: u32) -> LargeCommunity {
    LargeCommunity {
        global_admin,
        data_1,
        data_2,
    }
}

/// Announces IPv4 `prefixes` from the peer itself, with `more` attributes.
async fn announce_with(
    peer: &mut TestPeer,
    prefixes: &[&str],
    more: Vec<PathAttribute>,
) -> Result<()> {
    let mut path_attributes = vec![
        PathAttribute::Origin(OriginPathAttribute::IGP),
        AsPathAttribute::from_asns(vec![PEER_ASN]),
        PathAttribute::NextHop(NextHopPathAttribute(Ipv4Addr::new(192, 0, 2, 1))),
    ];
    path_attributes.extend(more);
    peer.send(&Message::Update(UpdateMessage {
        withdrawn_routes: vec![],
        path_attributes,
        announced_routes: prefixes.iter().map(|p| prefix(p).into()).collect(),
    }))
    .await
}