1 #[cfg(test)] 2 mod conn_map_test; 3 4 use crate::error::*; 5 use crate::vnet::conn::UdpConn; 6 use crate::Conn; 7 8 use std::collections::HashMap; 9 use std::net::SocketAddr; 10 use std::sync::Arc; 11 use tokio::sync::Mutex; 12 13 type PortMap = Mutex<HashMap<u16, Vec<Arc<UdpConn>>>>; 14 15 #[derive(Default)] 16 pub(crate) struct UdpConnMap { 17 port_map: PortMap, 18 } 19 20 impl UdpConnMap { new() -> Self21 pub(crate) fn new() -> Self { 22 UdpConnMap { 23 port_map: Mutex::new(HashMap::new()), 24 } 25 } 26 insert(&self, conn: Arc<UdpConn>) -> Result<()>27 pub(crate) async fn insert(&self, conn: Arc<UdpConn>) -> Result<()> { 28 let addr = conn.local_addr()?; 29 30 let mut port_map = self.port_map.lock().await; 31 if let Some(conns) = port_map.get(&addr.port()) { 32 if addr.ip().is_unspecified() { 33 return Err(Error::ErrAddressAlreadyInUse); 34 } 35 36 for c in conns { 37 let laddr = c.local_addr()?; 38 if laddr.ip().is_unspecified() || laddr.ip() == addr.ip() { 39 return Err(Error::ErrAddressAlreadyInUse); 40 } 41 } 42 } 43 44 if let Some(conns) = port_map.get_mut(&addr.port()) { 45 conns.push(conn); 46 } else { 47 port_map.insert(addr.port(), vec![conn]); 48 } 49 Ok(()) 50 } 51 find(&self, addr: &SocketAddr) -> Option<Arc<UdpConn>>52 pub(crate) async fn find(&self, addr: &SocketAddr) -> Option<Arc<UdpConn>> { 53 let port_map = self.port_map.lock().await; 54 if let Some(conns) = port_map.get(&addr.port()) { 55 if addr.ip().is_unspecified() { 56 // pick the first one appears in the iteration 57 if let Some(c) = conns.first() { 58 return Some(Arc::clone(c)); 59 } else { 60 return None; 61 } 62 } 63 64 for c in conns { 65 let laddr = { 66 match c.local_addr() { 67 Ok(laddr) => laddr, 68 Err(_) => return None, 69 } 70 }; 71 if laddr.ip().is_unspecified() || laddr.ip() == addr.ip() { 72 return Some(Arc::clone(c)); 73 } 74 } 75 } 76 77 None 78 } 79 delete(&self, addr: &SocketAddr) -> Result<()>80 pub(crate) async fn delete(&self, addr: &SocketAddr) -> Result<()> { 81 let mut port_map = self.port_map.lock().await; 82 let mut new_conns = vec![]; 83 if let Some(conns) = port_map.get(&addr.port()) { 84 if !addr.ip().is_unspecified() { 85 for c in conns { 86 let laddr = c.local_addr()?; 87 if laddr.ip().is_unspecified() { 88 // This can't happen! 89 return Err(Error::ErrCannotRemoveUnspecifiedIp); 90 } 91 92 if laddr.ip() == addr.ip() { 93 continue; 94 } 95 new_conns.push(Arc::clone(c)); 96 } 97 } 98 } else { 99 return Err(Error::ErrNoSuchUdpConn); 100 } 101 102 if new_conns.is_empty() { 103 port_map.remove(&addr.port()); 104 } else { 105 port_map.insert(addr.port(), new_conns); 106 } 107 108 Ok(()) 109 } 110 len(&self) -> usize111 pub(crate) async fn len(&self) -> usize { 112 let port_map = self.port_map.lock().await; 113 let mut n = 0; 114 for conns in port_map.values() { 115 n += conns.len(); 116 } 117 n 118 } 119 } 120