1 use super::*; 2 use crate::vnet::chunk::*; 3 use crate::vnet::conn::*; 4 5 use async_trait::async_trait; 6 use std::net::IpAddr; 7 use std::str::FromStr; 8 9 #[derive(Default)] 10 struct DummyObserver; 11 12 #[async_trait] 13 impl ConnObserver for DummyObserver { 14 async fn write(&self, _c: Box<dyn Chunk + Send + Sync>) -> Result<()> { 15 Ok(()) 16 } 17 18 async fn on_closed(&self, _addr: SocketAddr) {} 19 20 fn determine_source_ip(&self, loc_ip: IpAddr, _dst_ip: IpAddr) -> Option<IpAddr> { 21 Some(loc_ip) 22 } 23 } 24 25 #[tokio::test] 26 async fn test_udp_conn_map_insert_remove() -> Result<()> { 27 let conn_map = UdpConnMap::new(); 28 29 let obs: Arc<Mutex<dyn ConnObserver + Send + Sync>> = 30 Arc::new(Mutex::new(DummyObserver::default())); 31 32 let conn_in = Arc::new(UdpConn::new( 33 SocketAddr::from_str("127.0.0.1:1234")?, 34 None, 35 obs, 36 )); 37 38 conn_map.insert(Arc::clone(&conn_in)).await?; 39 40 let conn_out = conn_map.find(&conn_in.local_addr().await?).await; 41 assert!(conn_out.is_some(), "should succeed"); 42 if let Some(conn_out) = conn_out { 43 assert_eq!( 44 conn_in.local_addr().await?, 45 conn_out.local_addr().await?, 46 "should match" 47 ); 48 let port_map = conn_map.port_map.lock().await; 49 assert_eq!(1, port_map.len(), "should match"); 50 } 51 52 conn_map.delete(&conn_in.local_addr().await?).await?; 53 { 54 let port_map = conn_map.port_map.lock().await; 55 assert_eq!(0, port_map.len(), "should match"); 56 } 57 58 let result = conn_map.delete(&conn_in.local_addr().await?).await; 59 assert!(result.is_err(), "should fail"); 60 61 Ok(()) 62 } 63 64 #[tokio::test] 65 async fn test_udp_conn_map_insert_0_remove() -> Result<()> { 66 let conn_map = UdpConnMap::new(); 67 68 let obs: Arc<Mutex<dyn ConnObserver + Send + Sync>> = 69 Arc::new(Mutex::new(DummyObserver::default())); 70 71 let conn_in = Arc::new(UdpConn::new( 72 SocketAddr::from_str("0.0.0.0:1234")?, 73 None, 74 obs, 75 )); 76 77 conn_map.insert(Arc::clone(&conn_in)).await?; 78 79 let conn_out = conn_map.find(&conn_in.local_addr().await?).await; 80 assert!(conn_out.is_some(), "should succeed"); 81 if let Some(conn_out) = conn_out { 82 assert_eq!( 83 conn_in.local_addr().await?, 84 conn_out.local_addr().await?, 85 "should match" 86 ); 87 let port_map = conn_map.port_map.lock().await; 88 assert_eq!(1, port_map.len(), "should match"); 89 } 90 91 conn_map.delete(&conn_in.local_addr().await?).await?; 92 { 93 let port_map = conn_map.port_map.lock().await; 94 assert_eq!(0, port_map.len(), "should match"); 95 } 96 97 let result = conn_map.delete(&conn_in.local_addr().await?).await; 98 assert!(result.is_err(), "should fail"); 99 100 Ok(()) 101 } 102 103 #[tokio::test] 104 async fn test_udp_conn_map_find_0() -> Result<()> { 105 let conn_map = UdpConnMap::new(); 106 107 let obs: Arc<Mutex<dyn ConnObserver + Send + Sync>> = 108 Arc::new(Mutex::new(DummyObserver::default())); 109 110 let conn_in = Arc::new(UdpConn::new( 111 SocketAddr::from_str("0.0.0.0:1234")?, 112 None, 113 obs, 114 )); 115 116 conn_map.insert(Arc::clone(&conn_in)).await?; 117 118 let addr = SocketAddr::from_str("192.168.0.1:1234")?; 119 let conn_out = conn_map.find(&addr).await; 120 assert!(conn_out.is_some(), "should succeed"); 121 if let Some(conn_out) = conn_out { 122 let addr_in = conn_in.local_addr().await?; 123 let addr_out = conn_out.local_addr().await?; 124 assert_eq!(addr_in, addr_out, "should match"); 125 let port_map = conn_map.port_map.lock().await; 126 assert_eq!(1, port_map.len(), "should match"); 127 } 128 129 Ok(()) 130 } 131 132 #[tokio::test] 133 async fn test_udp_conn_map_insert_many_ips_with_same_port() -> Result<()> { 134 let conn_map = UdpConnMap::new(); 135 136 let obs: Arc<Mutex<dyn ConnObserver + Send + Sync>> = 137 Arc::new(Mutex::new(DummyObserver::default())); 138 139 let conn_in1 = Arc::new(UdpConn::new( 140 SocketAddr::from_str("10.1.2.1:5678")?, 141 None, 142 Arc::clone(&obs), 143 )); 144 145 let conn_in2 = Arc::new(UdpConn::new( 146 SocketAddr::from_str("10.1.2.2:5678")?, 147 None, 148 Arc::clone(&obs), 149 )); 150 151 conn_map.insert(Arc::clone(&conn_in1)).await?; 152 conn_map.insert(Arc::clone(&conn_in2)).await?; 153 154 let addr1 = SocketAddr::from_str("10.1.2.1:5678")?; 155 let conn_out1 = conn_map.find(&addr1).await; 156 assert!(conn_out1.is_some(), "should succeed"); 157 if let Some(conn_out1) = conn_out1 { 158 let addr_in = conn_in1.local_addr().await?; 159 let addr_out = conn_out1.local_addr().await?; 160 assert_eq!(addr_in, addr_out, "should match"); 161 let port_map = conn_map.port_map.lock().await; 162 assert_eq!(1, port_map.len(), "should match"); 163 } 164 165 let addr2 = SocketAddr::from_str("10.1.2.2:5678")?; 166 let conn_out2 = conn_map.find(&addr2).await; 167 assert!(conn_out2.is_some(), "should succeed"); 168 if let Some(conn_out2) = conn_out2 { 169 let addr_in = conn_in2.local_addr().await?; 170 let addr_out = conn_out2.local_addr().await?; 171 assert_eq!(addr_in, addr_out, "should match"); 172 let port_map = conn_map.port_map.lock().await; 173 assert_eq!(1, port_map.len(), "should match"); 174 } 175 176 Ok(()) 177 } 178 179 #[tokio::test] 180 async fn test_udp_conn_map_already_inuse_when_insert_0() -> Result<()> { 181 let conn_map = UdpConnMap::new(); 182 183 let obs: Arc<Mutex<dyn ConnObserver + Send + Sync>> = 184 Arc::new(Mutex::new(DummyObserver::default())); 185 186 let conn_in1 = Arc::new(UdpConn::new( 187 SocketAddr::from_str("10.1.2.1:5678")?, 188 None, 189 Arc::clone(&obs), 190 )); 191 let conn_in2 = Arc::new(UdpConn::new( 192 SocketAddr::from_str("0.0.0.0:5678")?, 193 None, 194 Arc::clone(&obs), 195 )); 196 197 conn_map.insert(Arc::clone(&conn_in1)).await?; 198 let result = conn_map.insert(Arc::clone(&conn_in2)).await; 199 assert!(result.is_err(), "should fail"); 200 201 Ok(()) 202 } 203 204 #[tokio::test] 205 async fn test_udp_conn_map_already_inuse_when_insert_a_specified_ip() -> Result<()> { 206 let conn_map = UdpConnMap::new(); 207 208 let obs: Arc<Mutex<dyn ConnObserver + Send + Sync>> = 209 Arc::new(Mutex::new(DummyObserver::default())); 210 211 let conn_in1 = Arc::new(UdpConn::new( 212 SocketAddr::from_str("0.0.0.0:5678")?, 213 None, 214 Arc::clone(&obs), 215 )); 216 let conn_in2 = Arc::new(UdpConn::new( 217 SocketAddr::from_str("192.168.0.1:5678")?, 218 None, 219 Arc::clone(&obs), 220 )); 221 222 conn_map.insert(Arc::clone(&conn_in1)).await?; 223 let result = conn_map.insert(Arc::clone(&conn_in2)).await; 224 assert!(result.is_err(), "should fail"); 225 226 Ok(()) 227 } 228 229 #[tokio::test] 230 async fn test_udp_conn_map_already_inuse_when_insert_same_specified_ip() -> Result<()> { 231 let conn_map = UdpConnMap::new(); 232 233 let obs: Arc<Mutex<dyn ConnObserver + Send + Sync>> = 234 Arc::new(Mutex::new(DummyObserver::default())); 235 236 let conn_in1 = Arc::new(UdpConn::new( 237 SocketAddr::from_str("192.168.0.1:5678")?, 238 None, 239 Arc::clone(&obs), 240 )); 241 let conn_in2 = Arc::new(UdpConn::new( 242 SocketAddr::from_str("192.168.0.1:5678")?, 243 None, 244 Arc::clone(&obs), 245 )); 246 247 conn_map.insert(Arc::clone(&conn_in1)).await?; 248 let result = conn_map.insert(Arc::clone(&conn_in2)).await; 249 assert!(result.is_err(), "should fail"); 250 251 Ok(()) 252 } 253 254 #[tokio::test] 255 async fn test_udp_conn_map_find_failure_1() -> Result<()> { 256 let conn_map = UdpConnMap::new(); 257 258 let obs: Arc<Mutex<dyn ConnObserver + Send + Sync>> = 259 Arc::new(Mutex::new(DummyObserver::default())); 260 261 let conn_in = Arc::new(UdpConn::new( 262 SocketAddr::from_str("192.168.0.1:5678")?, 263 None, 264 obs, 265 )); 266 267 conn_map.insert(Arc::clone(&conn_in)).await?; 268 269 let addr = SocketAddr::from_str("192.168.0.2:5678")?; 270 let result = conn_map.find(&addr).await; 271 assert!(result.is_none(), "should be none"); 272 273 Ok(()) 274 } 275 276 #[tokio::test] 277 async fn test_udp_conn_map_find_failure_2() -> Result<()> { 278 let conn_map = UdpConnMap::new(); 279 280 let obs: Arc<Mutex<dyn ConnObserver + Send + Sync>> = 281 Arc::new(Mutex::new(DummyObserver::default())); 282 283 let conn_in = Arc::new(UdpConn::new( 284 SocketAddr::from_str("192.168.0.1:5678")?, 285 None, 286 obs, 287 )); 288 289 conn_map.insert(Arc::clone(&conn_in)).await?; 290 291 let addr = SocketAddr::from_str("192.168.0.1:1234")?; 292 let result = conn_map.find(&addr).await; 293 assert!(result.is_none(), "should be none"); 294 295 Ok(()) 296 } 297 298 #[tokio::test] 299 async fn test_udp_conn_map_insert_two_on_same_port_then_remove() -> Result<()> { 300 let conn_map = UdpConnMap::new(); 301 302 let obs: Arc<Mutex<dyn ConnObserver + Send + Sync>> = 303 Arc::new(Mutex::new(DummyObserver::default())); 304 305 let conn_in1 = Arc::new(UdpConn::new( 306 SocketAddr::from_str("192.168.0.1:5678")?, 307 None, 308 Arc::clone(&obs), 309 )); 310 let conn_in2 = Arc::new(UdpConn::new( 311 SocketAddr::from_str("192.168.0.2:5678")?, 312 None, 313 Arc::clone(&obs), 314 )); 315 316 conn_map.insert(Arc::clone(&conn_in1)).await?; 317 conn_map.insert(Arc::clone(&conn_in2)).await?; 318 319 conn_map.delete(&conn_in1.local_addr().await?).await?; 320 conn_map.delete(&conn_in2.local_addr().await?).await?; 321 322 Ok(()) 323 } 324