diff --git a/src/ip_pubkey.rs b/src/ip_pubkey.rs new file mode 100644 index 0000000..3577ddf --- /dev/null +++ b/src/ip_pubkey.rs @@ -0,0 +1,58 @@ +//! A module to abstract the way known public keys for destinations are stored. + +use std::net::{IpAddr, Ipv6Addr}; + +use ip_network_table_deps_treebitmap::IpLookupTable; + +use crate::{ + crypto::{PublicKey, SharedSecret}, + subnet::Subnet, +}; + +pub struct IpPubkeyMap { + inner: IpLookupTable, +} + +impl IpPubkeyMap { + /// Create a new empty `IpPubkeyMap`. + pub fn new() -> Self { + Self { + inner: IpLookupTable::new(), + } + } + + /// Insert a new [`PublicKey`] and [`SharedSecret`] for a [`Subnet`] in the map. + pub fn insert(&mut self, subnet: Subnet, pk: PublicKey, ss: SharedSecret) { + let addr = match subnet.network() { + IpAddr::V6(addr) => addr, + _ => return, + }; + self.inner + .insert(addr, subnet.prefix_len() as u32, (pk, ss)); + } + + /// Remove the [`PublicKey`] and [`SharedSecret`] for a given [`Subnet`] from the map. + pub fn remove(&mut self, subnet: Subnet) { + let addr = match subnet.network() { + IpAddr::V6(addr) => addr, + _ => return, + }; + self.inner.remove(addr, subnet.prefix_len() as u32); + } + + /// Look up the [`PublicKey`] and [`SharedSecret`] associated with an [`IpAddr`], if any is + /// known in the map. + pub fn lookup(&self, addr: Ipv6Addr) -> Option<&(PublicKey, SharedSecret)> { + self.inner.longest_match(addr).map(|(_, _, v)| v) + } +} + +impl Clone for IpPubkeyMap { + fn clone(&self) -> Self { + let mut new = IpPubkeyMap::new(); + for (addr, plen, value) in self.inner.iter() { + new.inner.insert(addr, plen, value.clone()); + } + new + } +} diff --git a/src/main.rs b/src/main.rs index 997e0cb..7b44624 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,3 @@ -use crate::filters; use crate::router::StaticRoute; use crate::{packet::DataPacket, subnet::Subnet}; use bytes::BytesMut; @@ -20,6 +19,7 @@ mod codec; mod crypto; mod filters; mod interval; +mod ip_pubkey; mod metric; mod packet; mod peer;