ferrix/crates/server/src/fsm.rs

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

use bgp_packet::errors::{BgpError, CeaseSubcode, FsmSubcode, OpenMsgSubcode};
use bgp_packet::message::{BGP_HEADER_LEN, Message, NotificationMessage, PathAttribute};
use bgp_packet::open::{Capability, OpenMessage, OpenOption};
use bgp_primitives::constants::{AddressFamilyId, SubsequentAfi};
use bgp_primitives::parser::ParserContext;
use bgp_rib::{AdjRibIn, PeerId, RibEntry};
use bytes::BytesMut;
use eyre::{Result, bail, eyre};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
use tokio::sync::{mpsc, oneshot};
use tokio::time::{Instant, interval_at, sleep_until};
use tracing::{info, warn};

use bgp_packet::message::{AS_TRANS, BGP4_VERSION};

use crate::collision::{CollisionTable, Direction};
use crate::metrics::{MESSAGES_RECEIVED, MESSAGES_SENT, SESSIONS_ESTABLISHED_TOTAL, message_type_name};
use crate::session::{SessionHandle, SessionRegistry};

/// Number of outbound messages that may be queued for a session before
/// senders are back-pressured.
const OUTBOUND_QUEUE_CAPACITY: usize = 1024;

/// Minimal per-session config extracted from the server config.
#[derive(Clone)]
pub struct SessionConfig {
    pub router_id: Ipv4Addr,
    pub local_asn: u32,
    pub hold_time_secs: u16,
    pub adj_rib_in: AdjRibIn,
    pub registry: SessionRegistry,
    /// Shared table used to resolve connection collisions (RFC 4271 §6.8).
    pub collisions: CollisionTable,
    /// Which side opened this TCP connection.
    pub direction: Direction,
    /// Configured name of the peer, if known.
    pub peer_name: Option<String>,
    /// If set, the peer's OPEN must carry this ASN or the session is
    /// refused with a Bad Peer AS notification.
    pub expected_asn: Option<u32>,
}

/// Drive the BGP FSM for a single accepted TCP connection.
///
/// States visited: OpenSent → OpenConfirm → Established.
pub async fn run_session(
    stream: TcpStream,
    peer_addr: SocketAddr,
    config: SessionConfig,
) -> Result<()> {
    let local_addr = stream.local_addr()?;
    let (mut reader, mut writer) = stream.into_split();

    // --- OpenSent -----------------------------------------------------------
    // Send our OPEN immediately after the TCP connection is established.
    let our_open = build_open(&config);
    send_message(&mut writer, &our_open, &ParserContext::default()).await?;

    // RFC 4271 §8.1: use a large initial hold timer (4 minutes) before we
    // know the negotiated value.
    let peer_msg = recv_with_timeout(
        &mut reader,
        &ParserContext::default(),
        Duration::from_secs(240),
    )
    .await;

    let peer_open = match peer_msg {
        Ok(Message::Open(o)) => o,
        Ok(other) => {
            send_notification(
                &mut writer,
                BgpError::FsmError(FsmSubcode::UnexpectedOpenSent),
                &ParserContext::default(),
            )
            .await;
            bail!("Expected OPEN in OpenSent state, got {:?}", other);
        }
        Err(e) => {
            send_notification(
                &mut writer,
                BgpError::FsmError(FsmSubcode::UnexpectedOpenSent),
                &ParserContext::default(),
            )
            .await;
            return Err(e);
        }
    };

    if let Err(bgp_err) = validate_open(&peer_open, config.router_id) {
        send_notification(&mut writer, bgp_err, &ParserContext::default()).await;
        bail!("Peer {} sent invalid OPEN", peer_addr);
    }

    let negotiated_hold_time = config.hold_time_secs.min(peer_open.hold_time);
    let peer_four_byte_asn = peer_open.options.iter().any(|o| {
        matches!(o, OpenOption::Capabilities(caps)
            if caps.iter().any(|c| matches!(c, Capability::FourByteAsn { .. })))
    });
    let ctx = ParserContext {
        four_octet_asn: Some(peer_four_byte_asn),
        address_family: None,
    };

    let peer_asn = if peer_four_byte_asn {
        peer_open
            .options
            .iter()
            .find_map(|o| match o {
                OpenOption::Capabilities(caps) => caps.iter().find_map(|c| match c {
                    Capability::FourByteAsn { asn } => Some(*asn),
                    _ => None,
                }),
            })
            .unwrap_or(peer_open.asn as u32)
    } else {
        peer_open.asn as u32
    };

    if let Some(expected) = config.expected_asn
        && peer_asn != expected
    {
        send_notification(&mut writer, BgpError::OpenMsg(OpenMsgSubcode::BadPeerAs), &ctx).await;
        bail!("Peer {peer_addr} announced AS{peer_asn} but AS{expected} was configured");
    }

    // Connection collision detection (RFC 4271 §6.8). From here on the
    // guard keeps this connection visible to later ones, and `killed` fires
    // if one of them wins.
    let Some(mut collision_guard) =
        config.collisions.admit(config.router_id, peer_open.identifier, config.direction)
    else {
        return close_for_collision(&mut writer, &ctx, peer_addr).await;
    };

    info!(
        %peer_addr,
        peer_asn,
        negotiated_hold_time,
        direction = config.direction.as_str(),
        "OPEN received, transitioning to OpenConfirm"
    );

    // --- OpenConfirm --------------------------------------------------------
    send_message(&mut writer, &Message::KeepAlive, &ctx).await?;

    // With a negotiated hold time of zero the hold timer does not run
    // (RFC 4271 §4.2); fall back to the large pre-negotiation value so a
    // silent peer is still cleaned up eventually.
    let hold_duration = Duration::from_secs(if negotiated_hold_time == 0 { 240 } else { negotiated_hold_time as u64 });

    let confirm = tokio::select! {
        msg = recv_with_timeout(&mut reader, &ctx, hold_duration) => msg,
        _ = &mut collision_guard.killed => {
            return close_for_collision(&mut writer, &ctx, peer_addr).await;
        }
    };

    match confirm {
        Ok(Message::KeepAlive) => {}
        Ok(other) => {
            send_notification(
                &mut writer,
                BgpError::FsmError(FsmSubcode::UnexpectedOpenConfirm),
                &ctx,
            )
            .await;
            bail!("Expected KEEPALIVE in OpenConfirm state, got {:?}", other);
        }
        Err(e) => {
            send_notification(&mut writer, BgpError::HoldTimer, &ctx).await;
            return Err(e);
        }
    }

    info!(%peer_addr, peer_asn, "Session established");
    SESSIONS_ESTABLISHED_TOTAL.inc();

    let peer_id = PeerId { addr: peer_addr, asn: peer_asn };

    // --- Established --------------------------------------------------------
    // Register so the rest of the server can see this session and queue
    // messages for it. The guard deregisters on any exit path.
    let (outbound_tx, mut outbound_rx) = mpsc::channel(OUTBOUND_QUEUE_CAPACITY);
    collision_guard.set_established();
    let _guard = config.registry.register(SessionHandle::new(
        peer_id,
        config.peer_name.clone(),
        config.direction,
        local_addr,
        peer_four_byte_asn,
        outbound_tx,
    ));

    let result = established_loop(
        &mut reader,
        &mut writer,
        &ctx,
        negotiated_hold_time,
        peer_id,
        &config.adj_rib_in,
        SessionInputs { outbound: &mut outbound_rx, killed: &mut collision_guard.killed, registry: &config.registry },
    )
    .await;

    // Whatever ended the session, routes learned over it are gone
    // (RFC 4271 §8.2.2). Do this before the guard drops so subscribers see
    // the withdrawals before the Down event.
    let removed = config.adj_rib_in.remove_peer(peer_id).await;
    info!(peer = %peer_id, removed_routes = removed, "Session ended");
    result
}

/// Signals from the rest of the server into an established session.
struct SessionInputs<'a> {
    /// Messages queued for transmission to the peer.
    outbound: &'a mut mpsc::Receiver<Message>,
    /// Fires when a later connection won collision resolution.
    killed: &'a mut oneshot::Receiver<()>,
    /// Where to report a peer's route-refresh request.
    registry: &'a SessionRegistry,
}

async fn established_loop(
    reader: &mut OwnedReadHalf,
    writer: &mut OwnedWriteHalf,
    ctx: &ParserContext,
    hold_time_secs: u16,
    peer_id: PeerId,
    adj_rib_in: &AdjRibIn,
    inputs: SessionInputs<'_>,
) -> Result<()> {
    let SessionInputs { outbound, killed, registry } = inputs;
    // When hold_time is 0, keepalives and the hold timer are disabled
    // (RFC 4271 §4.4). The timer branches below are gated off in that case;
    // the placeholder durations are never observed.
    let timers_enabled = hold_time_secs != 0;
    let secs = if timers_enabled { hold_time_secs as u64 } else { 3 };
    let hold_duration = Duration::from_secs(secs);
    let keepalive_interval = Duration::from_secs(secs / 3);

    // interval_at avoids firing immediately; first keepalive is sent after one interval.
    let mut keepalive_timer = interval_at(Instant::now() + keepalive_interval, keepalive_interval);

    // sleep_until must be pinned so we can call .as_mut().reset() inside select!.
    let hold_timer = sleep_until(Instant::now() + hold_duration);
    tokio::pin!(hold_timer);

    loop {
        tokio::select! {
            result = read_message(reader, ctx) => {
                match result? {
                    Message::KeepAlive => {
                        hold_timer.as_mut().reset(Instant::now() + hold_duration);
                    }
                    Message::Update(update) => {
                        hold_timer.as_mut().reset(Instant::now() + hold_duration);
                        info!(
                            peer = %peer_id,
                            withdrawn = update.withdrawn_routes.len(),
                            announced = update.announced_routes.len(),
                            "Received UPDATE",
                        );
                        apply_update(adj_rib_in, peer_id, &update.path_attributes,
                            &update.withdrawn_routes, &update.announced_routes).await;
                    }
                    Message::Notification(notif) => {
                        info!(peer = %peer_id, error = %notif.error, "Received NOTIFICATION, closing session");
                        return Ok(());
                    }
                    Message::Open(_) => {
                        send_notification(writer, BgpError::FsmError(FsmSubcode::UnexpectedEstablished), ctx).await;
                        bail!("Received OPEN in Established state from {}", peer_id);
                    }
                    Message::RouteRefresh => {
                        hold_timer.as_mut().reset(Instant::now() + hold_duration);
                        info!(peer = %peer_id, "Received ROUTE-REFRESH, re-sending our routes");
                        registry.request_refresh(&peer_id);
                    }
                }
            }
            Some(msg) = outbound.recv() => {
                if let Message::Update(update) = &msg {
                    info!(
                        peer = %peer_id,
                        withdrawn = update.withdrawn_routes.len(),
                        announced = update.announced_routes.len(),
                        "Sending UPDATE",
                    );
                }
                send_message(writer, &msg, ctx).await?;
                // A NOTIFICATION from our side ends the session (RFC 4271 §6).
                if let Message::Notification(notif) = &msg {
                    info!(peer = %peer_id, error = %notif.error, "Sent NOTIFICATION, closing session");
                    return Ok(());
                }
            }
            _ = &mut *killed => {
                return close_for_collision(writer, ctx, peer_id.addr).await;
            }
            _ = keepalive_timer.tick(), if timers_enabled => {
                send_message(writer, &Message::KeepAlive, ctx).await?;
            }
            _ = &mut hold_timer, if timers_enabled => {
                warn!(peer = %peer_id, "Hold timer expired");
                send_notification(writer, BgpError::HoldTimer, ctx).await;
                bail!("Hold timer expired for peer {}", peer_id);
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Wire helpers
// ---------------------------------------------------------------------------

async fn read_message(reader: &mut OwnedReadHalf, ctx: &ParserContext) -> Result<Message> {
    let mut header = [0u8; BGP_HEADER_LEN as usize];
    reader.read_exact(&mut header).await?;

    let total_len = u16::from_be_bytes([header[16], header[17]]) as usize;
    if total_len < BGP_HEADER_LEN as usize {
        bail!("BGP message length {total_len} below minimum");
    }

    let mut buf = vec![0u8; total_len];
    buf[..BGP_HEADER_LEN as usize].copy_from_slice(&header);
    reader
        .read_exact(&mut buf[BGP_HEADER_LEN as usize..])
        .await?;

    let (_, msg) =
        Message::from_wire(ctx, &buf).map_err(|e| eyre!("Failed to parse BGP message: {e:?}"))?;
    MESSAGES_RECEIVED.with_label_values(&[message_type_name(&msg)]).inc();
    Ok(msg)
}

async fn recv_with_timeout(
    reader: &mut OwnedReadHalf,
    ctx: &ParserContext,
    timeout: Duration,
) -> Result<Message> {
    tokio::time::timeout(timeout, read_message(reader, ctx))
        .await
        .map_err(|_| eyre!("Timed out waiting for BGP message"))?
}

async fn send_message(
    writer: &mut OwnedWriteHalf,
    msg: &Message,
    ctx: &ParserContext,
) -> Result<()> {
    let mut buf = BytesMut::new();
    msg.to_wire(ctx, &mut buf)?;
    writer.write_all(&buf).await?;
    MESSAGES_SENT.with_label_values(&[message_type_name(msg)]).inc();
    Ok(())
}

/// Best-effort NOTIFICATION send — errors are logged and swallowed since we
/// are already tearing down the session.
async fn send_notification(writer: &mut OwnedWriteHalf, error: BgpError, ctx: &ParserContext) {
    let msg = Message::Notification(NotificationMessage {
        error,
        data: vec![],
    });
    if let Err(e) = send_message(writer, &msg, ctx).await {
        warn!("Failed to send NOTIFICATION: {e}");
    }
}

/// A later connection to the same BGP Identifier won collision resolution:
/// tell the peer why and end this session.
async fn close_for_collision(
    writer: &mut OwnedWriteHalf,
    ctx: &ParserContext,
    peer_addr: SocketAddr,
) -> Result<()> {
    info!(%peer_addr, "Closing connection: lost collision resolution");
    send_notification(writer, BgpError::Cease(CeaseSubcode::ConnCollisionResolution), ctx).await;
    bail!("Connection to {peer_addr} closed by collision resolution");
}

// ---------------------------------------------------------------------------
// RIB helpers
// ---------------------------------------------------------------------------

/// Applies a received UPDATE to the Adj-RIB-In.
///
/// Handles both the legacy IPv4 NLRI fields and the MP_REACH/MP_UNREACH path
/// attributes used for all other address families (RFC 2858).
async fn apply_update(
    rib: &AdjRibIn,
    peer: PeerId,
    path_attributes: &[PathAttribute],
    withdrawn: &[bgp_primitives::ip_prefix::IpPrefix],
    announced: &[bgp_primitives::ip_prefix::IpPrefix],
) {
    // Legacy IPv4 fields in the UPDATE body.
    for prefix in withdrawn {
        rib.withdraw(peer, prefix.clone()).await;
    }
    for prefix in announced {
        rib.announce(peer, prefix.clone(), RibEntry::new(path_attributes.to_vec())).await;
    }

    // Multiprotocol attributes carry NLRI for all other AFIs.
    for attr in path_attributes {
        match attr {
            PathAttribute::MpReachNlri(mp) => {
                for prefix in &mp.prefixes {
                    rib.announce(peer, prefix.clone(), RibEntry::new(path_attributes.to_vec())).await;
                }
            }
            PathAttribute::MpUnreachNlri(mp) => {
                for prefix in &mp.prefixes {
                    rib.withdraw(peer, prefix.clone()).await;
                }
            }
            _ => {}
        }
    }
}

// ---------------------------------------------------------------------------
// OPEN construction and validation
// ---------------------------------------------------------------------------

fn build_open(config: &SessionConfig) -> Message {
    // Use AS_TRANS in the 2-byte field when the ASN doesn't fit; the real ASN
    // is carried in the FourByteAsn capability (RFC 6793).
    let asn_field = if config.local_asn > u16::MAX as u32 {
        AS_TRANS
    } else {
        config.local_asn as u16
    };

    Message::Open(OpenMessage {
        version: BGP4_VERSION,
        asn: asn_field,
        hold_time: config.hold_time_secs,
        identifier: config.router_id,
        options: vec![OpenOption::Capabilities(vec![
            Capability::FourByteAsn {
                asn: config.local_asn,
            },
            Capability::MultiProtocol {
                afi: AddressFamilyId::Ipv4,
                safi: SubsequentAfi::Unicast,
            },
            // Without this a peer will not accept the IPv6 routes we
            // originate via MP_REACH_NLRI (RFC 4760 §8).
            Capability::MultiProtocol {
                afi: AddressFamilyId::Ipv6,
                safi: SubsequentAfi::Unicast,
            },
            Capability::RouteRefresh,
        ])],
    })
}

fn validate_open(open: &OpenMessage, local_id: Ipv4Addr) -> Result<(), BgpError> {
    if open.version != BGP4_VERSION {
        return Err(BgpError::OpenMsg(OpenMsgSubcode::UnsupportedVersion));
    }
    // A peer cannot share our identifier or use the unspecified address;
    // collision resolution depends on identifiers being distinct.
    if open.identifier == local_id || open.identifier.is_unspecified() {
        return Err(BgpError::OpenMsg(OpenMsgSubcode::BadBgpId));
    }
    if open.hold_time != 0 && open.hold_time < 3 {
        return Err(BgpError::OpenMsg(OpenMsgSubcode::UnacceptableHoldTime));
    }
    Ok(())
}