diff --git a/rust/connlib/shared/src/lib.rs b/rust/connlib/shared/src/lib.rs index 653ce86913..37ec6fbf64 100644 --- a/rust/connlib/shared/src/lib.rs +++ b/rust/connlib/shared/src/lib.rs @@ -27,12 +27,7 @@ pub use error::ConnlibError as Error; pub use error::Result; pub use phoenix_channel::{LoginUrl, LoginUrlError}; -use ip_network::Ipv4Network; -use ip_network::Ipv6Network; use rand_core::OsRng; -use std::net::IpAddr; -use std::net::Ipv4Addr; -use std::net::Ipv6Addr; pub type Dname = domain::base::Dname>; @@ -61,35 +56,6 @@ pub fn keypair() -> (StaticSecret, PublicKey) { (private_key, public_key) } -pub struct IpProvider { - ipv4: Box + Send + Sync>, - ipv6: Box + Send + Sync>, -} - -impl IpProvider { - pub fn new(ipv4: Ipv4Network, ipv6: Ipv6Network) -> Self { - Self { - ipv4: Box::new(ipv4.hosts()), - ipv6: Box::new(ipv6.subnets_with_prefix(128).map(|ip| ip.network_address())), - } - } - - pub fn get_proxy_ip_for(&mut self, ip: &IpAddr) -> Option { - let proxy_ip = match ip { - IpAddr::V4(_) => self.ipv4.next().map(Into::into), - IpAddr::V6(_) => self.ipv6.next().map(Into::into), - }; - - if proxy_ip.is_none() { - // TODO: we might want to make the iterator cyclic or another strategy to prevent ip exhaustion - // this might happen in ipv4 if tokens are too long lived. - tracing::error!("IP exhaustion: Please reset your client"); - } - - proxy_ip - } -} - pub fn get_user_agent(os_version_override: Option) -> String { // Note: we could switch to sys-info and get the hostname // but we lose the arch diff --git a/rust/connlib/tunnel/src/client.rs b/rust/connlib/tunnel/src/client.rs index 0b5f48eaff..9502c1b059 100644 --- a/rust/connlib/tunnel/src/client.rs +++ b/rust/connlib/tunnel/src/client.rs @@ -9,9 +9,9 @@ use connlib_shared::messages::{ IpDnsServer, Key, Offer, Relay, RequestConnection, ResourceDescription, ResourceDescriptionCidr, ResourceDescriptionDns, ResourceId, ReuseConnection, }; -use connlib_shared::{Callbacks, Dname, IpProvider, PublicKey, StaticSecret}; +use connlib_shared::{Callbacks, Dname, PublicKey, StaticSecret}; use domain::base::Rtype; -use ip_network::IpNetwork; +use ip_network::{IpNetwork, Ipv4Network, Ipv6Network}; use ip_network_table::IpNetworkTable; use itertools::Itertools; @@ -511,6 +511,8 @@ impl ClientState { ip_provider: IpProvider::new( IPV4_RESOURCES.parse().unwrap(), IPV6_RESOURCES.parse().unwrap(), + Some(DNS_SENTINELS_V4.parse().unwrap()), + Some(DNS_SENTINELS_V6.parse().unwrap()), ), dns_resources_internal_ips: Default::default(), dns_resources: Default::default(), @@ -969,6 +971,8 @@ fn sentinel_dns_mapping(dns: &[DnsServer]) -> BiMap { let mut ip_provider = IpProvider::new( DNS_SENTINELS_V4.parse().unwrap(), DNS_SENTINELS_V6.parse().unwrap(), + None, + None, ); dns.iter() @@ -1007,6 +1011,47 @@ fn is_definitely_not_a_resource(ip: IpAddr) -> bool { false } +pub struct IpProvider { + ipv4: Box + Send + Sync>, + ipv6: Box + Send + Sync>, +} + +impl IpProvider { + pub fn new( + ipv4: Ipv4Network, + ipv6: Ipv6Network, + exclusion_v4: Option, + exclusion_v6: Option, + ) -> Self { + Self { + ipv4: Box::new( + ipv4.hosts() + .filter(move |ip| !exclusion_v4.is_some_and(|e| e.contains(*ip))), + ), + ipv6: Box::new( + ipv6.subnets_with_prefix(128) + .map(|ip| ip.network_address()) + .filter(move |ip| !exclusion_v6.is_some_and(|e| e.contains(*ip))), + ), + } + } + + pub fn get_proxy_ip_for(&mut self, ip: &IpAddr) -> Option { + let proxy_ip = match ip { + IpAddr::V4(_) => self.ipv4.next().map(Into::into), + IpAddr::V6(_) => self.ipv6.next().map(Into::into), + }; + + if proxy_ip.is_none() { + // TODO: we might want to make the iterator cyclic or another strategy to prevent ip exhaustion + // this might happen in ipv4 if tokens are too long lived. + tracing::error!("IP exhaustion: Please reset your client"); + } + + proxy_ip + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/rust/connlib/tunnel/src/peer.rs b/rust/connlib/tunnel/src/peer.rs index 3fe38d006c..6f8b03239f 100644 --- a/rust/connlib/tunnel/src/peer.rs +++ b/rust/connlib/tunnel/src/peer.rs @@ -5,12 +5,12 @@ use std::time::Instant; use bimap::BiMap; use chrono::{DateTime, Utc}; use connlib_shared::messages::{DnsServer, ResourceId}; -use connlib_shared::IpProvider; use connlib_shared::{Error, Result}; use ip_network::IpNetwork; use ip_network_table::IpNetworkTable; use pnet_packet::Packet; +use crate::client::IpProvider; use crate::ip_packet::MutableIpPacket; type ExpiryingResource = (ResourceId, Option>);