diff --git a/Cargo.toml b/Cargo.toml index ed10daf..21e8c59 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,6 +4,6 @@ version = "0.1.0" edition = "2021" [dependencies] -tokio = { version = "1.43.0", features = ["macros", "rt-multi-thread", "net"] } +tokio = { version = "1.43.0", features = ["macros", "rt-multi-thread", "net", "sync"] } pnet = "0.35.0" ipnet = { version = "2", features = ["serde"] } diff --git a/src/global.rs b/src/global.rs index ea643dc..6a1f6cb 100644 --- a/src/global.rs +++ b/src/global.rs @@ -1,19 +1,21 @@ use core::net::{IpAddr, SocketAddr}; use ipnet::IpNet; use pnet::util::MacAddr; -use std::{collections::HashMap, sync::Mutex}; use tokio::net::UdpSocket; +use crate::redirect::MacTable; + pub struct Env { pub socket: UdpSocket, - pub vteps: Vec, - pub mac_table: Mutex>, -// pub local_nets: Vec, + pub used_vnis: (u32, u32), + // pub local_nets: Vec, pub reserved_nets: Vec, -// pub gateway: IpAddr, + // pub gateway: IpAddr, pub gateway_mac: MacAddr, pub gateway_vtep: IpAddr, pub gateway_vni: u32, + + pub mac_table: MacTable, } impl Env { diff --git a/src/main.rs b/src/main.rs index f9bfc64..4db857d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,12 +1,7 @@ use core::net::{IpAddr, SocketAddr}; use ipnet::IpNet; use pnet::util::MacAddr; -use std::{ - collections::HashMap, - io, - str::FromStr, - sync::{Arc, Mutex}, -}; +use std::{io, str::FromStr, sync::Arc}; use tokio::{net::UdpSocket, select, task}; mod global; @@ -24,40 +19,46 @@ fn net_parsing(net: &str) -> IpNet { } async fn treat_packet(env: Arc, raw: Vec, src_vtep: SocketAddr) -> Option<()> { + use packet::Action::*; let reading = packet::read_packet(&env, raw.as_slice(), src_vtep)?; - redirect::send(env, reading, raw).await; - Some(()) + match reading.action { + InternalTransmit(dst) => redirect::send(env, reading, dst, raw).await, + InternalBroadcast => redirect::broadcast(env, reading, raw).await, + GatewayArp => todo!(), + GatewayTransmit => todo!(), + GatewayEnter => todo!(), + } } #[tokio::main] async fn main() -> io::Result<()> { // TODO: config file let addr = "10.0.0.1"; - let vteps = ["10.0.0.2", "10.0.0.3"]; // will probably end in vtep learning let reserved_nets = [ "10.0.0.0/27", "10.0.253.0/24", "10.0.254.0/24", "10.0.255.0/24", ]; -// let local_nets = ["10.0.0.0/16"]; + // let local_nets = ["10.0.0.0/16"]; let gateway_mac = "02:00:00:00:00:01"; -// let gateway = "10.0.0.5"; + // let gateway = "10.0.0.5"; let gateway_vtep = "10.0.0.5"; let gateway_vni = 2; + let vni_range: (u32, u32) = (1000, 10000); let env = Arc::new(Env { socket: UdpSocket::bind(SocketAddr::new(addr_parsing(addr), 4789)).await?, - vteps: vteps.into_iter().map(addr_parsing).collect(), -// local_nets: local_nets.into_iter().map(net_parsing).collect(), + // local_nets: local_nets.into_iter().map(net_parsing).collect(), + used_vnis: vni_range, reserved_nets: reserved_nets.into_iter().map(net_parsing).collect(), -// gateway: addr_parsing(gateway), + // gateway: addr_parsing(gateway), gateway_vtep: addr_parsing(gateway_vtep), gateway_vni, gateway_mac: MacAddr::from_str(gateway_mac) .unwrap_or_else(|e| panic!("Error parsing {:?}: {}", gateway_mac, e)), - mac_table: Mutex::new(HashMap::new()), + mac_table: redirect::new_mactable(vni_range), }); let mut buf = vec![0u8; 2000]; diff --git a/src/packet.rs b/src/packet.rs index 00328cf..b8f19c1 100644 --- a/src/packet.rs +++ b/src/packet.rs @@ -17,7 +17,7 @@ use crate::redirect::mac_learn; /// Action to be taken for a packet pub enum Action { /// Transmit for the non-broadcast destination MAC - InternalTransmit, + InternalTransmit(MacAddr), /// Transmit to all known device on this VNI InternalBroadcast, @@ -36,8 +36,6 @@ pub enum Action { pub struct PacketReading { pub vtep: IpAddr, pub vni: u32, - pub src: MacAddr, - pub dst: MacAddr, pub action: Action, } @@ -47,11 +45,16 @@ pub fn read_packet(env: &Arc, payload: &[u8], src: SocketAddr) -> Option, payload: &[u8], src: SocketAddr) -> Option, mac: MacAddr, vni: u32, vtep: IpAddr) { - env.mac_table.lock().unwrap().insert((mac, vni), vtep); +pub type MacTable = Vec>>; + +pub fn new_mactable(vni_range: (u32, u32)) -> MacTable { + let mut mac_table = Vec::new(); + mac_table.resize_with((vni_range.1 - vni_range.0) as usize, || { + RwLock::new(HashMap::new()) + }); + mac_table } -pub async fn send(env: Arc, reading: PacketReading, packet: Vec) { - let svtep = env - .mac_table - .lock() - .unwrap() - .get(&(reading.dst, reading.vni)) - .copied(); - if let Some(vtep) = svtep { - env.send(&packet, vtep.clone()).await; +fn mac_table(env: &Arc, vni: u32) -> Option<&RwLock>> { + if env.used_vnis.0 <= vni && vni < env.used_vnis.1 { + Some(&env.mac_table[(vni - env.used_vnis.0) as usize]) } else { - for vtep in env.vteps.iter().filter(|vtep| **vtep != reading.vtep) { - env.send(&packet, *vtep).await; + None + } +} + +pub async fn mac_learn(env: Arc, mac: MacAddr, vni: u32, vtep: IpAddr) { + if let Some(vni_table) = mac_table(&env, vni) { + // avoid taking a write lock when unnecessary + if vni_table.read().await.get(&mac) != Some(&vtep) { + vni_table.write().await.insert(mac, vtep); } } } + +pub async fn send( + env: Arc, + reading: PacketReading, + dst: MacAddr, + packet: Vec, +) -> Option<()> { + let vtep = *mac_table(&env, reading.vni)?.read().await.get(&dst)?; + if vtep != reading.vtep { + env.send(&packet, vtep).await; + } + Some(()) +} + +pub async fn broadcast(env: Arc, reading: PacketReading, packet: Vec) -> Option<()> { + let table = mac_table(&env, reading.vni)?.read().await; + for vtep in table.values() { + if *vtep != reading.vtep { + // TODO: parallele send + env.send(&packet, *vtep).await; + } + } + Some(()) +}