use std::collections::HashMap;
use ts_keys::NodePublicKey;
use ts_packet::PacketMut;
use ts_transport::UnderlayTransportId;
#[derive(Default)]
pub struct Router {
pub table: HashMap<NodePublicKey, UnderlayTransportId>,
}
pub type Result = HashMap<(UnderlayTransportId, NodePublicKey), Vec<PacketMut>>;
impl Router {
pub fn route(
&self,
batches: impl IntoIterator<Item = (NodePublicKey, Vec<PacketMut>)>,
) -> Result {
let mut ret = Result::default();
for (peer_id, packets) in batches {
if let Some(transport) = self.table.get(&peer_id) {
ret.entry((*transport, peer_id))
.or_default()
.extend(packets);
}
}
ret
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_outbound_underlay() {
let peer_a = NodePublicKey::from([1u8; 32]);
let peer_b = NodePublicKey::from([2u8; 32]);
let peer_c = NodePublicKey::from([3u8; 32]);
let peer_d = NodePublicKey::from([4u8; 32]);
let peer_e = NodePublicKey::from([5u8; 32]);
let transport_a = 5.into();
let transport_b = 6.into();
let transport_c = 7.into();
let mut router = Router::default();
router.table.insert(peer_a, transport_b);
router.table.insert(peer_b, transport_a);
router.table.insert(peer_c, transport_b);
router.table.insert(peer_e, transport_c);
let mut got = router.route([
(peer_a, vec![b"foo".into(), b"bar".into()]),
(peer_b, vec![b"qux".into(), b"xyzzy".into()]),
(peer_c, vec![b"frobozz".into(), b"zork".into()]),
(peer_d, vec![b"frotz".into(), b"get lamp".into()]),
]);
let mut want = Result::from([
((transport_b, peer_a), vec![b"foo".into(), b"bar".into()]),
((transport_a, peer_b), vec![b"qux".into(), b"xyzzy".into()]),
(
(transport_b, peer_c),
vec![b"frobozz".into(), b"zork".into()],
),
]);
for pkts in got.values_mut() {
pkts.sort();
}
for pkts in want.values_mut() {
pkts.sort();
}
let mut got = got.into_iter().collect::<Vec<_>>();
let mut want = want.into_iter().collect::<Vec<_>>();
got.sort();
want.sort();
assert_eq!(got, want);
}
}