From 9f850bb92d53f42eb85870264c7d3268dca046d4 Mon Sep 17 00:00:00 2001
From: Gabi <gabrielalejandro7@gmail.com>
Date: Mon, 18 Mar 2024 21:05:05 -0300
Subject: [PATCH] fix(connlib): exclude sentinel dns range for resources ips
 (#4200)

In the future we will want to refactor this to a builder pattern to
prevent the number of parameters from growing and have them clearer but
this works simply for now.

Found while discussing #4174

---------

Co-authored-by: Thomas Eizinger <thomas@eizinger.io>
---
 rust/connlib/shared/src/lib.rs    | 34 ---------------------
 rust/connlib/tunnel/src/client.rs | 49 +++++++++++++++++++++++++++++--
 rust/connlib/tunnel/src/peer.rs   |  2 +-
 3 files changed, 48 insertions(+), 37 deletions(-)

diff --git a/rust/connlib/shared/src/lib.rs b/rust/connlib/shared/src/lib.rs
index 653ce8691..37ec6fbf6 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<Vec<u8>>;
 
@@ -61,35 +56,6 @@ pub fn keypair() -> (StaticSecret, PublicKey) {
     (private_key, public_key)
 }
 
-pub struct IpProvider {
-    ipv4: Box<dyn Iterator<Item = Ipv4Addr> + Send + Sync>,
-    ipv6: Box<dyn Iterator<Item = Ipv6Addr> + 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<IpAddr> {
-        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>) -> 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 d21cfce83..e2a0ce5f8 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;
 
@@ -352,6 +352,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(),
@@ -936,6 +938,8 @@ fn sentinel_dns_mapping(dns: &[DnsServer]) -> BiMap<IpAddr, DnsServer> {
     let mut ip_provider = IpProvider::new(
         DNS_SENTINELS_V4.parse().unwrap(),
         DNS_SENTINELS_V6.parse().unwrap(),
+        None,
+        None,
     );
 
     dns.iter()
@@ -974,6 +978,47 @@ fn is_definitely_not_a_resource(ip: IpAddr) -> bool {
     false
 }
 
+pub struct IpProvider {
+    ipv4: Box<dyn Iterator<Item = Ipv4Addr> + Send + Sync>,
+    ipv6: Box<dyn Iterator<Item = Ipv6Addr> + Send + Sync>,
+}
+
+impl IpProvider {
+    pub fn new(
+        ipv4: Ipv4Network,
+        ipv6: Ipv6Network,
+        exclusion_v4: Option<Ipv4Network>,
+        exclusion_v6: Option<Ipv6Network>,
+    ) -> 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<IpAddr> {
+        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 3fe38d006..6f8b03239 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<DateTime<Utc>>);