//! Roles and route leak prevention (RFC 9234), end to end: what the two
//! ends say they are to each other, and what becomes of routes that are
//! marked as only to go to customers.

use std::net::{Ipv4Addr, Ipv6Addr};
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, OpenMsgSubcode};
use bgp_packet::message::{
    AsPathAttribute, Message, MpReachNlriPathAttribute, NextHopPathAttribute, NlriNextHop,
    OriginPathAttribute, PathAttribute, UpdateMessage,
};
use bgp_packet::open::{BgpRole, Capability};
use bgp_primitives::constants::AfiSafi;
use bgp_primitives::ip_prefix::IpPrefix;
use bgp_rib::LocalRoute;
use bgp_server::config::PeerConfig;
use bgp_server::family::FamilySet;
use bgp_server::policy::PeerRole;
use eyre::{Result, bail};
use serde_json::Value;
use tokio::net::TcpListener;

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()
}

/// What a test holds on to: the server, and the listener its dialling
/// goes to.
struct Setup {
    server: TestServer,
    _listener: TcpListener,
}

/// A server with one configured peer at the address the test connects
/// from, which is `role` to it, and which originates 193.0.0.0/21.
async fn server_with(role: Option<PeerRole>, require_role: bool) -> Result<Setup> {
    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.originate = vec![LocalRoute::new(prefix("193.0.0.0/21"), Default::default())];
        c.peers = vec![PeerConfig {
            name: "r1".into(),
            endpoint: endpoint.into(),
            remote_asn: PEER_ASN,
            hold_time_secs: 90,
            local_address: None,
            families: FamilySet::supported(),
            add_path: None,
            import: default_policy(),
            role,
            require_role,
        }];
    })
    .await?;
    Ok(Setup {
        server,
        _listener: listener,
    })
}

/// The capabilities of an ordinary peer, and with `role` what it says it
/// is.
fn capabilities(role: Option<BgpRole>) -> Vec<Capability> {
    let mut capabilities = vec![
        Capability::FourByteAsn { asn: PEER_ASN },
        Capability::multiprotocol(AfiSafi::IPV4_UNICAST),
        Capability::multiprotocol(AfiSafi::IPV6_UNICAST),
        Capability::RouteRefresh,
    ];
    capabilities.extend(role.map(|role| Capability::BgpRole { role }));
    capabilities
}

/// A session with the server, on which the peer says it is `says`.
async fn established(server: &TestServer, says: Option<BgpRole>) -> Result<TestPeer> {
    let mut peer = TestPeer::connect(server.addr).await?;
    peer.handshake_with(PEER_ASN, Ipv4Addr::new(10, 0, 0, 2), 90, capabilities(says))
        .await?;
    tokio::time::timeout(TIMEOUT, async {
        while server.registry.list().is_empty() {
            tokio::time::sleep(Duration::from_millis(5)).await;
        }
    })
    .await?;
    Ok(peer)
}

/// Announces `prefixes` with the AS path `path` and, with `marked_by`, an
/// Only to Customer attribute with that AS number.
async fn announce(
    peer: &mut TestPeer,
    prefixes: &[&str],
    path: &[u32],
    marked_by: Option<u32>,
) -> Result<()> {
    let mut path_attributes = vec![
        PathAttribute::Origin(OriginPathAttribute::IGP),
        AsPathAttribute::from_asns(path.to_vec()),
        PathAttribute::NextHop(NextHopPathAttribute(Ipv4Addr::new(192, 0, 2, 1))),
    ];
    path_attributes.extend(marked_by.map(PathAttribute::OnlyToCustomer));
    peer.send(&Message::Update(UpdateMessage {
        withdrawn_routes: vec![],
        path_attributes,
        announced_routes: prefixes.iter().map(|p| prefix(p).into()).collect(),
    }))
    .await
}

/// Waits until the table holds `count` routes, and returns the mark each
/// has, by prefix.
async fn marks(server: &TestServer, count: usize) -> Result<Vec<(String, Option<u32>)>> {
    tokio::time::timeout(TIMEOUT, async {
        while server.received_routes().len() != count {
            tokio::time::sleep(Duration::from_millis(5)).await;
        }
    })
    .await
    .map_err(|_| {
        eyre::eyre!(
            "expected {count} routes, the table has {:?}",
            server.received_routes()
        )
    })?;
    let mut marks: Vec<(String, Option<u32>)> = server
        .received_routes()
        .iter()
        .map(|route| {
            let mark = route
                .attributes
                .iter()
                .find_map(|attribute| match attribute {
                    PathAttribute::OnlyToCustomer(asn) => Some(*asn),
                    _ => None,
                });
            (route.prefix.to_string(), mark)
        })
        .collect();
    marks.sort();
    Ok(marks)
}

/// The role the server's OPEN announces.
fn announced_role(open: &bgp_packet::open::OpenMessage) -> Option<BgpRole> {
    open.capabilities().find_map(|capability| match capability {
        Capability::BgpRole { role } => Some(*role),
        _ => None,
    })
}

/// Both ends say what they are, and it fits: the session comes up, and is
/// shown as one whose peer confirmed its role.
#[tokio::test]
async fn test_roles_that_fit_together_are_confirmed() -> Result<()> {
    let Setup { server, .. } = server_with(Some(PeerRole::Customer), false).await?;
    let (base, _http) = server.serve_http().await?;
    let mut peer = TestPeer::connect(server.addr).await?;
    let open = peer
        .handshake_with(
            PEER_ASN,
            Ipv4Addr::new(10, 0, 0, 2),
            90,
            capabilities(Some(BgpRole::Customer)),
        )
        .await?;
    // The peer is our customer: we are its provider.
    assert_eq!(announced_role(&open), Some(BgpRole::Provider));

    tokio::time::timeout(TIMEOUT, async {
        while server.registry.list().is_empty() {
            tokio::time::sleep(Duration::from_millis(5)).await;
        }
    })
    .await?;
    let peers: Value = reqwest::get(format!("{base}/api/peers"))
        .await?
        .json()
        .await?;
    assert_eq!(peers[0]["role"], "customer");
    assert_eq!(peers[0]["role_confirmed"], true);
    Ok(())
}

/// The commonest way to a leak is the two ends not agreeing on what they
/// are to each other. A peer that says it is something else than it is
/// configured as is refused, with the error RFC 9234 gives for it.
#[tokio::test]
async fn test_a_peer_that_says_it_is_something_else_is_refused() -> Result<()> {
    let Setup { server, .. } = server_with(Some(PeerRole::Customer), false).await?;
    let mut peer = TestPeer::connect(server.addr).await?;
    // It thinks it is our peer; we have it down as a customer.
    peer.send_open_with(
        PEER_ASN,
        Ipv4Addr::new(10, 0, 0, 2),
        90,
        capabilities(Some(BgpRole::Peer)),
    )
    .await?;
    peer.recv_open().await?;
    match tokio::time::timeout(TIMEOUT, peer.recv()).await?? {
        Message::Notification(n) => {
            assert_eq!(n.error, BgpError::OpenMsg(OpenMsgSubcode::RoleMismatch));
            // Subcode 11 on the wire.
            assert_eq!(n.error.code_subcode(), (2, 11));
        }
        other => bail!("expected NOTIFICATION, got {other:?}"),
    }
    assert!(server.registry.list().is_empty());
    Ok(())
}

/// A peer that says nothing does not implement RFC 9234. It is taken at
/// our word for what it is, unless the configuration requires it to say.
#[tokio::test]
async fn test_a_peer_that_does_not_say_is_refused_only_where_required() -> Result<()> {
    let Setup { server, .. } = server_with(Some(PeerRole::Customer), false).await?;
    let (base, _http) = server.serve_http().await?;
    let _peer = established(&server, None).await?;
    let peers: Value = reqwest::get(format!("{base}/api/peers"))
        .await?
        .json()
        .await?;
    assert_eq!(peers[0]["role"], "customer");
    assert_eq!(peers[0]["role_confirmed"], false);

    let Setup { server, .. } = server_with(Some(PeerRole::Customer), true).await?;
    let mut peer = TestPeer::connect(server.addr).await?;
    peer.send_open_with(PEER_ASN, Ipv4Addr::new(10, 0, 0, 2), 90, capabilities(None))
        .await?;
    peer.recv_open().await?;
    match tokio::time::timeout(TIMEOUT, peer.recv()).await?? {
        Message::Notification(n) => {
            assert_eq!(n.error, BgpError::OpenMsg(OpenMsgSubcode::RoleMismatch));
        }
        other => bail!("expected NOTIFICATION, got {other:?}"),
    }
    Ok(())
}

/// A route marked as only to go to customers has been handed down or
/// across once. Coming up from a customer it is a leak, and is refused
/// whatever else the peer's policy takes.
#[tokio::test]
async fn test_a_marked_route_from_a_customer_is_a_leak() -> Result<()> {
    let Setup { server, .. } = server_with(Some(PeerRole::Customer), false).await?;
    let mut peer = established(&server, Some(BgpRole::Customer)).await?;

    // The customer's own route, and one its other provider gave it.
    announce(&mut peer, &["194.0.0.0/24"], &[PEER_ASN], None).await?;
    announce(
        &mut peer,
        &["8.8.8.0/24"],
        &[PEER_ASN, 3356, 15169],
        Some(3356),
    )
    .await?;
    // A customer's routes may go anywhere: they are not marked on the way
    // in.
    assert_eq!(
        marks(&server, 1).await?,
        [("194.0.0.0/24".to_string(), None)]
    );
    tokio::time::timeout(TIMEOUT, async {
        while server.refusals.peers().is_empty() {
            tokio::time::sleep(Duration::from_millis(5)).await;
        }
    })
    .await?;
    let refused = &server.refusals.peers()[0].1;
    assert_eq!(refused.total(), 1);
    assert_eq!(refused.recent()[0].prefix, prefix("8.8.8.0/24"));
    assert_eq!(
        refused.recent()[0].reason.as_str(),
        "route leak (Only to Customer)"
    );
    Ok(())
}

/// From a peer a marked route is in order only if the peer marked it
/// itself, on its way to us. A mark of someone else's means the peer got
/// the route from above or from another peer.
#[tokio::test]
async fn test_a_route_marked_by_someone_else_is_a_leak_from_a_peer() -> Result<()> {
    let Setup { server, .. } = server_with(Some(PeerRole::Peer), false).await?;
    let mut peer = established(&server, Some(BgpRole::Peer)).await?;

    announce(&mut peer, &["194.0.0.0/24"], &[PEER_ASN], Some(PEER_ASN)).await?;
    announce(
        &mut peer,
        &["8.8.8.0/24"],
        &[PEER_ASN, 3356, 15169],
        Some(3356),
    )
    .await?;
    // A peer that does not mark its routes has them marked for it.
    announce(&mut peer, &["195.0.0.0/24"], &[PEER_ASN, 2121], None).await?;
    assert_eq!(
        marks(&server, 2).await?,
        [
            ("194.0.0.0/24".to_string(), Some(PEER_ASN)),
            ("195.0.0.0/24".to_string(), Some(PEER_ASN)),
        ]
    );
    assert_eq!(server.refusals.peers()[0].1.total(), 1);
    Ok(())
}

/// A route from a provider is marked on the way in, if it is not already:
/// from here it may only go down. That holds for both address families,
/// and whether or not the peer's attributes could be stored as they came.
#[tokio::test]
async fn test_routes_from_a_provider_are_marked_on_the_way_in() -> Result<()> {
    for four_octet in [true, false] {
        let Setup { server, .. } = server_with(Some(PeerRole::Provider), false).await?;
        let mut peer = TestPeer::connect(server.addr).await?;
        let mut offered = capabilities(Some(BgpRole::Provider));
        if !four_octet {
            offered.retain(|capability| !matches!(capability, Capability::FourByteAsn { .. }));
        }
        peer.handshake_with(PEER_ASN, Ipv4Addr::new(10, 0, 0, 2), 90, offered)
            .await?;

        announce(&mut peer, &["194.0.0.0/24"], &[PEER_ASN, 2121], None).await?;
        // One that someone further up has marked already keeps that mark.
        announce(
            &mut peer,
            &["8.8.8.0/24"],
            &[PEER_ASN, 3356, 15169],
            Some(3356),
        )
        .await?;
        peer.send(&Message::Update(UpdateMessage {
            withdrawn_routes: vec![],
            path_attributes: vec![
                PathAttribute::Origin(OriginPathAttribute::IGP),
                AsPathAttribute::from_asns(vec![PEER_ASN, 2121]),
                PathAttribute::MpReachNlri(MpReachNlriPathAttribute {
                    family: AfiSafi::IPV6_UNICAST,
                    next_hop: NlriNextHop::Ipv6(Ipv6Addr::new(0x2001, 0x7f8, 0, 0, 0, 0, 0, 1)),
                    prefixes: vec![prefix("2a00:1450::/32").into()],
                }),
            ],
            announced_routes: vec![],
        }))
        .await?;
        assert_eq!(
            marks(&server, 3).await?,
            [
                ("194.0.0.0/24".to_string(), Some(PEER_ASN)),
                ("2a00:1450::/32".to_string(), Some(PEER_ASN)),
                ("8.8.8.0/24".to_string(), Some(3356)),
            ],
            "four-octet AS numbers: {four_octet}"
        );
        assert!(server.refusals.peers().is_empty());
    }
    Ok(())
}

/// Our own routes are marked on their way down to a customer or across to
/// a peer, and not on their way up to a provider.
#[tokio::test]
async fn test_our_routes_are_marked_on_the_way_down_or_across() -> Result<()> {
    for (peer_is, says, marked) in [
        (PeerRole::Customer, BgpRole::Customer, true),
        (PeerRole::Peer, BgpRole::Peer, true),
        (
            PeerRole::RouteServerClient,
            BgpRole::RouteServerClient,
            true,
        ),
        (PeerRole::Provider, BgpRole::Provider, false),
        (PeerRole::RouteServer, BgpRole::RouteServer, false),
    ] {
        let Setup { server, .. } = server_with(Some(peer_is), false).await?;
        let mut peer = established(&server, Some(says)).await?;
        let update = tokio::time::timeout(TIMEOUT, peer.recv_update()).await??;
        assert_eq!(update.announced_routes, [prefix("193.0.0.0/21")]);
        let mark = update
            .path_attributes
            .iter()
            .find_map(|attribute| match attribute {
                PathAttribute::OnlyToCustomer(asn) => Some(*asn),
                _ => None,
            });
        assert_eq!(mark, marked.then_some(LOCAL_ASN), "to our {peer_is:?}");
    }
    // Nothing is said to a peer whose role is not known.
    let Setup { server, .. } = server_with(None, false).await?;
    let mut peer = established(&server, None).await?;
    let update = tokio::time::timeout(TIMEOUT, peer.recv_update()).await??;
    assert!(
        !update
            .path_attributes
            .iter()
            .any(|attribute| matches!(attribute, PathAttribute::OnlyToCustomer(_)))
    );
    Ok(())
}