From 953ddeace6faa00e5ec6bfdb2122467ac5d36431 Mon Sep 17 00:00:00 2001
From: Gabi <gabrielalejandro7@gmail.com>
Date: Fri, 3 Nov 2023 02:16:03 -0300
Subject: [PATCH] connlib: update upstream dns format configuration (#2543)

fixes #2297
---
 rust/connlib/clients/shared/src/control.rs | 18 ++++++++++++++----
 rust/connlib/shared/src/messages.rs        | 15 +++++++++++++--
 2 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/rust/connlib/clients/shared/src/control.rs b/rust/connlib/clients/shared/src/control.rs
index 91b5830c2..4a188ca8c 100644
--- a/rust/connlib/clients/shared/src/control.rs
+++ b/rust/connlib/clients/shared/src/control.rs
@@ -1,5 +1,5 @@
 use async_compression::tokio::bufread::GzipEncoder;
-use std::net::{IpAddr, SocketAddr};
+use connlib_shared::messages::{DnsServer, IpDnsServer};
 use std::path::PathBuf;
 use std::{io, sync::Arc};
 
@@ -37,7 +37,7 @@ pub struct ControlPlane<CB: Callbacks> {
 }
 
 fn create_resolver(
-    upstream_dns: Vec<IpAddr>,
+    upstream_dns: Vec<DnsServer>,
     callbacks: &impl Callbacks,
 ) -> Option<TokioAsyncResolver> {
     let dns_servers = if upstream_dns.is_empty() {
@@ -48,13 +48,23 @@ fn create_resolver(
             return None;
         }
         dns_servers
+            .into_iter()
+            .map(|ip| {
+                DnsServer::IpPort(IpDnsServer {
+                    address: (ip, DNS_PORT).into(),
+                })
+            })
+            .collect()
     } else {
         upstream_dns
     };
 
     let mut resolver_config = ResolverConfig::new();
-    for ip in dns_servers.iter() {
-        let name_server = NameServerConfig::new(SocketAddr::new(*ip, DNS_PORT), Protocol::Udp);
+    for srv in dns_servers.iter() {
+        let name_server = match srv {
+            DnsServer::IpPort(srv) => NameServerConfig::new(srv.address, Protocol::Udp),
+        };
+
         resolver_config.add_name_server(name_server);
     }
 
diff --git a/rust/connlib/shared/src/messages.rs b/rust/connlib/shared/src/messages.rs
index 65ede6c42..c3ca5799d 100644
--- a/rust/connlib/shared/src/messages.rs
+++ b/rust/connlib/shared/src/messages.rs
@@ -1,5 +1,5 @@
 //! Message types that are used by both the gateway and client.
-use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
+use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
 
 use chrono::{serde::ts_seconds, DateTime, Utc};
 use ip_network::IpNetwork;
@@ -203,6 +203,17 @@ pub struct ResourceDescriptionCidr {
     pub name: String,
 }
 
+#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq)]
+#[serde(tag = "protocol", rename_all = "snake_case")]
+pub enum DnsServer {
+    IpPort(IpDnsServer),
+}
+
+#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq)]
+pub struct IpDnsServer {
+    pub address: SocketAddr,
+}
+
 /// Represents a wireguard interface configuration.
 ///
 /// Note that the ips are /32 for ipv4 and /128 for ipv6.
@@ -216,7 +227,7 @@ pub struct Interface {
     /// DNS that will be used to query for DNS that aren't within our resource list.
     #[serde(skip_serializing_if = "Vec::is_empty")]
     #[serde(default)]
-    pub upstream_dns: Vec<IpAddr>,
+    pub upstream_dns: Vec<DnsServer>,
 }
 
 /// A single relay