1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90
use futures::channel::mpsc; use futures::future::Future; use futures::stream::Stream; use netsim_embed_core::{Ipv4Route, Plug}; use pnet_packet::ipv4::Ipv4Packet; use std::net::Ipv4Addr; use std::pin::Pin; use std::task::{Context, Poll}; #[derive(Debug)] pub struct Ipv4Router { addr: Ipv4Addr, rxs: Vec<mpsc::UnboundedReceiver<Vec<u8>>>, txs: Vec<(mpsc::UnboundedSender<Vec<u8>>, Vec<Ipv4Route>)>, } impl Ipv4Router { pub fn new(addr: Ipv4Addr) -> Self { Self { addr, rxs: Default::default(), txs: Default::default(), } } pub fn add_connection(&mut self, plug: Plug, routes: Vec<Ipv4Route>) { let (tx, rx) = plug.split(); self.rxs.push(rx); self.txs.push((tx, routes)); } fn process_packet(&mut self, bytes: Vec<u8>) { let packet = if let Some(packet) = Ipv4Packet::new(&bytes) { packet } else { log::info!("router {}: dropping invalid ipv4 packet", self.addr); return; }; let dest = packet.get_destination(); if dest == self.addr { log::info!("router {}: dropping packet addressed to me", self.addr); return; } for (tx, routes) in &self.txs { for route in routes { if route.dest().contains(dest) { log::debug!("router {}: routing packet on route {:?}", self.addr, route); let _ = tx.unbounded_send(bytes); return; } } } log::info!( "router {}: dropping unroutable packet to {}", self.addr, dest ); } } impl Future for Ipv4Router { type Output = (); fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> { let mut i = 0; while i < self.rxs.len() { loop { let packet = match Pin::new(&mut self.rxs[i]).poll_next(cx) { Poll::Pending => { i += 1; break; } Poll::Ready(None) => { self.rxs.remove(i); self.txs.remove(i); break; } Poll::Ready(Some(packet)) => packet, }; self.process_packet(packet) } } if self.rxs.is_empty() { return Poll::Ready(()); } Poll::Pending } }