From dabe493e9ea46d4a4a345cfb058ca1f3e6764824 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Thu, 20 Jun 2024 10:36:11 +1000
Subject: [PATCH] feat(connlib): short-circuit access request to DNS resources
 (#5438)

Currently, we always emit a connection intent whenever we see a DNS
query for a domain of one of our DNS resources. However, especially for
wildcard DNS resources, we are very likely already connected to the
corresponding gateway. In that case, sending a connection intent
triggers another handshake with the portal only to learn that - surprise
- we should reuse a connection that we already have to that gateway.

We can short-circuit this by checking if we are already connected to the
gateway for this resource and directly requested access for the domain
name in question. We reuse the same event here as we do for refreshing
DNS resources. At a later stage, we should rename this to something else
to make this clearer.

Co-authored-by: Gabi <gabrielalejandro7@gmail.com>
---
 rust/connlib/tunnel/src/client.rs    | 17 ++++++-
 rust/connlib/tunnel/src/lib.rs       |  1 +
 rust/connlib/tunnel/src/tests/sut.rs | 72 +++++++++++++++++++++++-----
 3 files changed, 75 insertions(+), 15 deletions(-)

diff --git a/rust/connlib/tunnel/src/client.rs b/rust/connlib/tunnel/src/client.rs
index fa1460734..8f288e574 100644
--- a/rust/connlib/tunnel/src/client.rs
+++ b/rust/connlib/tunnel/src/client.rs
@@ -801,9 +801,22 @@ impl ClientState {
                 Ok(None)
             }
             Some(dns::ResolveStrategy::DeferredResponse((resource, r_type))) => {
-                self.on_not_connected_resource(resource.id, Some(resource.address.clone()), now);
                 self.deferred_dns_queries
-                    .insert((resource, r_type), packet.as_immutable().to_owned());
+                    .insert((resource.clone(), r_type), packet.as_immutable().to_owned());
+
+                let Some(peer) = self.peer_by_resource(resource.id) else {
+                    self.on_not_connected_resource(resource.id, Some(resource.address), now);
+                    return Ok(None);
+                };
+
+                self.buffered_events
+                    .push_back(ClientEvent::RefreshResources {
+                        connections: vec![ReuseConnection {
+                            resource_id: resource.id,
+                            gateway_id: peer.id(),
+                            payload: Some(resource.address),
+                        }],
+                    });
 
                 Ok(None)
             }
diff --git a/rust/connlib/tunnel/src/lib.rs b/rust/connlib/tunnel/src/lib.rs
index aba11b07b..f2b4ffbfd 100644
--- a/rust/connlib/tunnel/src/lib.rs
+++ b/rust/connlib/tunnel/src/lib.rs
@@ -275,6 +275,7 @@ pub enum ClientEvent {
         resource: ResourceId,
         connected_gateway_ids: HashSet<GatewayId>,
     },
+    // FIXME: Refactor this to "request-access" or something similar.
     RefreshResources {
         connections: Vec<ReuseConnection>,
     },
diff --git a/rust/connlib/tunnel/src/tests/sut.rs b/rust/connlib/tunnel/src/tests/sut.rs
index 51d9624a3..9c209dd7f 100644
--- a/rust/connlib/tunnel/src/tests/sut.rs
+++ b/rust/connlib/tunnel/src/tests/sut.rs
@@ -703,39 +703,85 @@ impl TunnelTest {
                             .unwrap();
                     }
                     Request::ReuseConnection(reuse_connection) => {
-                        if let (Ok(()), Some(name)) = (
-                            self.gateway.span.in_scope(|| {
+                        let domain = reuse_connection.payload;
+
+                        self.gateway
+                            .span
+                            .in_scope(|| {
                                 self.gateway.state.allow_access(
                                     resource,
                                     self.client.id,
                                     None,
-                                    reuse_connection
-                                        .payload
-                                        .as_ref()
-                                        .map(|d| (d.clone(), Vec::new())),
+                                    domain.clone().map(|d| (d, Vec::new())),
                                     now,
                                 )
-                            }),
-                            reuse_connection.payload,
-                        ) {
+                            })
+                            .unwrap();
+
+                        if let Some(domain) = domain {
                             self.client
                                 .span
                                 .in_scope(|| {
                                     self.client.state.received_domain_parameters(
                                         resource_id,
                                         DomainResponse {
-                                            domain: name.clone(),
+                                            domain,
                                             address: resolved_ips,
                                         },
                                     )
                                 })
                                 .unwrap();
-                        };
+                        }
                     }
                 };
             }
-            ClientEvent::RefreshResources { .. } => {
-                tracing::warn!("Unimplemented");
+            ClientEvent::RefreshResources { connections } => {
+                for reuse_connection in connections {
+                    let domain = reuse_connection.payload.clone();
+                    let resource_id = reuse_connection.resource_id;
+
+                    let resolved_ips = reuse_connection
+                        .payload
+                        .into_iter()
+                        .flat_map(|domain| global_dns_records.get(&domain).cloned())
+                        .flatten()
+                        .collect::<Vec<_>>();
+
+                    let resource = map_client_resource_to_gateway_resource(
+                        client_cidr_resources,
+                        client_dns_resource,
+                        resolved_ips.clone(),
+                        resource_id,
+                    );
+
+                    self.gateway
+                        .span
+                        .in_scope(|| {
+                            self.gateway.state.allow_access(
+                                resource,
+                                self.client.id,
+                                None,
+                                domain.clone().map(|d| (d, Vec::new())),
+                                now,
+                            )
+                        })
+                        .unwrap();
+
+                    if let Some(domain) = domain {
+                        self.client
+                            .span
+                            .in_scope(|| {
+                                self.client.state.received_domain_parameters(
+                                    resource_id,
+                                    DomainResponse {
+                                        domain,
+                                        address: resolved_ips,
+                                    },
+                                )
+                            })
+                            .unwrap();
+                    }
+                }
             }
             ClientEvent::ResourcesChanged { .. } => {
                 tracing::warn!("Unimplemented");