From 5c1f5e1ece0b0e778f30a870800d7169dc9b36cf Mon Sep 17 00:00:00 2001
From: Gabi <gabrielalejandro7@gmail.com>
Date: Wed, 17 Jul 2024 22:27:07 -0300
Subject: [PATCH] fix(connlib): prioritize dns resources to CIDR ones in case
 of an overlap (#5840)

For full route this happens always and if we don't prioritize DNS
resources any packet for DNS IPs will get routed to the full route
gateway which might not have the correct resource.

TODO: this still needs unit tests
TODO: Waiting on #5891
---
 rust/connlib/tunnel/src/client.rs           | 10 ++++++---
 rust/connlib/tunnel/src/tests/reference.rs  |  6 +-----
 rust/connlib/tunnel/src/tests/transition.rs | 24 +++++++++++++++++++++
 3 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/rust/connlib/tunnel/src/client.rs b/rust/connlib/tunnel/src/client.rs
index 02289c4d0..d9b496eca 100644
--- a/rust/connlib/tunnel/src/client.rs
+++ b/rust/connlib/tunnel/src/client.rs
@@ -696,6 +696,10 @@ impl ClientState {
                 vacant.insert(AwaitingConnectionDetails {
                     gateways: gateways.clone(),
                     last_intent_sent_at: now,
+                    // Note: in case of an overlapping CIDR resource this should be None instead of Some if the resource_id
+                    // is for a CIDR resource.
+                    // But this should never happen as DNS resources are always preferred, so we don't encode the logic here.
+                    // Tests will prevent this from ever happening.
                     domain: self.stub_resolver.get_fqdn(destination).map(|(fqdn, ips)| {
                         ResolveRequest {
                             name: fqdn.clone(),
@@ -745,14 +749,14 @@ impl ClientState {
     }
 
     fn get_resource_by_destination(&self, destination: IpAddr) -> Option<ResourceId> {
+        let maybe_dns_resource_id = self.stub_resolver.resolve_resource_by_ip(&destination);
+
         let maybe_cidr_resource_id = self
             .cidr_resources
             .longest_match(destination)
             .map(|(_, res)| res.id);
 
-        let maybe_dns_resource_id = self.stub_resolver.resolve_resource_by_ip(&destination);
-
-        maybe_cidr_resource_id.or(maybe_dns_resource_id)
+        maybe_dns_resource_id.or(maybe_cidr_resource_id)
     }
 
     #[must_use]
diff --git a/rust/connlib/tunnel/src/tests/reference.rs b/rust/connlib/tunnel/src/tests/reference.rs
index d1703a3e0..91ecc2e2d 100644
--- a/rust/connlib/tunnel/src/tests/reference.rs
+++ b/rust/connlib/tunnel/src/tests/reference.rs
@@ -145,11 +145,7 @@ impl ReferenceStateMachine for ReferenceState {
             )
             .with(
                 1,
-                cidr_resource(
-                    any_ip_network(8),
-                    sample::select(state.portal.all_sites()).prop_map(|s| vec![s]),
-                )
-                .prop_map(|resource| Transition::AddCidrResource { resource }),
+                add_cidr_resource(sample::select(state.portal.all_sites()).prop_map(|s| vec![s])),
             )
             .with(1, roam_client())
             .with(
diff --git a/rust/connlib/tunnel/src/tests/transition.rs b/rust/connlib/tunnel/src/tests/transition.rs
index e6ac13f76..2c068dc18 100644
--- a/rust/connlib/tunnel/src/tests/transition.rs
+++ b/rust/connlib/tunnel/src/tests/transition.rs
@@ -1,3 +1,5 @@
+use crate::client::{IPV4_RESOURCES, IPV6_RESOURCES};
+
 use super::{
     sim_net::{any_ip_stack, any_port},
     strategies::*,
@@ -11,10 +13,12 @@ use connlib_shared::{
     DomainName,
 };
 use hickory_proto::rr::RecordType;
+use ip_network::IpNetwork;
 use proptest::{prelude::*, sample};
 use std::{
     collections::{HashMap, HashSet},
     net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
+    str::FromStr,
 };
 
 /// The possible transitions of the state machine.
@@ -175,6 +179,26 @@ where
         )
 }
 
+// Adds a non-overlapping CIDR resource with the gateway
+pub(crate) fn add_cidr_resource(
+    sites: impl Strategy<Value = Vec<Site>>,
+) -> impl Strategy<Value = Transition> {
+    cidr_resource(any_ip_network(8), sites)
+        .prop_filter(
+            "tests doesn't support yet CIDR resources overlapping DNS resources",
+            |r| {
+                // This works because CIDR resourc's host mask is always <8 while IP resource is 21
+                !IpNetwork::from_str(IPV4_RESOURCES)
+                    .unwrap()
+                    .contains(r.address.network_address())
+                    && !IpNetwork::from_str(IPV6_RESOURCES)
+                        .unwrap()
+                        .contains(r.address.network_address())
+            },
+        )
+        .prop_map(|resource| Transition::AddCidrResource { resource })
+}
+
 pub(crate) fn non_wildcard_dns_resource(
     site: impl Strategy<Value = Site>,
 ) -> impl Strategy<Value = Transition> {