ferrix/crates/rib/src/table/attrs.rs

//! Attribute sets, stored once.
//!
//! A full table has millions of routes and far fewer distinct lists of path
//! attributes: a peer announces prefixes that share their attributes
//! together, about six to an UPDATE on the Internet today. A route in the
//! table is therefore a pointer to an [`AttrSet`], and each distinct set
//! exists once, in the [`AttrStore`].
//!
//! A set is kept as its bytes on the wire, which is a fifth the size of the
//! parsed form. What choosing between routes needs from it is worked out
//! once, when the set is stored, and kept beside the bytes.

use std::hash::BuildHasher;
use std::net::{Ipv4Addr, Ipv6Addr};
use std::sync::{Arc, Mutex, PoisonError};

use bgp_packet::message::{NlriNextHop, OriginPathAttribute, PathAttribute};
use bgp_primitives::parser::{EncodeError, ParserContext};
use bytes::{BufMut, BytesMut};
use foldhash::fast::RandomState;
use hashbrown::HashTable;

const MP_REACH_NLRI: u8 = 14;
const OPTIONAL: u8 = 0x80;

/// What route selection reads from a set of attributes (RFC 4271 §9.1.2.2).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SelectionKey {
    /// LOCAL_PREF, if the set has one.
    pub local_pref: Option<u32>,
    /// MULTI_EXIT_DISC, if the set has one.
    pub med: Option<u32>,
    /// The AS that the route was learned from: the first AS of the path, or
    /// zero if the path is empty or begins with a set. MEDs are comparable
    /// only between routes with the same one.
    pub neighbor_as: u32,
    /// The length of the AS path as selection counts it: an AS_SET is one
    /// hop however many AS numbers it holds.
    pub as_path_len: u16,
    /// ORIGIN: 0 for IGP, 1 for EGP, 2 for INCOMPLETE.
    pub origin: u8,
}

impl SelectionKey {
    fn of(attributes: &[PathAttribute]) -> Self {
        let mut key = Self {
            local_pref: None,
            med: None,
            neighbor_as: 0,
            as_path_len: 0,
            // A route without ORIGIN does not get this far; if one did, it
            // would be the least preferred.
            origin: 2,
        };
        for attribute in attributes {
            match attribute {
                PathAttribute::Origin(origin) => {
                    key.origin = match origin {
                        OriginPathAttribute::IGP => 0,
                        OriginPathAttribute::EGP => 1,
                        OriginPathAttribute::INCOMPLETE => 2,
                    }
                }
                PathAttribute::AsPath(path) => {
                    let hops: usize = path
                        .segments
                        .iter()
                        .map(|s| if s.ordered { s.path.len() } else { 1 })
                        .sum();
                    key.as_path_len = u16::try_from(hops).unwrap_or(u16::MAX);
                    key.neighbor_as = path
                        .segments
                        .first()
                        .filter(|s| s.ordered)
                        .and_then(|s| s.path.first().copied())
                        .unwrap_or(0);
                }
                PathAttribute::MultiExitDisc(med) => key.med = Some(med.0),
                PathAttribute::LocalPref(local_pref) => key.local_pref = Some(local_pref.0),
                _ => {}
            }
        }
        key
    }
}

/// Reads the selection key straight off attributes in the form an
/// [`AttrSet`] keeps, without building anything. `None` if the bytes are
/// not a well-formed list of attributes.
fn scan(mut wire: &[u8]) -> Option<SelectionKey> {
    const EXTENDED_LENGTH: u8 = 0x10;
    let mut key = SelectionKey {
        local_pref: None,
        med: None,
        neighbor_as: 0,
        as_path_len: 0,
        origin: 2,
    };
    let number = |value: &[u8]| Some(u32::from_be_bytes(<[u8; 4]>::try_from(value).ok()?));
    while !wire.is_empty() {
        let (&flags, rest) = wire.split_first()?;
        let (&type_code, rest) = rest.split_first()?;
        let (length, rest) = if flags & EXTENDED_LENGTH != 0 {
            let (length, rest) = rest.split_first_chunk::<2>()?;
            (usize::from(u16::from_be_bytes(*length)), rest)
        } else {
            let (&length, rest) = rest.split_first()?;
            (usize::from(length), rest)
        };
        let (value, rest) = rest.split_at_checked(length)?;
        wire = rest;
        match type_code {
            1 => key.origin = *value.first()?,
            2 => {
                let mut segments = value;
                let mut hops = 0usize;
                let mut first = true;
                while let [kind, count, rest @ ..] = segments {
                    let (asns, rest) = rest.split_at_checked(usize::from(*count) * 4)?;
                    let ordered = *kind == 2;
                    hops += if ordered { usize::from(*count) } else { 1 };
                    if first && ordered {
                        key.neighbor_as = asns
                            .first_chunk::<4>()
                            .map_or(0, |a| u32::from_be_bytes(*a));
                    }
                    first = false;
                    segments = rest;
                }
                if !segments.is_empty() {
                    return None;
                }
                key.as_path_len = u16::try_from(hops).unwrap_or(u16::MAX);
            }
            4 => key.med = Some(number(value)?),
            5 => key.local_pref = Some(number(value)?),
            _ => {}
        }
    }
    Some(key)
}

/// One distinct list of path attributes, with the next hop that goes with
/// it. Made by [`AttrStore::intern`]; two routes have the same attributes
/// exactly when they point to the same `AttrSet`.
#[derive(Debug)]
pub struct AttrSet {
    hash: u64,
    key: SelectionKey,
    /// The attributes as on the wire, with four-octet AS numbers. A next
    /// hop that arrived in MP_REACH_NLRI is kept as an MP_REACH_NLRI that
    /// holds nothing else, the form RFC 6396 §4.3.4 uses in table dumps.
    wire: Box<[u8]>,
}

impl AttrSet {
    pub fn key(&self) -> &SelectionKey {
        &self.key
    }

    /// The attributes as they are on the wire.
    pub fn wire(&self) -> &[u8] {
        &self.wire
    }

    /// The attributes, parsed, and the next hop if it came in
    /// MP_REACH_NLRI.
    pub fn attributes(&self) -> (Vec<PathAttribute>, Option<NlriNextHop>) {
        // The bytes were written by `encode`, so they parse.
        parse(&self.wire).unwrap_or_default()
    }
}

fn four_octet() -> ParserContext {
    ParserContext {
        four_octet_asn: Some(true),
        ibgp: true,
        add_path: Vec::new(),
    }
}

/// Writes `attributes` and `next_hop` in the form an [`AttrSet`] keeps.
fn encode(
    attributes: &[PathAttribute],
    next_hop: Option<&NlriNextHop>,
    out: &mut BytesMut,
) -> Result<(), EncodeError> {
    let ctx = four_octet();
    for attribute in attributes {
        // These name prefixes, which are not a property of the route.
        if matches!(
            attribute,
            PathAttribute::MpReachNlri(_) | PathAttribute::MpUnreachNlri(_)
        ) {
            continue;
        }
        attribute.to_wire(&ctx, out)?;
    }
    if let Some(next_hop) = next_hop {
        out.put_u8(OPTIONAL);
        out.put_u8(MP_REACH_NLRI);
        out.put_u8(1 + next_hop.wire_len());
        out.put_u8(next_hop.wire_len());
        next_hop.to_wire(&ctx, out)?;
    }
    Ok(())
}

/// Parses path attributes in the form [`AttrSet`] keeps and table dumps
/// use (RFC 6396 §4.3.4): four-octet AS numbers, and an MP_REACH_NLRI that
/// holds only a next hop. `None` if they are malformed.
pub fn parse(mut wire: &[u8]) -> Option<(Vec<PathAttribute>, Option<NlriNextHop>)> {
    let ctx = four_octet();
    let mut attributes = Vec::new();
    let mut next_hop = None;
    while !wire.is_empty() {
        if wire.get(1) == Some(&MP_REACH_NLRI) {
            let extended = wire[0] & 0x10 != 0;
            let (length, header) = if extended {
                (
                    usize::from(u16::from_be_bytes([*wire.get(2)?, *wire.get(3)?])),
                    4,
                )
            } else {
                (usize::from(*wire.get(2)?), 3)
            };
            let value = wire.get(header..header + length)?;
            let address = |bytes: &[u8]| Some(Ipv6Addr::from(<[u8; 16]>::try_from(bytes).ok()?));
            next_hop = match value {
                [4, a, b, c, d] => Some(NlriNextHop::Ipv4(Ipv4Addr::new(*a, *b, *c, *d))),
                [16, rest @ ..] => Some(NlriNextHop::Ipv6(address(rest)?)),
                [32, rest @ ..] if rest.len() == 32 => Some(NlriNextHop::Ipv6WithLl {
                    global: address(&rest[..16])?,
                    link_local: address(&rest[16..])?,
                }),
                _ => return None,
            };
            wire = &wire[header + length..];
            continue;
        }
        let (rest, attribute) = PathAttribute::from_wire(&ctx, wire).ok()?;
        attributes.push(attribute);
        wire = rest;
    }
    Some((attributes, next_hop))
}

/// How many locks the store is split over. Sessions intern concurrently,
/// once per UPDATE; this many locks keep them out of each other's way.
const STRIPES: usize = 64;

/// The one place attribute sets are kept. Safe to share between threads.
pub struct AttrStore {
    stripes: Box<[Mutex<HashTable<Arc<AttrSet>>>]>,
    hasher: RandomState,
}

impl Default for AttrStore {
    fn default() -> Self {
        Self::new()
    }
}

impl AttrStore {
    pub fn new() -> Self {
        Self {
            stripes: (0..STRIPES).map(|_| Mutex::default()).collect(),
            hasher: RandomState::default(),
        }
    }

    /// The stored set equal to `attributes` and `next_hop`, added if it was
    /// not there. `next_hop` is the one from MP_REACH_NLRI; routes whose
    /// next hop is the NEXT_HOP attribute pass `None`.
    ///
    /// Fails only if the attributes cannot be encoded, which attributes that
    /// were decoded from a message always can.
    pub fn intern(
        &self,
        attributes: &[PathAttribute],
        next_hop: Option<&NlriNextHop>,
    ) -> Result<Arc<AttrSet>, EncodeError> {
        let mut wire = BytesMut::with_capacity(256);
        encode(attributes, next_hop, &mut wire)?;
        let hash = self.hasher.hash_one(&wire[..]);

        let mut stripe = self.stripe(hash);
        if let Some(existing) = stripe.find(hash, |set| set.wire[..] == wire[..]) {
            return Ok(Arc::clone(existing));
        }
        let key = SelectionKey::of(attributes);
        debug_assert_eq!(scan(&wire), Some(key), "both ways of reading the key agree");
        let set = Arc::new(AttrSet {
            hash,
            key,
            wire: wire[..].into(),
        });
        stripe.insert_unique(hash, Arc::clone(&set), |set| set.hash);
        Ok(set)
    }

    /// As [`AttrStore::intern`], for attributes that are already in the
    /// form a set keeps: on the wire, with four-octet AS numbers, no
    /// MP_UNREACH_NLRI, and an MP_REACH_NLRI (if any) that holds only the
    /// next hop. That is how a table dump has them (RFC 6396 §4.3.4).
    ///
    /// Nothing is parsed: finding a set that is already stored costs a hash
    /// of the bytes and a comparison. `None` if the bytes are not a
    /// well-formed list of attributes.
    pub fn intern_wire(&self, wire: &[u8]) -> Option<Arc<AttrSet>> {
        let hash = self.hasher.hash_one(wire);
        let mut stripe = self.stripe(hash);
        if let Some(existing) = stripe.find(hash, |set| *set.wire == *wire) {
            return Some(Arc::clone(existing));
        }
        let set = Arc::new(AttrSet {
            hash,
            key: scan(wire)?,
            wire: wire.into(),
        });
        stripe.insert_unique(hash, Arc::clone(&set), |set| set.hash);
        Some(set)
    }

    /// The set for attributes as a session received them: `wire` is what
    /// the decoder kept of the UPDATE's attributes
    /// (`DecodedUpdate::attributes_wire`), and `next_hop` the next hop of
    /// the routes if they came in MP_REACH_NLRI.
    pub fn intern_received(
        &self,
        wire: &[u8],
        next_hop: Option<&NlriNextHop>,
    ) -> Option<Arc<AttrSet>> {
        let Some(next_hop) = next_hop else {
            return self.intern_wire(wire);
        };
        let mut with_next_hop = BytesMut::with_capacity(wire.len() + 36);
        with_next_hop.put_slice(wire);
        encode(&[], Some(next_hop), &mut with_next_hop).ok()?;
        self.intern_wire(&with_next_hop)
    }

    /// How many distinct sets are stored.
    pub fn len(&self) -> usize {
        (0..STRIPES).map(|i| self.stripe_at(i).len()).sum()
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Drops the sets that no route refers to any more, and returns how
    /// many that was. Sets are not dropped as their last route goes, since
    /// that would mean a lock here for every route removed; this is called
    /// now and then instead.
    pub fn sweep(&self) -> usize {
        let mut dropped = 0;
        for i in 0..STRIPES {
            let mut stripe = self.stripe_at(i);
            let before = stripe.len();
            // The store's own reference is the only one left. Nobody can
            // be taking another: that is done under this lock.
            stripe.retain(|set| Arc::strong_count(set) > 1);
            dropped += before - stripe.len();
        }
        dropped
    }

    fn stripe(&self, hash: u64) -> std::sync::MutexGuard<'_, HashTable<Arc<AttrSet>>> {
        // The high bits: the table inside uses the low ones.
        self.stripe_at((hash >> 56) as usize % STRIPES)
    }

    fn stripe_at(&self, index: usize) -> std::sync::MutexGuard<'_, HashTable<Arc<AttrSet>>> {
        // A panic while holding the lock cannot leave the table half
        // changed in a way that matters: see `IgnorePoison` in the server.
        self.stripes[index]
            .lock()
            .unwrap_or_else(PoisonError::into_inner)
    }
}

#[cfg(test)]
mod tests {
    use bgp_packet::message::{
        AsPathAttribute, AsPathSegment, CommunitiesPathAttribute, LocalPrefPathAttribute,
        MultiExitDiscPathAttribute, NextHopPathAttribute,
    };

    use super::*;

    fn attributes(path: Vec<u32>) -> Vec<PathAttribute> {
        vec![
            PathAttribute::Origin(OriginPathAttribute::EGP),
            AsPathAttribute::from_asns(path),
            PathAttribute::NextHop(NextHopPathAttribute(Ipv4Addr::new(192, 0, 2, 1))),
            PathAttribute::MultiExitDisc(MultiExitDiscPathAttribute(50)),
            PathAttribute::Communities(CommunitiesPathAttribute(vec![(65000, 1)])),
        ]
    }

    #[test]
    fn equal_attributes_are_one_set() {
        let store = AttrStore::new();
        let a = store.intern(&attributes(vec![65001, 65002]), None).unwrap();
        let b = store.intern(&attributes(vec![65001, 65002]), None).unwrap();
        let c = store.intern(&attributes(vec![65001, 65003]), None).unwrap();
        assert!(Arc::ptr_eq(&a, &b));
        assert!(!Arc::ptr_eq(&a, &c));
        assert_eq!(store.len(), 2);

        // The next hop is part of what a set is.
        let next_hop = NlriNextHop::Ipv6("2001:db8::1".parse().unwrap());
        let d = store
            .intern(&attributes(vec![65001, 65002]), Some(&next_hop))
            .unwrap();
        assert!(!Arc::ptr_eq(&a, &d));
        assert_eq!(store.len(), 3);
    }

    #[test]
    fn a_set_gives_back_what_went_in() {
        let store = AttrStore::new();
        let next_hop = NlriNextHop::Ipv6WithLl {
            global: "2001:db8::1".parse().unwrap(),
            link_local: "fe80::1".parse().unwrap(),
        };
        let original = attributes(vec![65001, 4_200_000_001]);
        let set = store.intern(&original, Some(&next_hop)).unwrap();
        assert_eq!(set.attributes(), (original, Some(next_hop)));
    }

    #[test]
    fn the_selection_key_is_read_off_the_attributes() {
        let store = AttrStore::new();
        let mut attrs = vec![
            PathAttribute::Origin(OriginPathAttribute::INCOMPLETE),
            PathAttribute::AsPath(AsPathAttribute {
                segments: vec![
                    AsPathSegment {
                        ordered: true,
                        path: vec![65001, 65002],
                    },
                    // A set counts as one hop.
                    AsPathSegment {
                        ordered: false,
                        path: vec![65003, 65004, 65005],
                    },
                ],
            }),
            PathAttribute::LocalPref(LocalPrefPathAttribute(200)),
        ];
        let key = *store.intern(&attrs, None).unwrap().key();
        assert_eq!(
            key,
            SelectionKey {
                local_pref: Some(200),
                med: None,
                neighbor_as: 65001,
                as_path_len: 3,
                origin: 2,
            }
        );

        attrs.push(PathAttribute::MultiExitDisc(MultiExitDiscPathAttribute(7)));
        assert_eq!(store.intern(&attrs, None).unwrap().key().med, Some(7));
    }

    #[test]
    fn attributes_on_the_wire_are_interned_without_parsing_them() {
        let store = AttrStore::new();
        let next_hop = NlriNextHop::Ipv6("2001:db8::1".parse().unwrap());
        let mut attrs = attributes(vec![65001, 65002, 65003]);
        attrs.push(PathAttribute::LocalPref(LocalPrefPathAttribute(300)));
        let parsed = store.intern(&attrs, Some(&next_hop)).unwrap();

        // The same set by its bytes is the same set, key and all.
        let by_wire = store.intern_wire(parsed.wire()).unwrap();
        assert!(Arc::ptr_eq(&parsed, &by_wire));

        // And one that arrives as bytes first gets the same key as it would
        // have had from the parsed form.
        let other = AttrStore::new();
        let fresh = other.intern_wire(parsed.wire()).unwrap();
        assert_eq!(fresh.key(), parsed.key());
        assert_eq!(fresh.key().as_path_len, 3);
        assert_eq!(fresh.key().neighbor_as, 65001);
        assert_eq!(fresh.attributes(), (attrs, Some(next_hop)));

        // Bytes that are not attributes are refused.
        assert!(other.intern_wire(&[0x40, 0x01, 0x05, 0x00]).is_none());
        assert!(
            other
                .intern_wire(&[0x40, 0x02, 0x03, 0x02, 0x01, 0x00])
                .is_none()
        );
    }

    #[test]
    fn sweeping_drops_only_sets_nobody_holds() {
        let store = AttrStore::new();
        let kept = store.intern(&attributes(vec![65001]), None).unwrap();
        drop(store.intern(&attributes(vec![65002]), None).unwrap());
        assert_eq!(store.len(), 2);

        assert_eq!(store.sweep(), 1);
        assert_eq!(store.len(), 1);
        // The one that is left is still the one that is found.
        let again = store.intern(&attributes(vec![65001]), None).unwrap();
        assert!(Arc::ptr_eq(&kept, &again));
    }
}