1 use super::*; 2 use crate::relay::relay_none::*; 3 4 use std::{net::IpAddr, str::FromStr}; 5 use tokio::{ 6 net::UdpSocket, 7 time::{Duration, Instant}, 8 }; 9 use util::vnet::net::*; 10 11 const STATIC_KEY: &str = "ABC"; 12 13 #[tokio::test] 14 async fn test_allocation_lifetime_parsing() -> Result<()> { 15 let lifetime = Lifetime(Duration::from_secs(5)); 16 17 let mut m = Message::new(); 18 let lifetime_duration = allocation_lifetime(&m); 19 20 assert_eq!( 21 lifetime_duration, DEFAULT_LIFETIME, 22 "Allocation lifetime should be default time duration" 23 ); 24 25 lifetime.add_to(&mut m)?; 26 27 let lifetime_duration = allocation_lifetime(&m); 28 assert_eq!( 29 lifetime_duration, lifetime.0, 30 "Expect lifetime_duration is {lifetime}, but {lifetime_duration:?}" 31 ); 32 33 Ok(()) 34 } 35 36 #[tokio::test] 37 async fn test_allocation_lifetime_overflow() -> Result<()> { 38 let lifetime = Lifetime(MAXIMUM_ALLOCATION_LIFETIME * 2); 39 40 let mut m2 = Message::new(); 41 lifetime.add_to(&mut m2)?; 42 43 let lifetime_duration = allocation_lifetime(&m2); 44 assert_eq!( 45 lifetime_duration, DEFAULT_LIFETIME, 46 "Expect lifetime_duration is {DEFAULT_LIFETIME:?}, but {lifetime_duration:?}" 47 ); 48 49 Ok(()) 50 } 51 52 struct TestAuthHandler; 53 impl AuthHandler for TestAuthHandler { 54 fn auth_handle(&self, _username: &str, _realm: &str, _src_addr: SocketAddr) -> Result<Vec<u8>> { 55 Ok(STATIC_KEY.as_bytes().to_vec()) 56 } 57 } 58 59 #[tokio::test] 60 async fn test_allocation_lifetime_deletion_zero_lifetime() -> Result<()> { 61 //env_logger::init(); 62 63 let l = Arc::new(UdpSocket::bind("0.0.0.0:0").await?); 64 65 let allocation_manager = Arc::new(Manager::new(ManagerConfig { 66 relay_addr_generator: Box::new(RelayAddressGeneratorNone { 67 address: "0.0.0.0".to_owned(), 68 net: Arc::new(Net::new(None)), 69 }), 70 })); 71 72 let socket = SocketAddr::new(IpAddr::from_str("127.0.0.1")?, 5000); 73 74 let mut r = Request::new(l, socket, allocation_manager, Arc::new(TestAuthHandler {})); 75 76 { 77 let mut nonces = r.nonces.lock().await; 78 nonces.insert(STATIC_KEY.to_owned(), Instant::now()); 79 } 80 81 let five_tuple = FiveTuple { 82 src_addr: r.src_addr, 83 dst_addr: r.conn.local_addr()?, 84 protocol: PROTO_UDP, 85 }; 86 87 r.allocation_manager 88 .create_allocation( 89 five_tuple, 90 Arc::clone(&r.conn), 91 0, 92 Duration::from_secs(3600), 93 TextAttribute::new(ATTR_USERNAME, "user".into()), 94 ) 95 .await?; 96 assert!(r 97 .allocation_manager 98 .get_allocation(&five_tuple) 99 .await 100 .is_some()); 101 102 let mut m = Message::new(); 103 Lifetime::default().add_to(&mut m)?; 104 MessageIntegrity(STATIC_KEY.as_bytes().to_vec()).add_to(&mut m)?; 105 Nonce::new(ATTR_NONCE, STATIC_KEY.to_owned()).add_to(&mut m)?; 106 Realm::new(ATTR_REALM, STATIC_KEY.to_owned()).add_to(&mut m)?; 107 Username::new(ATTR_USERNAME, STATIC_KEY.to_owned()).add_to(&mut m)?; 108 109 r.handle_refresh_request(&m).await?; 110 assert!(r 111 .allocation_manager 112 .get_allocation(&five_tuple) 113 .await 114 .is_none()); 115 116 Ok(()) 117 } 118