From af7f2f2ce7e9d577401bed1dfc8380baf5066ea8 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Thu, 17 Oct 2024 10:29:33 +1100
Subject: [PATCH] fix(connlib): index mangled queries by ID and upstream
 (#7069)

When forwarding UDP DNS queries through the tunnel, `connlib` needs to
mangle the IP header to set upstream as the correct destination of the
packet. In order to know, which ones need to be mangled back, we keep a
map of query ID to timestamp. This isn't good enough as the added
proptest failure case shows. If there are two concurrent queries with
the same query ID to different resolvers, we only handle one of those
and don't mangle the 2nd one.
---
 rust/connlib/tunnel/proptest-regressions/tests.txt |  1 +
 rust/connlib/tunnel/src/client.rs                  | 12 ++++++------
 rust/connlib/tunnel/src/tests/sim_client.rs        |  2 +-
 3 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/rust/connlib/tunnel/proptest-regressions/tests.txt b/rust/connlib/tunnel/proptest-regressions/tests.txt
index 3e1fd2cae..d849ffdb7 100644
--- a/rust/connlib/tunnel/proptest-regressions/tests.txt
+++ b/rust/connlib/tunnel/proptest-regressions/tests.txt
@@ -115,3 +115,4 @@ cc eb667925801c37305d723dc1460cf0c94b4ad6be0f5e0392216dad33a43e892f
 cc a7038222559dc89bf16f94a6293e59a57a7585289b0490e57710def11b8901a2
 cc c431bc91c5f7ab88885d67da8785e69fc23baf290dfdb433cce1050f65911145
 cc b4dd2a98e4e6aa29f875fa7b8e3af451b1ce6ef8b4e9d6c4cd29fcb68e9249de
+cc 0a717e57a998e97be9134007c6a102c5ebaba5c477c95003eaa8f3c4503f88f1
diff --git a/rust/connlib/tunnel/src/client.rs b/rust/connlib/tunnel/src/client.rs
index 9f6a5bd65..f66b7557d 100644
--- a/rust/connlib/tunnel/src/client.rs
+++ b/rust/connlib/tunnel/src/client.rs
@@ -112,7 +112,7 @@ pub struct ClientState {
     /// DNS queries that had their destination IP mangled because the servers is a CIDR resource.
     ///
     /// The [`Instant`] tracks when the DNS query expires.
-    mangled_dns_queries: HashMap<u16, Instant>,
+    mangled_dns_queries: HashMap<(SocketAddr, u16), Instant>,
     /// Manages internal dns records and emits forwarding event when not internally handled
     stub_resolver: StubResolver,
 
@@ -621,14 +621,13 @@ impl ClientState {
                     tracing::trace!(server = %upstream, %query_id, "Forwarding UDP DNS query via tunnel");
 
                     self.mangled_dns_queries
-                        .insert(message.header().id(), now + IDS_EXPIRE);
+                        .insert((upstream, message.header().id()), now + IDS_EXPIRE);
                     packet.set_dst(upstream.ip());
                     packet.update_checksum();
 
                     return ControlFlow::Continue(packet);
                 }
 
-                let query_id = message.header().id();
                 let source = SocketAddr::new(packet.source(), datagram.source_port());
 
                 tracing::trace!(server = %upstream, %query_id, "Forwarding UDP DNS query directly via host");
@@ -1359,7 +1358,7 @@ fn is_definitely_not_a_resource(ip: IpAddr) -> bool {
 fn maybe_mangle_dns_response_from_cidr_resource(
     mut packet: IpPacket,
     dns_mapping: &BiMap<IpAddr, DnsServer>,
-    mangeled_dns_queries: &mut HashMap<u16, Instant>,
+    mangeled_dns_queries: &mut HashMap<(SocketAddr, u16), Instant>,
     now: Instant,
 ) -> IpPacket {
     let src_ip = packet.source();
@@ -1369,8 +1368,9 @@ fn maybe_mangle_dns_response_from_cidr_resource(
     };
 
     let src_port = udp.source_port();
+    let src_socket = SocketAddr::new(src_ip, src_port);
 
-    let Some(sentinel) = dns_mapping.get_by_right(&DnsServer::from((src_ip, src_port))) else {
+    let Some(sentinel) = dns_mapping.get_by_right(&DnsServer::from(src_socket)) else {
         return packet;
     };
 
@@ -1379,7 +1379,7 @@ fn maybe_mangle_dns_response_from_cidr_resource(
     };
 
     let Some(query_sent_at) = mangeled_dns_queries
-        .remove(&message.header().id())
+        .remove(&(src_socket, message.header().id()))
         .map(|expires_at| expires_at - IDS_EXPIRE)
     else {
         return packet;
diff --git a/rust/connlib/tunnel/src/tests/sim_client.rs b/rust/connlib/tunnel/src/tests/sim_client.rs
index 592adfa65..30cc96cda 100644
--- a/rust/connlib/tunnel/src/tests/sim_client.rs
+++ b/rust/connlib/tunnel/src/tests/sim_client.rs
@@ -199,7 +199,7 @@ impl SimClient {
                 // Map back to upstream socket so we can assert on it correctly.
                 let sentinel = SocketAddr::from((packet.source(), udp.source_port()));
                 let Some(upstream) = self.upstream_dns_by_sentinel(&sentinel) else {
-                    tracing::error!(%sentinel, "Unknown DNS server");
+                    tracing::error!(%sentinel, mapping = ?self.dns_by_sentinel, "Unknown DNS server");
                     return;
                 };