93 lines
2.2 KiB
Rust
93 lines
2.2 KiB
Rust
use std::{net::Ipv4Addr, sync::Arc};
|
|
|
|
use arc_swap::ArcSwap;
|
|
|
|
#[derive(Default, Clone)]
|
|
pub struct TrieNode {
|
|
child: [Option<Box<TrieNode>>; 2],
|
|
prefix_len: u8,
|
|
nexthop: Option<Ipv4Addr>,
|
|
}
|
|
|
|
#[derive(Default, Clone)]
|
|
struct IpTrie {
|
|
root: TrieNode,
|
|
}
|
|
|
|
impl IpTrie {
|
|
fn new() -> Self {
|
|
Self {
|
|
root: TrieNode::default()
|
|
}
|
|
}
|
|
|
|
fn insert(&mut self, prefix: u32, prefix_len: u8, nexthop: Ipv4Addr) {
|
|
if prefix_len > 32 {
|
|
return;
|
|
}
|
|
|
|
let mut node = &mut self.root;
|
|
for i in 0..prefix_len {
|
|
let bit = ((prefix >> (31-i)) & 1) as usize;
|
|
node = node.child[bit].get_or_insert_with(|| Box::new(TrieNode::default()));
|
|
|
|
}
|
|
if prefix_len > node.prefix_len {
|
|
node.prefix_len = prefix_len;
|
|
node.nexthop = Some(nexthop);
|
|
}
|
|
}
|
|
|
|
fn lookup(&self, ip: u32) -> Option<(u8, Ipv4Addr)>{
|
|
let mut node = &self.root;
|
|
let mut best = None;
|
|
|
|
for i in 0..32 {
|
|
if node.nexthop.is_some() {
|
|
best = Some((node.prefix_len, node.nexthop.unwrap()));
|
|
}
|
|
let bit = ((ip>>(31-i)) & 1) as usize;
|
|
match &node.child[bit] {
|
|
Some(child) => {
|
|
node = child;
|
|
}
|
|
None => {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
if node.nexthop.is_some() {
|
|
best = Some((node.prefix_len, node.nexthop.unwrap()));
|
|
}
|
|
best
|
|
}
|
|
}
|
|
|
|
pub struct RouteTableTrie {
|
|
trie: ArcSwap<IpTrie>,
|
|
}
|
|
|
|
impl RouteTableTrie {
|
|
pub fn new() -> Self {
|
|
Self {
|
|
trie: ArcSwap::new(Arc::new(IpTrie::default()))
|
|
}
|
|
}
|
|
|
|
pub fn clear(&self) {
|
|
self.trie.store(Arc::new(IpTrie::default()));
|
|
}
|
|
|
|
pub fn lookup(&self, ip: u32) -> Option<(u8, Ipv4Addr)> {
|
|
let trie = self.trie.load();
|
|
trie.lookup(ip)
|
|
}
|
|
|
|
pub fn insert(&self, prefix: u32, prefix_len: u8, nexthop: Ipv4Addr) {
|
|
let old = self.trie.load();
|
|
let mut new_trie = (*(*old)).clone();
|
|
|
|
new_trie.insert(prefix, prefix_len, nexthop);
|
|
self.trie.store(Arc::new(new_trie));
|
|
}
|
|
} |