//! End-to-end tests for capability negotiation: what the two OPENs agree on
//! decides which routes are sent and accepted, and how they are encoded.

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

use bgp_integration_tests::test_peer::TestPeer;
use bgp_integration_tests::test_server::TestServer;
use bgp_packet::errors::{BgpError, OpenMsgSubcode};
use bgp_packet::message::{
    AsPathAttribute, Message, MpReachNlriPathAttribute, NextHopPathAttribute, NlriNextHop,
    OriginPathAttribute, PathAttribute, UpdateMessage,
};
use bgp_packet::open::Capability;
use bgp_primitives::constants::{Afi, AfiSafi, Safi};
use bgp_primitives::ip_prefix::IpPrefix;
use bgp_rib::{LocalRoute, LocalRouteAttrs};
use bgp_server::config::PeerConfig;
use bgp_server::family::{FamilySet, SUPPORTED};
use eyre::{Result, bail};

const LOCAL_ASN: u32 = 65000;
const PEER_ASN: u32 = 65001;
const PEER_ID: Ipv4Addr = Ipv4Addr::new(10, 0, 0, 1);
const TIMEOUT: Duration = Duration::from_secs(5);

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

fn only(family: AfiSafi) -> FamilySet {
    [family].into_iter().collect()
}

fn four_octet() -> Capability {
    Capability::FourByteAsn { asn: PEER_ASN }
}

async fn wait_for_session(server: &TestServer) -> Result<bgp_server::session::SessionHandle> {
    Ok(tokio::time::timeout(TIMEOUT, async {
        loop {
            if let Some(session) = server.registry.list().into_iter().next() {
                return session;
            }
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
    })
    .await?)
}

/// An UPDATE announcing one IPv4 route in the body and one IPv6 route in
/// MP_REACH_NLRI.
fn dual_stack_update() -> UpdateMessage {
    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::MpReachNlri(MpReachNlriPathAttribute {
                family: AfiSafi::IPV6_UNICAST,
                next_hop: NlriNextHop::Ipv6(Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 1)),
                prefixes: vec![prefix("2001:db8:1::/48").into()],
            }),
        ],
        announced_routes: vec![prefix("192.0.2.0/24").into()],
    }
}

/// The session carries the families both OPENs offered, and the server's
/// OPEN offers every family by default.
#[tokio::test]
async fn test_families_are_the_ones_both_sides_offered() -> Result<()> {
    let server = TestServer::start(LOCAL_ASN).await?;
    let mut peer = TestPeer::connect(server.addr).await?;
    let open = peer
        .handshake_with(
            PEER_ASN,
            PEER_ID,
            90,
            vec![
                four_octet(),
                Capability::multiprotocol(AfiSafi::IPV4_UNICAST),
            ],
        )
        .await?;

    for family in SUPPORTED {
        assert!(
            open.capabilities()
                .any(|c| *c == Capability::multiprotocol(family)),
            "the server's OPEN offers {family}"
        );
    }
    let session = wait_for_session(&server).await?;
    assert_eq!(session.families, only(AfiSafi::IPV4_UNICAST));
    Ok(())
}

/// A route of a family the peer did not offer is not accepted from it. The
/// rest of the UPDATE is, and the session stays up.
#[tokio::test]
async fn test_routes_of_a_family_not_negotiated_are_not_accepted() -> Result<()> {
    let server = TestServer::start(LOCAL_ASN).await?;
    let mut peer = TestPeer::connect(server.addr).await?;
    peer.handshake_with(
        PEER_ASN,
        PEER_ID,
        90,
        vec![
            four_octet(),
            Capability::multiprotocol(AfiSafi::IPV4_UNICAST),
        ],
    )
    .await?;
    wait_for_session(&server).await?;

    peer.send(&Message::Update(dual_stack_update())).await?;

    let routes = tokio::time::timeout(TIMEOUT, async {
        loop {
            let routes = server.received_routes();
            if !routes.is_empty() {
                return routes;
            }
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
    })
    .await?;
    // One UPDATE is applied as a whole, so the IPv6 route is not on its way.
    assert_eq!(routes.len(), 1);
    assert_eq!(routes[0].prefix, prefix("192.0.2.0/24"));
    assert_eq!(server.registry.list().len(), 1, "the session survives");
    Ok(())
}

/// A route of a family the peer did not offer is not sent to it.
#[tokio::test]
async fn test_routes_of_a_family_not_negotiated_are_not_sent() -> Result<()> {
    let server = TestServer::start_with(LOCAL_ASN, |c| {
        c.originate = vec![LocalRoute::new(
            prefix("10.100.0.0/16"),
            LocalRouteAttrs::default(),
        )];
    })
    .await?;

    // This peer wants IPv6 only; the server has an IPv4 route.
    let mut v6_only = TestPeer::connect(server.addr).await?;
    v6_only
        .handshake_with(
            PEER_ASN,
            PEER_ID,
            90,
            vec![
                four_octet(),
                Capability::multiprotocol(AfiSafi::IPV6_UNICAST),
            ],
        )
        .await?;
    let session = wait_for_session(&server).await?;
    assert_eq!(session.families, only(AfiSafi::IPV6_UNICAST));

    // A second peer that takes IPv4 is sent the route. The advertiser
    // handles sessions in the order they came up, so by the time this one
    // has its route the first has been dealt with.
    let mut dual = TestPeer::connect(server.addr).await?;
    dual.handshake(PEER_ASN, Ipv4Addr::new(10, 0, 0, 2), 90)
        .await?;
    let update = tokio::time::timeout(TIMEOUT, dual.recv_update()).await??;
    assert_eq!(update.announced_routes, vec![prefix("10.100.0.0/16")]);

    let nothing = tokio::time::timeout(Duration::from_millis(200), v6_only.recv_update()).await;
    assert!(nothing.is_err(), "got {nothing:?}");
    Ok(())
}

/// A configured peer is offered only the families configured for it, and a
/// peer that shares none of them is refused with the capability we needed.
#[tokio::test]
async fn test_no_family_in_common_refuses_the_session() -> Result<()> {
    let server = TestServer::start_with(LOCAL_ASN, |c| {
        // Inbound connections from this address inherit the peer's settings.
        c.peers = vec![PeerConfig {
            name: "v6-only".into(),
            endpoint: SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 1),
            remote_asn: PEER_ASN,
            hold_time_secs: 90,
            local_address: None,
            families: only(AfiSafi::IPV6_UNICAST),
            add_path: None,
        }];
        // Keep the dialer out of the way.
        c.connect_retry = Duration::from_secs(3600);
    })
    .await?;

    let mut peer = TestPeer::connect(server.addr).await?;
    peer.send_open_with(
        PEER_ASN,
        PEER_ID,
        90,
        vec![
            four_octet(),
            Capability::multiprotocol(AfiSafi::IPV4_UNICAST),
        ],
    )
    .await?;
    let open = peer.recv_open().await?;
    let offered: Vec<_> = open
        .capabilities()
        .filter(|c| matches!(c, Capability::MultiProtocol { .. }))
        .collect();
    assert_eq!(offered, [&Capability::multiprotocol(AfiSafi::IPV6_UNICAST)]);

    match tokio::time::timeout(TIMEOUT, peer.recv()).await?? {
        Message::Notification(n) => {
            assert_eq!(
                n.error,
                BgpError::OpenMsg(OpenMsgSubcode::UnsupportedCapability)
            );
            // The Multiprotocol capability for IPv6 unicast.
            assert_eq!(n.data, [1, 4, 0, 2, 0, 1]);
        }
        other => bail!("expected NOTIFICATION, got {other:?}"),
    }
    assert!(server.registry.list().is_empty());
    Ok(())
}

/// Real routers offer families we do not support. That must not stop
/// the families we do share from working.
#[tokio::test]
async fn test_a_peer_offering_an_unsupported_family_is_still_accepted() -> Result<()> {
    let server = TestServer::start(LOCAL_ASN).await?;
    let mut peer = TestPeer::connect(server.addr).await?;
    peer.handshake_with(
        PEER_ASN,
        PEER_ID,
        90,
        vec![
            four_octet(),
            Capability::multiprotocol(AfiSafi::new(Afi::L2VPN, Safi::EVPN)),
            Capability::multiprotocol(AfiSafi::IPV4_UNICAST),
        ],
    )
    .await?;

    let session = wait_for_session(&server).await?;
    assert_eq!(session.families, only(AfiSafi::IPV4_UNICAST));
    assert!(
        session.four_octet_asn,
        "capabilities after it are still read"
    );
    Ok(())
}

/// A peer that offers no capabilities at all is a plain RFC 4271 speaker:
/// IPv4 unicast, two-octet AS numbers. Our four-octet AS number reaches it
/// as AS_TRANS in AS_PATH and in full in AS4_PATH (RFC 6793 §4.2.2).
#[tokio::test]
async fn test_a_peer_without_capabilities_gets_two_octet_as_paths() -> Result<()> {
    const WIDE_ASN: u32 = 4_200_000_001;
    let server = TestServer::start_with(WIDE_ASN, |c| {
        c.originate = vec![LocalRoute::new(
            prefix("10.100.0.0/16"),
            LocalRouteAttrs::default(),
        )];
    })
    .await?;

    let mut peer = TestPeer::connect(server.addr).await?;
    peer.handshake_with(PEER_ASN, PEER_ID, 90, vec![]).await?;
    let session = wait_for_session(&server).await?;
    assert!(!session.four_octet_asn);
    assert_eq!(session.families, only(AfiSafi::IPV4_UNICAST));

    // The test peer now decodes AS_PATH two octets wide, as such a speaker
    // would; an UPDATE encoded four octets wide would not parse.
    let update = tokio::time::timeout(TIMEOUT, peer.recv_update()).await??;
    assert_eq!(update.announced_routes, vec![prefix("10.100.0.0/16")]);
    assert!(
        update
            .path_attributes
            .contains(&AsPathAttribute::from_asns(vec![23456])),
        "{update:?}"
    );
    let as4_path = update.path_attributes.iter().find_map(|a| match a {
        PathAttribute::As4Path(path) => Some(path.segments[0].path.clone()),
        _ => None,
    });
    assert_eq!(as4_path, Some(vec![WIDE_ASN]));

    // And its own two-octet paths are read correctly.
    peer.send(&Message::Update(UpdateMessage {
        withdrawn_routes: vec![],
        path_attributes: vec![
            PathAttribute::Origin(OriginPathAttribute::IGP),
            AsPathAttribute::from_asns(vec![PEER_ASN, 64512]),
            PathAttribute::NextHop(NextHopPathAttribute(Ipv4Addr::new(192, 0, 2, 1))),
        ],
        announced_routes: vec![prefix("192.0.2.0/24").into()],
    }))
    .await?;
    let routes = tokio::time::timeout(TIMEOUT, async {
        loop {
            let routes = server.received_routes();
            if !routes.is_empty() {
                return routes;
            }
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
    })
    .await?;
    assert!(
        routes[0]
            .attributes
            .contains(&AsPathAttribute::from_asns(vec![PEER_ASN, 64512]))
    );
    Ok(())
}