ferrix/crates/integration-tests/tests/basic_connection.rs

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

use bgp_integration_tests::test_peer::TestPeer;
use bgp_integration_tests::test_server::TestServer;
use bgp_packet::errors::{BgpError, CeaseSubcode, FsmSubcode, MsgHeaderSubcode, UpdateMsgSubcode};
use bgp_packet::message::{
    AsPathAttribute, Message, MpReachNlriPathAttribute, MpUnreachNlriPathAttribute,
    NextHopPathAttribute, NlriNextHop, NotificationMessage, OriginPathAttribute, PathAttribute,
    UpdateMessage,
};
use bgp_primitives::constants::{AddressFamilyId, Afi, Safi};
use bgp_primitives::ip_prefix::IpPrefix;
use bgp_rib::RibEvent;
use eyre::Result;

const TIMEOUT: Duration = Duration::from_secs(5);

/// The attributes every announcement of IPv4 routes must carry. An UPDATE
/// without them has its routes withdrawn instead (RFC 7606).
fn mandatory_attributes() -> Vec<PathAttribute> {
    let mut attributes = with_origin_and_as_path(PathAttribute::NextHop(NextHopPathAttribute(
        Ipv4Addr::new(192, 0, 2, 1),
    )));
    // In order of type code: ORIGIN, AS_PATH, NEXT_HOP.
    attributes.rotate_left(1);
    attributes
}

/// `first` followed by the attributes every announcement must carry.
/// Routes in MP_REACH_NLRI have their next hop there.
fn with_origin_and_as_path(first: PathAttribute) -> Vec<PathAttribute> {
    vec![
        first,
        PathAttribute::Origin(OriginPathAttribute::IGP),
        AsPathAttribute::from_asns(vec![65001]),
    ]
}

// ---------------------------------------------------------------------------
// Connectivity
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_server_starts() -> Result<()> {
    let server = TestServer::start(65000).await?;
    assert_ne!(server.addr.port(), 0);
    Ok(())
}

#[tokio::test]
async fn test_server_accepts_tcp_connection() -> Result<()> {
    let server = TestServer::start(65000).await?;
    let _peer = TestPeer::connect(server.addr).await?;
    Ok(())
}

#[tokio::test]
async fn test_server_accepts_multiple_connections() -> Result<()> {
    let server = TestServer::start(65000).await?;
    let mut peers = Vec::new();
    for _ in 0..5 {
        peers.push(TestPeer::connect(server.addr).await?);
    }
    assert_eq!(peers.len(), 5);
    Ok(())
}

// ---------------------------------------------------------------------------
// OpenSent state
// ---------------------------------------------------------------------------

/// Server sends an OPEN immediately after TCP is established.
#[tokio::test]
async fn test_server_sends_open() -> Result<()> {
    let server = TestServer::start(65000).await?;
    let mut peer = TestPeer::connect(server.addr).await?;

    let open = peer.recv_open().await?;
    assert_eq!(open.version, 4);
    assert_eq!(open.asn, 65000);
    Ok(())
}

/// Server sends a NOTIFICATION and closes if we send garbage instead of OPEN.
#[tokio::test]
async fn test_bad_marker_does_not_crash_server() -> Result<()> {
    let server = TestServer::start(65000).await?;

    let mut bad = TestPeer::connect(server.addr).await?;
    // Discard the server's OPEN, then send garbage.
    let _ = bad.recv_open().await?;
    bad.send_raw(&[0x00; 19]).await?;
    drop(bad);

    // New connections are still accepted.
    let _good = TestPeer::connect(server.addr).await?;
    Ok(())
}

/// Reads until a NOTIFICATION arrives, skipping KEEPALIVEs.
async fn expect_notification(peer: &mut TestPeer) -> Result<NotificationMessage> {
    tokio::time::timeout(TIMEOUT, async {
        loop {
            match peer.recv().await? {
                Message::Notification(n) => return Ok(n),
                Message::KeepAlive => {}
                other => eyre::bail!("expected NOTIFICATION, got {other:?}"),
            }
        }
    })
    .await?
}

/// A header with the given length and type, followed by `body`.
fn framed(length: u16, msg_type: u8, body: &[u8]) -> Vec<u8> {
    let mut buf = vec![0xFF; 16];
    buf.extend_from_slice(&length.to_be_bytes());
    buf.push(msg_type);
    buf.extend_from_slice(body);
    buf
}

/// What the peer sent decides the NOTIFICATION it gets back (RFC 4271 §6),
/// in every state, and the session ends.
#[tokio::test]
async fn test_malformed_messages_are_answered_with_the_matching_notification() -> Result<()> {
    let server = TestServer::start(65000).await?;
    let header = |subcode| BgpError::MsgHeader(subcode);
    let mut bad_marker = framed(19, 4, &[]);
    bad_marker[0] = 0;

    let cases: Vec<(&str, Vec<u8>, BgpError, Vec<u8>)> = vec![
        (
            "length beyond the maximum",
            framed(5000, 2, &[]),
            header(MsgHeaderSubcode::BadMessageLength),
            5000u16.to_be_bytes().to_vec(),
        ),
        (
            "length below a header",
            framed(5, 4, &[]),
            header(MsgHeaderSubcode::BadMessageLength),
            vec![0, 5],
        ),
        (
            "unknown message type",
            framed(19, 9, &[]),
            header(MsgHeaderSubcode::BadMessageType),
            vec![9],
        ),
        (
            "marker not all ones",
            bad_marker,
            header(MsgHeaderSubcode::ConnNotSynchronized),
            vec![],
        ),
        (
            "UPDATE with attributes running past its end",
            framed(23, 2, &[0, 0, 0, 40]),
            BgpError::UpdateMsg(UpdateMsgSubcode::MalformedAttrs),
            vec![],
        ),
    ];

    for (i, (what, bytes, error, data)) in cases.into_iter().enumerate() {
        // Established: a distinct router ID each time, so that a session
        // still closing does not collide with the next one.
        let mut peer = TestPeer::connect(server.addr).await?;
        peer.handshake(65001, Ipv4Addr::new(10, 0, 1, i as u8 + 1), 90)
            .await?;
        peer.send_raw(&bytes).await?;
        let notification = expect_notification(&mut peer).await?;
        assert_eq!(
            (notification.error, notification.data),
            (error.clone(), data.clone()),
            "{what}, Established"
        );

        // OpenSent: the same bytes instead of an OPEN.
        let mut peer = TestPeer::connect(server.addr).await?;
        let _ = peer.recv_open().await?;
        peer.send_raw(&bytes).await?;
        let notification = expect_notification(&mut peer).await?;
        assert_eq!(
            (notification.error, notification.data),
            (error, data),
            "{what}, OpenSent"
        );
    }

    // A well-formed message at the wrong time is an FSM error, not a
    // message error.
    let mut peer = TestPeer::connect(server.addr).await?;
    let _ = peer.recv_open().await?;
    peer.send(&Message::KeepAlive).await?;
    let notification = expect_notification(&mut peer).await?;
    assert_eq!(
        notification.error,
        BgpError::FsmError(FsmSubcode::UnexpectedOpenSent)
    );
    Ok(())
}

/// A message may arrive in pieces, with the server's own timers firing in
/// between. The server once threw away the piece it had already read when
/// that happened, so the message was silently lost (and the stream was only
/// back in step if the remainder happened to be thrown away as well).
#[tokio::test]
async fn test_message_split_across_a_keepalive_tick_is_still_one_message() -> Result<()> {
    // Hold time 3 s: the server sends a KEEPALIVE every second.
    let server = TestServer::start_with_hold_time(65000, 3).await?;
    let mut peer = TestPeer::connect(server.addr).await?;
    peer.handshake(65001, Ipv4Addr::new(10, 0, 0, 1), 3).await?;

    let prefix = IpPrefix::new(AddressFamilyId::Ipv4, vec![10], 8)?;
    let mut update = bytes::BytesMut::new();
    Message::Update(UpdateMessage {
        withdrawn_routes: vec![],
        path_attributes: mandatory_attributes(),
        announced_routes: vec![prefix.clone()],
    })
    .to_wire(&peer.ctx, &mut update)?;

    // The first half, a pause long enough for the server's keepalive timer
    // to fire, then the rest.
    peer.send_raw(&update[..10]).await?;
    tokio::time::sleep(Duration::from_millis(1300)).await;
    peer.send_raw(&update[10..]).await?;

    tokio::time::timeout(TIMEOUT, async {
        while !server
            .adj_rib_in
            .snapshot()
            .await
            .unwrap()
            .iter()
            .any(|e| e.prefix == prefix)
        {
            tokio::time::sleep(Duration::from_millis(20)).await;
        }
    })
    .await
    .map_err(|_| eyre::eyre!("the split UPDATE never reached the RIB"))?;
    assert_eq!(server.registry.list().len(), 1, "session was dropped");
    Ok(())
}

// ---------------------------------------------------------------------------
// OpenConfirm → Established
// ---------------------------------------------------------------------------

/// Full OPEN/KEEPALIVE handshake reaches Established.
#[tokio::test]
async fn test_session_established() -> Result<()> {
    let server = TestServer::start(65000).await?;
    let mut peer = TestPeer::connect(server.addr).await?;

    let server_open = peer
        .handshake(65001, Ipv4Addr::new(10, 0, 0, 1), 90)
        .await?;

    assert_eq!(server_open.version, 4);
    Ok(())
}

/// After Established, the server sends periodic KEEPALIVEs.
#[tokio::test]
async fn test_server_sends_keepalives() -> Result<()> {
    // Use a short hold time so the keepalive interval (hold/3 = 1 s) fires quickly.
    let server = TestServer::start_with_hold_time(65000, 3).await?;
    let mut peer = TestPeer::connect(server.addr).await?;
    peer.handshake(65001, Ipv4Addr::new(10, 0, 0, 1), 3).await?;

    // We should receive a KEEPALIVE within ~1 second (hold/3).
    let msg = tokio::time::timeout(Duration::from_secs(3), peer.recv()).await??;
    assert!(
        matches!(msg, Message::KeepAlive),
        "Expected KEEPALIVE, got {:?}",
        msg
    );
    Ok(())
}

/// Server closes with a NOTIFICATION when the hold timer expires.
#[tokio::test]
async fn test_hold_timer_expires() -> Result<()> {
    // hold_time = 3 s → server expects a KEEPALIVE within 3 s.
    let server = TestServer::start_with_hold_time(65000, 3).await?;
    let mut peer = TestPeer::connect(server.addr).await?;
    peer.handshake(65001, Ipv4Addr::new(10, 0, 0, 1), 3).await?;

    // Don't send any keepalives. The server will send keepalives until its own
    // hold timer fires and it sends a NOTIFICATION. Drain messages until we
    // see it (or the outer timeout fires).
    let notif = tokio::time::timeout(Duration::from_secs(6), async {
        loop {
            match peer.recv().await? {
                Message::Notification(n) => return Ok::<_, eyre::Error>(n),
                Message::KeepAlive => {} // server keepalive — keep waiting
                other => eyre::bail!("Unexpected message: {:?}", other),
            }
        }
    })
    .await??;
    assert!(
        matches!(notif.error, BgpError::HoldTimer),
        "Expected HoldTimer NOTIFICATION, got {:?}",
        notif
    );
    Ok(())
}

/// A peer-sent NOTIFICATION is handled gracefully (session closes cleanly).
#[tokio::test]
async fn test_peer_sends_notification() -> Result<()> {
    let server = TestServer::start(65000).await?;
    let mut peer = TestPeer::connect(server.addr).await?;
    peer.handshake(65001, Ipv4Addr::new(10, 0, 0, 1), 90)
        .await?;

    peer.send(&Message::Notification(NotificationMessage {
        error: BgpError::Cease(CeaseSubcode::AdminShutdown),
        data: vec![],
    }))
    .await?;

    // Server should close the connection cleanly.
    tokio::time::sleep(Duration::from_millis(100)).await;
    assert!(peer.is_closed().await);
    Ok(())
}

// ---------------------------------------------------------------------------
// UPDATE
// ---------------------------------------------------------------------------

/// An UPDATE message is accepted without crashing the server.
#[tokio::test]
async fn test_server_accepts_update() -> Result<()> {
    let server = TestServer::start(65000).await?;
    let mut peer = TestPeer::connect(server.addr).await?;
    peer.handshake(65001, Ipv4Addr::new(10, 0, 0, 1), 90)
        .await?;

    let update = Message::Update(UpdateMessage {
        withdrawn_routes: vec![],
        path_attributes: mandatory_attributes(),
        // 10.0.0.0/8 — one byte of prefix for a /8.
        announced_routes: vec![IpPrefix::new(AddressFamilyId::Ipv4, vec![10], 8)?],
    });
    peer.send(&update).await?;

    // Server should stay up — we can still exchange a keepalive.
    peer.send(&Message::KeepAlive).await?;
    Ok(())
}

/// An IPv6 UPDATE via MP_REACH_NLRI is accepted without crashing the server.
#[tokio::test]
async fn test_server_accepts_ipv6_update() -> Result<()> {
    let server = TestServer::start(65000).await?;
    let mut peer = TestPeer::connect(server.addr).await?;
    peer.handshake(65001, Ipv4Addr::new(10, 0, 0, 1), 90)
        .await?;

    // 2001:db8::/32
    let prefix = IpPrefix::new(AddressFamilyId::Ipv6, vec![0x20, 0x01, 0x0d, 0xb8], 32)?;
    let update = Message::Update(UpdateMessage {
        withdrawn_routes: vec![],
        path_attributes: with_origin_and_as_path(PathAttribute::MpReachNlri(
            MpReachNlriPathAttribute {
                afi: Afi::IPV6,
                safi: Safi::UNICAST,
                next_hop: NlriNextHop::Ipv6("2001:db8::1".parse().unwrap()),
                prefixes: vec![prefix],
            },
        )),
        announced_routes: vec![],
    });
    peer.send(&update).await?;

    // Server should stay up — we can still exchange a keepalive.
    peer.send(&Message::KeepAlive).await?;
    Ok(())
}

// ---------------------------------------------------------------------------
// Adj-RIB-In introspection
// ---------------------------------------------------------------------------

/// IPv4 prefix announced via the legacy NLRI field appears in the RIB.
#[tokio::test]
async fn test_rib_records_ipv4_announce() -> Result<()> {
    let server = TestServer::start(65000).await?;
    let mut events = server.adj_rib_in.subscribe();

    let mut peer = TestPeer::connect(server.addr).await?;
    peer.handshake(65001, Ipv4Addr::new(10, 0, 0, 1), 90)
        .await?;

    let prefix = IpPrefix::new(AddressFamilyId::Ipv4, vec![10], 8)?;
    peer.send(&Message::Update(UpdateMessage {
        withdrawn_routes: vec![],
        path_attributes: mandatory_attributes(),
        announced_routes: vec![prefix.clone()],
    }))
    .await?;

    let event = tokio::time::timeout(Duration::from_secs(1), events.recv()).await??;
    assert!(
        matches!(&event, RibEvent::Announced { prefix: p, .. } if p == &prefix),
        "Expected Announced for 10.0.0.0/8, got {event:?}",
    );

    let snapshot = server.adj_rib_in.snapshot().await.unwrap();
    assert!(
        snapshot.iter().any(|e| e.prefix == prefix),
        "10.0.0.0/8 missing from RIB snapshot",
    );
    Ok(())
}

/// Withdrawn IPv4 prefix is removed from the RIB.
#[tokio::test]
async fn test_rib_records_ipv4_withdraw() -> Result<()> {
    let server = TestServer::start(65000).await?;
    let mut events = server.adj_rib_in.subscribe();

    let mut peer = TestPeer::connect(server.addr).await?;
    peer.handshake(65001, Ipv4Addr::new(10, 0, 0, 1), 90)
        .await?;

    let prefix = IpPrefix::new(AddressFamilyId::Ipv4, vec![10], 8)?;

    // Announce first.
    peer.send(&Message::Update(UpdateMessage {
        withdrawn_routes: vec![],
        path_attributes: mandatory_attributes(),
        announced_routes: vec![prefix.clone()],
    }))
    .await?;
    tokio::time::timeout(Duration::from_secs(1), events.recv()).await??;

    // Now withdraw.
    peer.send(&Message::Update(UpdateMessage {
        withdrawn_routes: vec![prefix.clone()],
        path_attributes: vec![],
        announced_routes: vec![],
    }))
    .await?;

    let event = tokio::time::timeout(Duration::from_secs(1), events.recv()).await??;
    assert!(
        matches!(&event, RibEvent::Withdrawn { prefix: p, .. } if p == &prefix),
        "Expected Withdrawn for 10.0.0.0/8, got {event:?}",
    );

    let snapshot = server.adj_rib_in.snapshot().await.unwrap();
    assert!(
        !snapshot.iter().any(|e| e.prefix == prefix),
        "10.0.0.0/8 should have been removed from RIB snapshot",
    );
    Ok(())
}

/// IPv6 prefix announced via MP_REACH_NLRI appears in the RIB.
#[tokio::test]
async fn test_rib_records_ipv6_announce() -> Result<()> {
    let server = TestServer::start(65000).await?;
    let mut events = server.adj_rib_in.subscribe();

    let mut peer = TestPeer::connect(server.addr).await?;
    peer.handshake(65001, Ipv4Addr::new(10, 0, 0, 1), 90)
        .await?;

    let prefix = IpPrefix::new(AddressFamilyId::Ipv6, vec![0x20, 0x01, 0x0d, 0xb8], 32)?;
    peer.send(&Message::Update(UpdateMessage {
        withdrawn_routes: vec![],
        path_attributes: with_origin_and_as_path(PathAttribute::MpReachNlri(
            MpReachNlriPathAttribute {
                afi: Afi::IPV6,
                safi: Safi::UNICAST,
                next_hop: NlriNextHop::Ipv6("2001:db8::1".parse().unwrap()),
                prefixes: vec![prefix.clone()],
            },
        )),
        announced_routes: vec![],
    }))
    .await?;

    let event = tokio::time::timeout(Duration::from_secs(1), events.recv()).await??;
    assert!(
        matches!(&event, RibEvent::Announced { prefix: p, .. } if p == &prefix),
        "Expected Announced for 2001:db8::/32, got {event:?}",
    );

    let snapshot = server.adj_rib_in.snapshot().await.unwrap();
    assert!(
        snapshot.iter().any(|e| e.prefix == prefix),
        "2001:db8::/32 missing from RIB snapshot",
    );
    Ok(())
}

/// IPv6 prefix withdrawn via MP_UNREACH_NLRI is removed from the RIB.
#[tokio::test]
async fn test_rib_records_ipv6_withdraw() -> Result<()> {
    let server = TestServer::start(65000).await?;
    let mut events = server.adj_rib_in.subscribe();

    let mut peer = TestPeer::connect(server.addr).await?;
    peer.handshake(65001, Ipv4Addr::new(10, 0, 0, 1), 90)
        .await?;

    let prefix = IpPrefix::new(AddressFamilyId::Ipv6, vec![0x20, 0x01, 0x0d, 0xb8], 32)?;

    // Announce first.
    peer.send(&Message::Update(UpdateMessage {
        withdrawn_routes: vec![],
        path_attributes: with_origin_and_as_path(PathAttribute::MpReachNlri(
            MpReachNlriPathAttribute {
                afi: Afi::IPV6,
                safi: Safi::UNICAST,
                next_hop: NlriNextHop::Ipv6("2001:db8::1".parse().unwrap()),
                prefixes: vec![prefix.clone()],
            },
        )),
        announced_routes: vec![],
    }))
    .await?;
    tokio::time::timeout(Duration::from_secs(1), events.recv()).await??;

    // Now withdraw via MP_UNREACH_NLRI.
    peer.send(&Message::Update(UpdateMessage {
        withdrawn_routes: vec![],
        path_attributes: vec![PathAttribute::MpUnreachNlri(MpUnreachNlriPathAttribute {
            afi: Afi::IPV6,
            safi: Safi::UNICAST,
            prefixes: vec![prefix.clone()],
        })],
        announced_routes: vec![],
    }))
    .await?;

    let event = tokio::time::timeout(Duration::from_secs(1), events.recv()).await??;
    assert!(
        matches!(&event, RibEvent::Withdrawn { prefix: p, .. } if p == &prefix),
        "Expected Withdrawn for 2001:db8::/32, got {event:?}",
    );

    let snapshot = server.adj_rib_in.snapshot().await.unwrap();
    assert!(
        !snapshot.iter().any(|e| e.prefix == prefix),
        "2001:db8::/32 should have been removed from RIB snapshot",
    );
    Ok(())
}