feat: wip! wip! gateway

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"] }

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<IpAddr>,
-    pub mac_table: Mutex<HashMap<(MacAddr, u32), IpAddr>>,
-//    pub local_nets: Vec<IpNet>,
+    pub used_vnis: (u32, u32),
+    // pub local_nets: Vec<IpNet>,
     pub reserved_nets: Vec<IpNet>,
-//    pub gateway: IpAddr,
+    // pub gateway: IpAddr,
     pub gateway_mac: MacAddr,
     pub gateway_vtep: IpAddr,
     pub gateway_vni: u32,
+
+    pub mac_table: MacTable,
 }
 
 impl Env {

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<Env>, raw: Vec<u8>, 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];

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<Env>, payload: &[u8], src: SocketAddr) -> Option<Pa
     let vxlan = VxlanPacket::new(payload)?;
     let eth = EthernetPacket::new(vxlan.payload())?;
 
-    let mut action = InternalTransmit;
+    let mut action = InternalTransmit(eth.get_destination());
     if vxlan.get_vni() == env.gateway_vni {
         action = GatewayEnter;
     } else {
-        mac_learn(env, eth.get_source(), vxlan.get_vni(), src.ip());
+        tokio::spawn(mac_learn(
+            env.clone(),
+            eth.get_source(),
+            vxlan.get_vni(),
+            src.ip(),
+        ));
         if eth.get_destination().is_broadcast() {
             action = InternalBroadcast;
             if eth.get_ethertype() == EtherTypes::Arp {
@@ -78,8 +81,6 @@ pub fn read_packet(env: &Arc<Env>, payload: &[u8], src: SocketAddr) -> Option<Pa
     Some(PacketReading {
         vtep: src.ip(),
         vni: vxlan.get_vni(),
-        src: eth.get_source(),
-        dst: eth.get_destination(),
         action,
     })
 }

src/redirect.rs

@@ -1,25 +1,57 @@
-use crate::global::Env;
 use pnet::util::MacAddr;
-use std::{net::IpAddr, sync::Arc};
+use std::{collections::HashMap, net::IpAddr, sync::Arc};
+use tokio::sync::RwLock;
 
+use crate::global::Env;
 use crate::packet::PacketReading;
 
-pub fn mac_learn(env: &Arc<Env>, mac: MacAddr, vni: u32, vtep: IpAddr) {
-    env.mac_table.lock().unwrap().insert((mac, vni), vtep);
+pub type MacTable = Vec<RwLock<HashMap<MacAddr, IpAddr>>>;
+
+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<Env>, reading: PacketReading, packet: Vec<u8>) {
-    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<Env>, vni: u32) -> Option<&RwLock<HashMap<MacAddr, IpAddr>>> {
+    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<Env>, 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<Env>,
+    reading: PacketReading,
+    dst: MacAddr,
+    packet: Vec<u8>,
+) -> 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<Env>, reading: PacketReading, packet: Vec<u8>) -> 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(())
+}