1 use super::*;
2 
3 use crate::proto::lifetime::DEFAULT_LIFETIME;
4 use std::str::FromStr;
5 use stun::{attributes::ATTR_USERNAME, textattrs::TextAttribute};
6 use tokio::net::UdpSocket;
7 
8 #[tokio::test]
9 async fn test_has_permission() -> Result<()> {
10     let turn_socket = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
11     let relay_socket = Arc::clone(&turn_socket);
12     let relay_addr = relay_socket.local_addr()?;
13     let a = Allocation::new(
14         turn_socket,
15         relay_socket,
16         relay_addr,
17         FiveTuple::default(),
18         TextAttribute::new(ATTR_USERNAME, "user".into()),
19     );
20 
21     let addr1 = SocketAddr::from_str("127.0.0.1:3478")?;
22     let addr2 = SocketAddr::from_str("127.0.0.1:3479")?;
23     let addr3 = SocketAddr::from_str("127.0.0.2:3478")?;
24 
25     let p1 = Permission::new(addr1);
26     let p2 = Permission::new(addr2);
27     let p3 = Permission::new(addr3);
28 
29     a.add_permission(p1).await;
30     a.add_permission(p2).await;
31     a.add_permission(p3).await;
32 
33     let found_p1 = a.has_permission(&addr1).await;
34     assert!(found_p1, "Should keep the first one.");
35 
36     let found_p2 = a.has_permission(&addr2).await;
37     assert!(found_p2, "Second one should be ignored.");
38 
39     let found_p3 = a.has_permission(&addr3).await;
40     assert!(found_p3, "Permission with another IP should be found");
41 
42     Ok(())
43 }
44 
45 #[tokio::test]
46 async fn test_add_permission() -> Result<()> {
47     let turn_socket = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
48     let relay_socket = Arc::clone(&turn_socket);
49     let relay_addr = relay_socket.local_addr()?;
50     let a = Allocation::new(
51         turn_socket,
52         relay_socket,
53         relay_addr,
54         FiveTuple::default(),
55         TextAttribute::new(ATTR_USERNAME, "user".into()),
56     );
57 
58     let addr = SocketAddr::from_str("127.0.0.1:3478")?;
59     let p = Permission::new(addr);
60     a.add_permission(p).await;
61 
62     let found_p = a.has_permission(&addr).await;
63     assert!(found_p, "Should keep the first one.");
64 
65     Ok(())
66 }
67 
68 #[tokio::test]
69 async fn test_remove_permission() -> Result<()> {
70     let turn_socket = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
71     let relay_socket = Arc::clone(&turn_socket);
72     let relay_addr = relay_socket.local_addr()?;
73     let a = Allocation::new(
74         turn_socket,
75         relay_socket,
76         relay_addr,
77         FiveTuple::default(),
78         TextAttribute::new(ATTR_USERNAME, "user".into()),
79     );
80 
81     let addr = SocketAddr::from_str("127.0.0.1:3478")?;
82 
83     let p = Permission::new(addr);
84     a.add_permission(p).await;
85 
86     let found_p = a.has_permission(&addr).await;
87     assert!(found_p, "Should keep the first one.");
88 
89     a.remove_permission(&addr).await;
90 
91     let found_permission = a.has_permission(&addr).await;
92     assert!(
93         !found_permission,
94         "Got permission should be nil after removed."
95     );
96 
97     Ok(())
98 }
99 
100 #[tokio::test]
101 async fn test_add_channel_bind() -> Result<()> {
102     let turn_socket = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
103     let relay_socket = Arc::clone(&turn_socket);
104     let relay_addr = relay_socket.local_addr()?;
105     let a = Allocation::new(
106         turn_socket,
107         relay_socket,
108         relay_addr,
109         FiveTuple::default(),
110         TextAttribute::new(ATTR_USERNAME, "user".into()),
111     );
112 
113     let addr = SocketAddr::from_str("127.0.0.1:3478")?;
114     let c = ChannelBind::new(ChannelNumber(MIN_CHANNEL_NUMBER), addr);
115 
116     a.add_channel_bind(c, DEFAULT_LIFETIME).await?;
117 
118     let c2 = ChannelBind::new(ChannelNumber(MIN_CHANNEL_NUMBER + 1), addr);
119     let result = a.add_channel_bind(c2, DEFAULT_LIFETIME).await;
120     assert!(
121         result.is_err(),
122         "should failed with conflicted peer address"
123     );
124 
125     let addr2 = SocketAddr::from_str("127.0.0.1:3479")?;
126     let c3 = ChannelBind::new(ChannelNumber(MIN_CHANNEL_NUMBER), addr2);
127     let result = a.add_channel_bind(c3, DEFAULT_LIFETIME).await;
128     assert!(result.is_err(), "should fail with conflicted number.");
129 
130     Ok(())
131 }
132 
133 #[tokio::test]
134 async fn test_get_channel_by_number() -> Result<()> {
135     let turn_socket = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
136     let relay_socket = Arc::clone(&turn_socket);
137     let relay_addr = relay_socket.local_addr()?;
138     let a = Allocation::new(
139         turn_socket,
140         relay_socket,
141         relay_addr,
142         FiveTuple::default(),
143         TextAttribute::new(ATTR_USERNAME, "user".into()),
144     );
145 
146     let addr = SocketAddr::from_str("127.0.0.1:3478")?;
147     let c = ChannelBind::new(ChannelNumber(MIN_CHANNEL_NUMBER), addr);
148 
149     a.add_channel_bind(c, DEFAULT_LIFETIME).await?;
150 
151     let exist_channel_addr = a
152         .get_channel_addr(&ChannelNumber(MIN_CHANNEL_NUMBER))
153         .await
154         .unwrap();
155     assert_eq!(addr, exist_channel_addr);
156 
157     let not_exist_channel = a
158         .get_channel_addr(&ChannelNumber(MIN_CHANNEL_NUMBER + 1))
159         .await;
160     assert!(
161         not_exist_channel.is_none(),
162         "should be nil for not existed channel."
163     );
164 
165     Ok(())
166 }
167 
168 #[tokio::test]
169 async fn test_get_channel_by_addr() -> Result<()> {
170     let turn_socket = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
171     let relay_socket = Arc::clone(&turn_socket);
172     let relay_addr = relay_socket.local_addr()?;
173     let a = Allocation::new(
174         turn_socket,
175         relay_socket,
176         relay_addr,
177         FiveTuple::default(),
178         TextAttribute::new(ATTR_USERNAME, "user".into()),
179     );
180 
181     let addr = SocketAddr::from_str("127.0.0.1:3478")?;
182     let addr2 = SocketAddr::from_str("127.0.0.1:3479")?;
183     let c = ChannelBind::new(ChannelNumber(MIN_CHANNEL_NUMBER), addr);
184 
185     a.add_channel_bind(c, DEFAULT_LIFETIME).await?;
186 
187     let exist_channel_number = a.get_channel_number(&addr).await.unwrap();
188     assert_eq!(ChannelNumber(MIN_CHANNEL_NUMBER), exist_channel_number);
189 
190     let not_exist_channel = a.get_channel_number(&addr2).await;
191     assert!(
192         not_exist_channel.is_none(),
193         "should be nil for not existed channel."
194     );
195 
196     Ok(())
197 }
198 
199 #[tokio::test]
200 async fn test_remove_channel_bind() -> Result<()> {
201     let turn_socket = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
202     let relay_socket = Arc::clone(&turn_socket);
203     let relay_addr = relay_socket.local_addr()?;
204     let a = Allocation::new(
205         turn_socket,
206         relay_socket,
207         relay_addr,
208         FiveTuple::default(),
209         TextAttribute::new(ATTR_USERNAME, "user".into()),
210     );
211 
212     let addr = SocketAddr::from_str("127.0.0.1:3478")?;
213     let number = ChannelNumber(MIN_CHANNEL_NUMBER);
214     let c = ChannelBind::new(number, addr);
215 
216     a.add_channel_bind(c, DEFAULT_LIFETIME).await?;
217 
218     a.remove_channel_bind(number).await;
219 
220     let not_exist_channel = a.get_channel_addr(&number).await;
221     assert!(
222         not_exist_channel.is_none(),
223         "should be nil for not existed channel."
224     );
225 
226     let not_exist_channel = a.get_channel_number(&addr).await;
227     assert!(
228         not_exist_channel.is_none(),
229         "should be nil for not existed channel."
230     );
231 
232     Ok(())
233 }
234 
235 #[tokio::test]
236 async fn test_allocation_refresh() -> Result<()> {
237     let turn_socket = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
238     let relay_socket = Arc::clone(&turn_socket);
239     let relay_addr = relay_socket.local_addr()?;
240     let a = Allocation::new(
241         turn_socket,
242         relay_socket,
243         relay_addr,
244         FiveTuple::default(),
245         TextAttribute::new(ATTR_USERNAME, "user".into()),
246     );
247 
248     a.start(DEFAULT_LIFETIME).await;
249     a.refresh(Duration::from_secs(0)).await;
250 
251     assert!(!a.stop(), "lifetimeTimer has expired");
252 
253     Ok(())
254 }
255 
256 #[tokio::test]
257 async fn test_allocation_close() -> Result<()> {
258     let turn_socket = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
259     let relay_socket = Arc::clone(&turn_socket);
260     let relay_addr = relay_socket.local_addr()?;
261     let a = Allocation::new(
262         turn_socket,
263         relay_socket,
264         relay_addr,
265         FiveTuple::default(),
266         TextAttribute::new(ATTR_USERNAME, "user".into()),
267     );
268 
269     // add mock lifetimeTimer
270     a.start(DEFAULT_LIFETIME).await;
271 
272     // add channel
273     let addr = SocketAddr::from_str("127.0.0.1:3478")?;
274     let number = ChannelNumber(MIN_CHANNEL_NUMBER);
275     let c = ChannelBind::new(number, addr);
276 
277     a.add_channel_bind(c, DEFAULT_LIFETIME).await?;
278 
279     // add permission
280     a.add_permission(Permission::new(addr)).await;
281 
282     a.close().await?;
283 
284     Ok(())
285 }
286