From 9c0c1c141c0fae36044ea8422ee03f186afbeabb Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Fri, 7 Jun 2024 03:15:48 +1000
Subject: [PATCH] refactor(connlib): don't strinify domain name early (#5264)

Turning a query's `name` into a `String` as late as possible avoids
reparsing it in the tests.
---
 rust/connlib/tunnel/src/dns.rs   | 16 ++++++++--------
 rust/connlib/tunnel/src/io.rs    |  2 +-
 rust/connlib/tunnel/src/tests.rs |  6 ++----
 3 files changed, 11 insertions(+), 13 deletions(-)

diff --git a/rust/connlib/tunnel/src/dns.rs b/rust/connlib/tunnel/src/dns.rs
index 58e78671f..6e06486a9 100644
--- a/rust/connlib/tunnel/src/dns.rs
+++ b/rust/connlib/tunnel/src/dns.rs
@@ -37,7 +37,7 @@ pub(crate) enum ResolveStrategy<T, U, V> {
 
 #[derive(Debug)]
 pub struct DnsQuery<'a> {
-    pub name: String,
+    pub name: DomainName,
     pub record_type: RecordType,
     // We could be much more efficient with this field,
     // we only need the header to create the response.
@@ -64,7 +64,7 @@ impl<'a> DnsQuery<'a> {
 }
 
 struct DnsQueryParams {
-    name: String,
+    name: DomainName,
     record_type: RecordType,
 }
 
@@ -79,7 +79,7 @@ impl DnsQueryParams {
 }
 
 impl<T, V> ResolveStrategy<T, DnsQueryParams, V> {
-    fn forward(name: String, record_type: Rtype) -> ResolveStrategy<T, DnsQueryParams, V> {
+    fn forward(name: DomainName, record_type: Rtype) -> ResolveStrategy<T, DnsQueryParams, V> {
         ResolveStrategy::ForwardQuery(DnsQueryParams {
             name,
             record_type: u16::from(record_type).into(),
@@ -378,7 +378,7 @@ fn resource_from_question<N: ToName>(
     match qtype {
         Rtype::A => {
             let Some(description) = get_description(&name, dns_resources) else {
-                return Some(ResolveStrategy::forward(name.to_string(), qtype));
+                return Some(ResolveStrategy::forward(name, qtype));
             };
 
             let description = DnsResource::from_description(&description, name);
@@ -399,7 +399,7 @@ fn resource_from_question<N: ToName>(
         }
         Rtype::AAAA => {
             let Some(description) = get_description(&name, dns_resources) else {
-                return Some(ResolveStrategy::forward(name.to_string(), qtype));
+                return Some(ResolveStrategy::forward(name, qtype));
             };
             let description = DnsResource::from_description(&description, name);
 
@@ -419,10 +419,10 @@ fn resource_from_question<N: ToName>(
         }
         Rtype::PTR => {
             let Some(ip) = reverse_dns_addr(&name.to_string()) else {
-                return Some(ResolveStrategy::forward(name.to_string(), qtype));
+                return Some(ResolveStrategy::forward(name, qtype));
             };
             let Some(resource) = dns_resources_internal_ips.get(&ip) else {
-                return Some(ResolveStrategy::forward(name.to_string(), qtype));
+                return Some(ResolveStrategy::forward(name, qtype));
             };
             Some(ResolveStrategy::LocalResponse(RecordData::Ptr(
                 domain::rdata::Ptr::new(resource.address.clone()),
@@ -433,7 +433,7 @@ fn resource_from_question<N: ToName>(
                 return None;
             };
 
-            Some(ResolveStrategy::forward(name.to_string(), qtype))
+            Some(ResolveStrategy::forward(name, qtype))
         }
     }
 }
diff --git a/rust/connlib/tunnel/src/io.rs b/rust/connlib/tunnel/src/io.rs
index 4e4e1a50f..f23e1ef5a 100644
--- a/rust/connlib/tunnel/src/io.rs
+++ b/rust/connlib/tunnel/src/io.rs
@@ -139,7 +139,7 @@ impl Io {
             .forwarded_dns_queries
             .try_push(
                 {
-                    let name = query.name.clone();
+                    let name = query.name.clone().to_string();
                     let record_type = query.record_type;
 
                     async move { resolver.lookup(&name, record_type).await }
diff --git a/rust/connlib/tunnel/src/tests.rs b/rust/connlib/tunnel/src/tests.rs
index 0ed983caa..a6c049263 100644
--- a/rust/connlib/tunnel/src/tests.rs
+++ b/rust/connlib/tunnel/src/tests.rs
@@ -1336,14 +1336,12 @@ impl TunnelTest {
     // - hickory error?
     // - TTL?
     fn on_forwarded_dns_query(&mut self, query: DnsQuery<'static>, ref_state: &ReferenceState) {
-        let name = query.name.parse::<DomainName>().unwrap(); // TODO: Could `DnsQuery` hold a `DomainName` directly?
-
         let resolved_ips = &ref_state
             .global_dns_records
-            .get(&name)
+            .get(&query.name)
             .expect("Deferred DNS query to be for known domain");
 
-        let name = domain_to_hickory_name(name);
+        let name = domain_to_hickory_name(query.name.clone());
         let record_type = query.record_type;
 
         let record_data = resolved_ips