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