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 {
write(&self, _c: Box<dyn Chunk + Send + Sync>) -> Result<()>14 async fn write(&self, _c: Box<dyn Chunk + Send + Sync>) -> Result<()> {
15 Ok(())
16 }
17
on_closed(&self, _addr: SocketAddr)18 async fn on_closed(&self, _addr: SocketAddr) {}
19
determine_source_ip(&self, loc_ip: IpAddr, _dst_ip: IpAddr) -> Option<IpAddr>20 fn determine_source_ip(&self, loc_ip: IpAddr, _dst_ip: IpAddr) -> Option<IpAddr> {
21 Some(loc_ip)
22 }
23 }
24
25 #[tokio::test]
test_udp_conn_map_insert_remove() -> Result<()>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;
41 assert!(conn_out.is_some(), "should succeed");
42 if let Some(conn_out) = conn_out {
43 assert_eq!(
44 conn_in.local_addr()?,
45 conn_out.local_addr()?,
46 "should match"
47 );
48 let port_map = conn_map.port_map.lock().await;
49 assert_eq!(port_map.len(), 1, "should match");
50 }
51
52 conn_map.delete(&conn_in.local_addr()?).await?;
53 {
54 let port_map = conn_map.port_map.lock().await;
55 assert_eq!(port_map.len(), 0, "should match");
56 }
57
58 let result = conn_map.delete(&conn_in.local_addr()?).await;
59 assert!(result.is_err(), "should fail");
60
61 Ok(())
62 }
63
64 #[tokio::test]
test_udp_conn_map_insert_0_remove() -> Result<()>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;
80 assert!(conn_out.is_some(), "should succeed");
81 if let Some(conn_out) = conn_out {
82 assert_eq!(
83 conn_in.local_addr()?,
84 conn_out.local_addr()?,
85 "should match"
86 );
87 let port_map = conn_map.port_map.lock().await;
88 assert_eq!(port_map.len(), 1, "should match");
89 }
90
91 conn_map.delete(&conn_in.local_addr()?).await?;
92 {
93 let port_map = conn_map.port_map.lock().await;
94 assert_eq!(port_map.len(), 0, "should match");
95 }
96
97 let result = conn_map.delete(&conn_in.local_addr()?).await;
98 assert!(result.is_err(), "should fail");
99
100 Ok(())
101 }
102
103 #[tokio::test]
test_udp_conn_map_find_0() -> Result<()>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()?;
123 let addr_out = conn_out.local_addr()?;
124 assert_eq!(addr_in, addr_out, "should match");
125 let port_map = conn_map.port_map.lock().await;
126 assert_eq!(port_map.len(), 1, "should match");
127 }
128
129 Ok(())
130 }
131
132 #[tokio::test]
test_udp_conn_map_insert_many_ips_with_same_port() -> Result<()>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()?;
159 let addr_out = conn_out1.local_addr()?;
160 assert_eq!(addr_in, addr_out, "should match");
161 let port_map = conn_map.port_map.lock().await;
162 assert_eq!(port_map.len(), 1, "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()?;
170 let addr_out = conn_out2.local_addr()?;
171 assert_eq!(addr_in, addr_out, "should match");
172 let port_map = conn_map.port_map.lock().await;
173 assert_eq!(port_map.len(), 1, "should match");
174 }
175
176 Ok(())
177 }
178
179 #[tokio::test]
test_udp_conn_map_already_inuse_when_insert_0() -> Result<()>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]
test_udp_conn_map_already_inuse_when_insert_a_specified_ip() -> Result<()>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]
test_udp_conn_map_already_inuse_when_insert_same_specified_ip() -> Result<()>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]
test_udp_conn_map_find_failure_1() -> Result<()>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]
test_udp_conn_map_find_failure_2() -> Result<()>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]
test_udp_conn_map_insert_two_on_same_port_then_remove() -> Result<()>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?;
320 conn_map.delete(&conn_in2.local_addr()?).await?;
321
322 Ok(())
323 }
324