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 {}, but {:?}", 31 lifetime, lifetime_duration 32 ); 33 34 Ok(()) 35 } 36 37 #[tokio::test] 38 async fn test_allocation_lifetime_overflow() -> Result<()> { 39 let lifetime = Lifetime(MAXIMUM_ALLOCATION_LIFETIME * 2); 40 41 let mut m2 = Message::new(); 42 lifetime.add_to(&mut m2)?; 43 44 let lifetime_duration = allocation_lifetime(&m2); 45 assert_eq!( 46 lifetime_duration, DEFAULT_LIFETIME, 47 "Expect lifetime_duration is {:?}, but {:?}", 48 DEFAULT_LIFETIME, lifetime_duration 49 ); 50 51 Ok(()) 52 } 53 54 struct TestAuthHandler; 55 impl AuthHandler for TestAuthHandler { 56 fn auth_handle(&self, _username: &str, _realm: &str, _src_addr: SocketAddr) -> Result<Vec<u8>> { 57 Ok(STATIC_KEY.as_bytes().to_vec()) 58 } 59 } 60 61 #[tokio::test] 62 async fn test_allocation_lifetime_deletion_zero_lifetime() -> Result<()> { 63 //env_logger::init(); 64 65 let l = Arc::new(UdpSocket::bind("0.0.0.0:0").await?); 66 67 let allocation_manager = Arc::new(Manager::new(ManagerConfig { 68 relay_addr_generator: Box::new(RelayAddressGeneratorNone { 69 address: "0.0.0.0".to_owned(), 70 net: Arc::new(Net::new(None)), 71 }), 72 })); 73 74 let socket = SocketAddr::new(IpAddr::from_str("127.0.0.1")?, 5000); 75 76 let mut r = Request::new(l, socket, allocation_manager, Arc::new(TestAuthHandler {})); 77 78 { 79 let mut nonces = r.nonces.lock().await; 80 nonces.insert(STATIC_KEY.to_owned(), Instant::now()); 81 } 82 83 let five_tuple = FiveTuple { 84 src_addr: r.src_addr, 85 dst_addr: r.conn.local_addr()?, 86 protocol: PROTO_UDP, 87 }; 88 89 r.allocation_manager 90 .create_allocation( 91 five_tuple, 92 Arc::clone(&r.conn), 93 0, 94 Duration::from_secs(3600), 95 TextAttribute::new(ATTR_USERNAME, "user".into()), 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