1 use super::*; 2 use crate::error::*; 3 4 use async_trait::async_trait; 5 use std::net::IpAddr; 6 use util::vnet::net::*; 7 8 // RelayAddressGeneratorRanges can be used to only allocate connections inside a defined port range 9 pub struct RelayAddressGeneratorRanges { 10 // relay_address is the IP returned to the user when the relay is created 11 pub relay_address: IpAddr, 12 13 // min_port the minimum port to allocate 14 pub min_port: u16, 15 16 // max_port the maximum (inclusive) port to allocate 17 pub max_port: u16, 18 19 // max_retries the amount of tries to allocate a random port in the defined range 20 pub max_retries: u16, 21 22 // Address is passed to Listen/ListenPacket when creating the Relay 23 pub address: String, 24 25 pub net: Arc<Net>, 26 } 27 28 #[async_trait] 29 impl RelayAddressGenerator for RelayAddressGeneratorRanges { 30 // validate confirms that the RelayAddressGenerator is properly initialized validate(&self) -> Result<()>31 fn validate(&self) -> Result<()> { 32 if self.min_port == 0 { 33 Err(Error::ErrMinPortNotZero) 34 } else if self.max_port == 0 { 35 Err(Error::ErrMaxPortNotZero) 36 } else if self.max_port < self.min_port { 37 Err(Error::ErrMaxPortLessThanMinPort) 38 } else if self.address.is_empty() { 39 Err(Error::ErrListeningAddressInvalid) 40 } else { 41 Ok(()) 42 } 43 } 44 45 // Allocate a PacketConn (UDP) relay_address allocate_conn( &self, use_ipv4: bool, requested_port: u16, ) -> Result<(Arc<dyn Conn + Send + Sync>, SocketAddr)>46 async fn allocate_conn( 47 &self, 48 use_ipv4: bool, 49 requested_port: u16, 50 ) -> Result<(Arc<dyn Conn + Send + Sync>, SocketAddr)> { 51 let max_retries = if self.max_retries == 0 { 52 10 53 } else { 54 self.max_retries 55 }; 56 57 if requested_port != 0 { 58 let addr = self 59 .net 60 .resolve_addr(use_ipv4, &format!("{}:{}", self.address, requested_port)) 61 .await?; 62 let conn = self.net.bind(addr).await?; 63 let mut relay_addr = conn.local_addr()?; 64 relay_addr.set_ip(self.relay_address); 65 return Ok((conn, relay_addr)); 66 } 67 68 for _ in 0..max_retries { 69 let port = self.min_port + rand::random::<u16>() % (self.max_port - self.min_port + 1); 70 let addr = self 71 .net 72 .resolve_addr(use_ipv4, &format!("{}:{}", self.address, port)) 73 .await?; 74 let conn = match self.net.bind(addr).await { 75 Ok(conn) => conn, 76 Err(_) => continue, 77 }; 78 79 let mut relay_addr = conn.local_addr()?; 80 relay_addr.set_ip(self.relay_address); 81 return Ok((conn, relay_addr)); 82 } 83 84 Err(Error::ErrMaxRetriesExceeded) 85 } 86 } 87