use std::collections::{
hash_map::{IntoIter, Iter},
HashMap, HashSet,
};
use crate::{
bgp::BgpRoute,
network::Network,
ospf::OspfImpl,
types::{Prefix, PrefixMap, RouterId},
};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BgpStateRef<'n, P: Prefix> {
prefix: P,
g: BgpStateGraph<&'n BgpRoute<P>>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BgpState<P: Prefix> {
prefix: P,
g: BgpStateGraph<BgpRoute<P>>,
}
impl<P: Prefix> BgpState<P> {
#[inline]
pub fn from_net<Q, Ospf: OspfImpl>(net: &Network<P, Q, Ospf>, prefix: P) -> Self {
Self {
prefix,
g: BgpStateGraph::from_net(net, prefix, |e| e.clone()),
}
}
#[inline]
pub fn prefix(&self) -> P {
self.prefix
}
#[inline]
pub fn get(&self, router: RouterId) -> Option<(RouterId, &BgpRoute<P>)> {
self.g
.get(router)
.and_then(|node| node.node.as_ref())
.map(|(r, x)| (*x, r))
}
#[inline]
pub fn selected(&self, router: RouterId) -> Option<&BgpRoute<P>> {
self.g.get(router).and_then(|node| node.selected())
}
#[inline]
pub fn learned_from(&self, router: RouterId) -> Option<RouterId> {
self.g.get(router).and_then(|node| node.learned_from())
}
#[inline]
pub fn advertised(&self, src: RouterId, dst: RouterId) -> Option<&BgpRoute<P>> {
self.g.advertised(src, dst)
}
#[inline]
pub fn iter(&self) -> impl Iterator<Item = (RouterId, Option<&BgpRoute<P>>)> {
self.g
.iter()
.map(|(router, node)| (*router, node.selected()))
}
#[inline]
pub fn outgoing(&self, router: RouterId) -> impl Iterator<Item = (RouterId, &BgpRoute<P>)> {
self.g.outgoing(router).map(|(n, r)| (*n, r))
}
#[inline]
pub fn incoming(&self, router: RouterId) -> impl Iterator<Item = (RouterId, &BgpRoute<P>)> {
self.g.incoming(router)
}
#[inline]
pub fn peers_outgoing(&self, router: RouterId) -> impl Iterator<Item = RouterId> + '_ {
self.g.peers_outgoing(router)
}
#[inline]
pub fn peers_incoming(&self, router: RouterId) -> impl Iterator<Item = RouterId> + '_ {
self.g.peers_incoming(router)
}
#[inline]
pub fn propagation_path(&self, router: RouterId) -> Vec<RouterId> {
self.g.propagation_path(router)
}
#[inline]
pub fn ingress_session(&self, router: RouterId) -> Option<(RouterId, RouterId)> {
self.g.ingress_session(router)
}
#[inline]
pub fn reach(&self, router: RouterId) -> HashSet<RouterId> {
self.g.reach(router)
}
}
impl<'n, P: Prefix> BgpStateRef<'n, P> {
#[inline]
pub fn from_net<Q, Ospf: OspfImpl>(net: &'n Network<P, Q, Ospf>, prefix: P) -> Self {
Self {
prefix,
g: BgpStateGraph::from_net(net, prefix, |e| e),
}
}
#[inline]
pub fn prefix(&self) -> P {
self.prefix
}
#[inline]
pub fn get(&self, router: RouterId) -> Option<(RouterId, &'n BgpRoute<P>)> {
self.g
.get(router)
.and_then(|node| node.node.as_ref())
.map(|(r, x)| (*x, *r))
}
#[inline]
pub fn selected(&self, router: RouterId) -> Option<&'n BgpRoute<P>> {
self.g.get(router).and_then(|node| node.selected()).copied()
}
#[inline]
pub fn learned_from(&self, router: RouterId) -> Option<RouterId> {
self.g.get(router).and_then(|node| node.learned_from())
}
#[inline]
pub fn advertised(&self, src: RouterId, dst: RouterId) -> Option<&'n BgpRoute<P>> {
self.g.advertised(src, dst).copied()
}
#[inline]
pub fn iter(&self) -> impl Iterator<Item = (RouterId, Option<&'n BgpRoute<P>>)> + '_ {
self.g
.iter()
.map(|(router, node)| (*router, node.selected().copied()))
}
#[inline]
pub fn outgoing(
&self,
router: RouterId,
) -> impl Iterator<Item = (RouterId, &'n BgpRoute<P>)> + '_ {
self.g.outgoing(router).map(|(n, r)| (*n, *r))
}
#[inline]
pub fn incoming(
&self,
router: RouterId,
) -> impl Iterator<Item = (RouterId, &'n BgpRoute<P>)> + '_ {
self.g.incoming(router).map(|(n, r)| (n, *r))
}
#[inline]
pub fn peers_outgoing(&self, router: RouterId) -> impl Iterator<Item = RouterId> + '_ {
self.g.peers_outgoing(router)
}
#[inline]
pub fn peers_incoming(&self, router: RouterId) -> impl Iterator<Item = RouterId> + '_ {
self.g.peers_incoming(router)
}
#[inline]
pub fn propagation_path(&self, router: RouterId) -> Vec<RouterId> {
self.g.propagation_path(router)
}
#[inline]
pub fn ingress_session(&self, router: RouterId) -> Option<(RouterId, RouterId)> {
self.g.ingress_session(router)
}
#[inline]
pub fn reach(&self, router: RouterId) -> HashSet<RouterId> {
self.g.reach(router)
}
}
impl<'n, P: Prefix> From<BgpStateRef<'n, P>> for BgpState<P> {
fn from(val: BgpStateRef<'n, P>) -> Self {
BgpState {
prefix: val.prefix,
g: val
.g
.into_iter()
.map(|(r, n)| (r, n.into_owned()))
.collect(),
}
}
}
impl<P: Prefix> BgpState<P> {
pub fn as_state_ref(&self) -> BgpStateRef<'_, P> {
BgpStateRef {
prefix: self.prefix,
g: self.g.iter().map(|(k, v)| (*k, v.as_node_ref())).collect(),
}
}
}
impl<P: Prefix> BgpStateRef<'_, P> {
pub fn as_owned(&self) -> BgpState<P> {
BgpState {
prefix: self.prefix,
g: self.g.iter().map(|(r, n)| (*r, n.as_owned())).collect(),
}
}
pub fn into_owned(self) -> BgpState<P> {
BgpState {
prefix: self.prefix,
g: self
.g
.into_iter()
.map(|(r, n)| (r, n.into_owned()))
.collect(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct BgpStateGraph<T>(HashMap<RouterId, BgpStateNode<T>>);
impl<T> FromIterator<(RouterId, BgpStateNode<T>)> for BgpStateGraph<T> {
fn from_iter<I: IntoIterator<Item = (RouterId, BgpStateNode<T>)>>(iter: I) -> Self {
Self(HashMap::from_iter(iter))
}
}
impl<T> BgpStateGraph<T> {
fn from_net<'n, P: Prefix, Ospf: OspfImpl, F: Fn(&'n BgpRoute<P>) -> T, Q>(
net: &'n Network<P, Q, Ospf>,
prefix: P,
f: F,
) -> Self {
let mut g = net
.indices()
.map(|id| (id, BgpStateNode::default()))
.collect::<HashMap<_, _>>();
for r in net.routers() {
let id = r.router_id();
if let Some(entry) = r.bgp.get_route(prefix) {
g.get_mut(&id).unwrap().node = Some((f(&entry.route), entry.from_id));
}
r.bgp
.get_rib_out()
.get(&prefix)
.into_iter()
.flatten()
.for_each(|(peer, entry)| {
g.get_mut(&id)
.unwrap()
.edges_out
.insert(*peer, f(&entry.route));
g.get_mut(peer).unwrap().edges_in.insert(id);
});
}
Self(g)
}
#[inline]
fn iter(&self) -> Iter<'_, RouterId, BgpStateNode<T>> {
self.0.iter()
}
#[inline]
fn get(&self, k: RouterId) -> Option<&BgpStateNode<T>> {
self.0.get(&k)
}
#[inline]
fn advertised(&self, src: RouterId, dst: RouterId) -> Option<&T> {
self.0.get(&src).and_then(|r| r.edges_out.get(&dst))
}
#[inline]
fn outgoing(&self, router: RouterId) -> impl Iterator<Item = (&RouterId, &T)> + '_ {
self.get(router)
.into_iter()
.flat_map(|x| x.edges_out.iter())
}
#[inline]
fn incoming(&self, router: RouterId) -> BgpGraphIncomingIterator<'_, T> {
BgpGraphIncomingIterator {
origin: router,
iter: self.0[&router].edges_in.iter(),
g: self,
}
}
#[inline]
fn peers_outgoing(&self, router: RouterId) -> impl Iterator<Item = RouterId> + '_ {
self.get(router)
.into_iter()
.flat_map(|x| x.edges_out.keys().copied())
}
#[inline]
fn peers_incoming(&self, router: RouterId) -> impl Iterator<Item = RouterId> + '_ {
self.get(router)
.into_iter()
.flat_map(|x| x.edges_in.iter().copied())
}
fn propagation_path(&self, mut router: RouterId) -> Vec<RouterId> {
if self.get(router).and_then(|r| r.learned_from()).is_none() {
return Vec::new();
}
let mut path = vec![router];
while let Some(nh) = self.get(router).and_then(|r| r.learned_from()) {
if nh == router {
break;
}
router = nh;
path.push(router)
}
path.reverse();
path
}
pub fn ingress_session(&self, router: RouterId) -> Option<(RouterId, RouterId)> {
let path = self.propagation_path(router);
if path.len() <= 1 {
None
} else {
Some((path[0], path[1]))
}
}
fn reach(&self, router: RouterId) -> HashSet<RouterId> {
let mut set = HashSet::new();
let mut to_visit = vec![router];
while let Some(cur) = to_visit.pop() {
set.insert(cur);
to_visit.extend(
self.peers_outgoing(cur)
.filter(|r| !set.contains(r))
.filter(|r| self.0[r].learned_from() == Some(cur)),
);
}
set
}
}
struct BgpGraphIncomingIterator<'a, T> {
origin: RouterId,
iter: std::collections::hash_set::Iter<'a, RouterId>,
g: &'a BgpStateGraph<T>,
}
impl<'a, T> Iterator for BgpGraphIncomingIterator<'a, T> {
type Item = (RouterId, &'a T);
fn next(&mut self) -> Option<Self::Item> {
self.iter.next().copied().map(|r| {
(
r,
self.g
.get(r)
.and_then(|x| x.edges_out.get(&self.origin))
.unwrap(),
)
})
}
}
impl<T> IntoIterator for BgpStateGraph<T> {
type Item = (RouterId, BgpStateNode<T>);
type IntoIter = IntoIter<RouterId, BgpStateNode<T>>;
#[inline]
fn into_iter(self) -> Self::IntoIter {
self.0.into_iter()
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct BgpStateNode<T> {
pub(self) node: Option<(T, RouterId)>,
edges_out: HashMap<RouterId, T>,
edges_in: HashSet<RouterId>,
}
impl<T> Default for BgpStateNode<T> {
fn default() -> Self {
Self {
node: Default::default(),
edges_out: Default::default(),
edges_in: Default::default(),
}
}
}
impl<T> BgpStateNode<T> {
#[inline]
fn selected(&self) -> Option<&T> {
self.node.as_ref().map(|(route, _)| route)
}
#[inline]
fn learned_from(&self) -> Option<RouterId> {
self.node.as_ref().map(|(_, r)| *r)
}
}
impl<T: Clone> BgpStateNode<&T> {
fn into_owned(self) -> BgpStateNode<T> {
BgpStateNode {
node: self.node.map(|(route, from)| (route.clone(), from)),
edges_out: self
.edges_out
.into_iter()
.map(|(k, v)| (k, v.clone()))
.collect(),
edges_in: self.edges_in,
}
}
fn as_owned(&self) -> BgpStateNode<T> {
BgpStateNode {
node: self
.node
.as_ref()
.map(|(route, from)| ((*route).clone(), *from)),
edges_out: self
.edges_out
.iter()
.map(|(k, v)| (*k, (*v).clone()))
.collect(),
edges_in: self.edges_in.clone(),
}
}
}
impl<T> BgpStateNode<T> {
fn as_node_ref(&self) -> BgpStateNode<&T> {
BgpStateNode {
node: self.node.as_ref().map(|(route, from)| (route, *from)),
edges_out: self.edges_out.iter().map(|(k, v)| (*k, v)).collect(),
edges_in: self.edges_in.clone(),
}
}
}