From a509ee9198548786e05b2c99d5be582597a19277 Mon Sep 17 00:00:00 2001 From: link2xt Date: Mon, 14 Oct 2024 14:07:50 +0000 Subject: [PATCH] feat: prioritize cached results if DNS resolver returns many results This ensures we do not get stuck trying DNS resolver results when we have a known to work IP address in the cache and DNS resolver returns garbage either because it is a captive portal or if it maliciously wants to get us stuck trying a long list of unresponsive IP addresses. This also limits the number of results we try to 10 overall. If there are more results, we will retry later with new resolution results. --- src/net/dns.rs | 178 ++++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 162 insertions(+), 16 deletions(-) diff --git a/src/net/dns.rs b/src/net/dns.rs index d77081877a..147074924b 100644 --- a/src/net/dns.rs +++ b/src/net/dns.rs @@ -597,9 +597,14 @@ pub(crate) async fn lookup_host_with_cache( load_cache: bool, ) -> Result> { let now = time(); - let mut resolved_addrs = match lookup_host_and_update_cache(context, hostname, port, now).await - { - Ok(res) => res, + let resolved_addrs = match lookup_host_and_update_cache(context, hostname, port, now).await { + Ok(res) => { + if alpn.is_empty() { + res + } else { + sort_by_connection_timestamp(context, res, alpn, hostname).await? + } + } Err(err) => { warn!( context, @@ -608,29 +613,43 @@ pub(crate) async fn lookup_host_with_cache( Vec::new() } }; - if !alpn.is_empty() { - resolved_addrs = - sort_by_connection_timestamp(context, resolved_addrs, alpn, hostname).await?; - } if load_cache { - for addr in lookup_cache(context, hostname, port, alpn, now).await? { - if !resolved_addrs.contains(&addr) { - resolved_addrs.push(addr); - } - } - + let mut cache = lookup_cache(context, hostname, port, alpn, now).await?; if let Some(ips) = DNS_PRELOAD.get(hostname) { for ip in ips { let addr = SocketAddr::new(*ip, port); - if !resolved_addrs.contains(&addr) { - resolved_addrs.push(addr); + if !cache.contains(&addr) { + cache.push(addr); } } } + + Ok(merge_with_cache(resolved_addrs, cache)) + } else { + Ok(resolved_addrs) } +} - Ok(resolved_addrs) +/// Merges results received from DNS with cached results. +/// +/// At most 10 results are returned. +fn merge_with_cache( + mut resolved_addrs: Vec, + cache: Vec, +) -> Vec { + let rest = resolved_addrs.split_off(std::cmp::min(resolved_addrs.len(), 2)); + + for addr in cache.into_iter().chain(rest.into_iter()) { + if !resolved_addrs.contains(&addr) { + resolved_addrs.push(addr); + if resolved_addrs.len() >= 10 { + break; + } + } + } + + resolved_addrs } #[cfg(test)] @@ -866,4 +885,131 @@ mod tests { ], ); } + + #[test] + fn test_merge_with_cache() { + let first_addr = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1)); + let second_addr = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 2)); + + // If there is no cache, just return resolved addresses. + { + let resolved_addrs = vec![ + SocketAddr::new(first_addr, 993), + SocketAddr::new(second_addr, 993), + ]; + let cache = vec![]; + assert_eq!( + merge_with_cache(resolved_addrs.clone(), cache), + resolved_addrs + ); + } + + // If cache contains address that is not in resolution results, + // it is inserted in the merged result. + { + let resolved_addrs = vec![SocketAddr::new(first_addr, 993)]; + let cache = vec![SocketAddr::new(second_addr, 993)]; + assert_eq!( + merge_with_cache(resolved_addrs, cache), + vec![ + SocketAddr::new(first_addr, 993), + SocketAddr::new(second_addr, 993), + ] + ); + } + + // If cache contains address that is already in resolution results, + // it is not duplicated. + { + let resolved_addrs = vec![ + SocketAddr::new(first_addr, 993), + SocketAddr::new(second_addr, 993), + ]; + let cache = vec![SocketAddr::new(second_addr, 993)]; + assert_eq!( + merge_with_cache(resolved_addrs, cache), + vec![ + SocketAddr::new(first_addr, 993), + SocketAddr::new(second_addr, 993), + ] + ); + } + + // If DNS resolvers returns a lot of results, + // we should try cached results before going through all + // the resolver results. + { + let resolved_addrs = vec![ + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 2)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 3)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 4)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 5)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 6)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 7)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 8)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 9)), 993), + ]; + let cache = vec![SocketAddr::new(second_addr, 993)]; + assert_eq!( + merge_with_cache(resolved_addrs, cache), + vec![ + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 2)), 993), + SocketAddr::new(second_addr, 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 3)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 4)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 5)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 6)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 7)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 8)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 9)), 993), + ] + ); + } + + // Even if cache already contains all the incorrect results + // that resolver returns, this should not result in them being sorted to the top. + // Cache has known to work result returned first, + // so we should try it after the second result. + { + let resolved_addrs = vec![ + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 2)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 3)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 4)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 5)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 6)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 7)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 8)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 9)), 993), + ]; + let cache = vec![ + SocketAddr::new(second_addr, 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 9)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 8)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 7)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 6)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 5)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 4)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 3)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 2)), 993), + ]; + assert_eq!( + merge_with_cache(resolved_addrs, cache), + vec![ + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 2)), 993), + SocketAddr::new(second_addr, 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 9)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 8)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 7)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 6)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 5)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 4)), 993), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 3)), 993), + ] + ); + } + } }