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.
This commit is contained in:
Thomas Eizinger
2024-10-17 10:29:33 +11:00
committed by GitHub
parent b96244efd8
commit af7f2f2ce7
3 changed files with 8 additions and 7 deletions

View File

@@ -115,3 +115,4 @@ cc eb667925801c37305d723dc1460cf0c94b4ad6be0f5e0392216dad33a43e892f
cc a7038222559dc89bf16f94a6293e59a57a7585289b0490e57710def11b8901a2
cc c431bc91c5f7ab88885d67da8785e69fc23baf290dfdb433cce1050f65911145
cc b4dd2a98e4e6aa29f875fa7b8e3af451b1ce6ef8b4e9d6c4cd29fcb68e9249de
cc 0a717e57a998e97be9134007c6a102c5ebaba5c477c95003eaa8f3c4503f88f1

View File

@@ -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;

View File

@@ -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;
};