use crate::{DstAddr, LinkType, LockedStream, Peer, SourceAddr};
use async_std::sync::{Arc, RwLock};
use std::collections::BTreeMap;
use tracing::{trace, warn};
#[derive(Clone, Debug, Default)]
pub(crate) struct Routes {
port: u16,
peers: Arc<RwLock<BTreeMap<usize, Arc<Peer>>>>,
src_map: Arc<RwLock<BTreeMap<SourceAddr, usize>>>,
dst_map: Arc<RwLock<BTreeMap<DstAddr, usize>>>,
}
impl Routes {
pub(crate) fn new(port: u16) -> Arc<Self> {
Arc::new(Self {
port,
..Self::default()
})
}
pub(crate) async fn stop_all(self: &Arc<Self>) {
for (_, peer) in self.peers.read().await.iter() {
peer.stop();
}
}
pub(crate) async fn all_dst(self: &Arc<Self>) -> Vec<Arc<Peer>> {
self.peers
.read()
.await
.iter()
.filter_map(|(_, p)| match p.get_dst() {
Some(_) => Some(Arc::clone(&p)),
None => None,
})
.collect()
}
pub(crate) async fn get_peer(self: &Arc<Self>, id: usize) -> Option<Arc<Peer>> {
self.peers.read().await.get(&id).map(|p| Arc::clone(&p))
}
pub(crate) async fn purge_src(self: &Arc<Self>, src: SourceAddr) {
if let Some(id) = self.find_via_src(&src).await {
let peers = self.peers.read().await;
let p = peers.get(&id).unwrap();
p.set_src(None);
}
trace!("Removing existing SRC accociation: {:?}", src);
self.src_map.write().await.remove(&src);
}
pub(crate) async fn add_via_dst(self: &Arc<Self>, dst: DstAddr, _type: LinkType) -> usize {
let p = Peer::open(dst.clone(), self.port, _type);
let id = p.id;
self.peers.write().await.insert(id, p);
self.dst_map.write().await.insert(dst, id);
id
}
pub(crate) async fn add_via_src(self: &Arc<Self>, src: &SourceAddr) -> usize {
let p = Peer::from_src(src.clone());
let id = p.id;
self.peers.write().await.insert(id, p);
self.src_map.write().await.insert(src.clone(), id);
id
}
pub(crate) async fn add_src(self: &Arc<Self>, id: usize, src: SourceAddr) {
self.src_map.write().await.insert(src, id);
}
pub(crate) async fn find_via_src(self: &Arc<Self>, src: &SourceAddr) -> Option<usize> {
self.src_map.read().await.get(src).map(|id| *id)
}
pub(crate) async fn find_via_srcport(
self: &Arc<Self>,
src: &SourceAddr,
port: u16,
) -> Option<usize> {
let imply_dst = DstAddr::new(src.ip(), port);
self.dst_map.read().await.get(&imply_dst).map(|id| *id)
}
pub(crate) async fn upgrade(
self: &Arc<Self>,
id: usize,
port: u16,
stream: Option<LockedStream>,
) -> usize {
let mut peers = self.peers.write().await;
let mut src_map = self.src_map.write().await;
let mut dst_map = self.dst_map.write().await;
let peer = peers
.remove(&id)
.expect(&format!("Peer with id {} wasn't found!", id));
let src = peer
.get_src()
.expect("Invalid variant: peer must have SRC at this point");
src_map.remove(&src);
if let Some(ref dst) = peer.get_dst() {
trace!("Trying to upgrade a DUPLEX connection. Was there a connection drop?");
dst_map.remove(dst);
}
let dst = DstAddr::new(src.ip(), port);
match dst_map.get(&dst) {
Some(id) => {
if stream.is_some() {
warn!("An outgoing stream exists for a LIMITED incoming stream! ignoring...");
}
trace!("Upgrading peer {} with SRC address", id);
let peer = peers.get(&id).unwrap();
src_map.insert(src, peer.id);
peer.set_src(src);
peer.id
}
None => {
let p = Peer::open(dst, port, LinkType::Bidirect);
p.set_src(src);
if let Some(s) = stream {
p.set_stream(s).await;
}
src_map.insert(src, p.id);
dst_map.insert(dst, p.id);
peers.insert(p.id, p);
peer.id
}
}
}
}