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, + 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, - mangeled_dns_queries: &mut HashMap, + 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; };