feat(dns): Add timeout for tcp handshake for Dns over HTTPS/TLS

This commit is contained in:
Σrebe - Romain GERARD 2024-06-25 20:27:45 +02:00
parent cd87698b02
commit ef1ca16e4a
No known key found for this signature in database
GPG key ID: 7A42B4B97E0332F4

View file

@ -1,5 +1,6 @@
use crate::tcp;
use anyhow::{anyhow, Context}; use anyhow::{anyhow, Context};
use futures_util::FutureExt; use futures_util::{FutureExt, TryFutureExt};
use hickory_resolver::config::{NameServerConfig, Protocol, ResolverConfig, ResolverOpts}; use hickory_resolver::config::{NameServerConfig, Protocol, ResolverConfig, ResolverOpts};
use hickory_resolver::name_server::{GenericConnector, RuntimeProvider, TokioRuntimeProvider}; use hickory_resolver::name_server::{GenericConnector, RuntimeProvider, TokioRuntimeProvider};
use hickory_resolver::proto::iocompat::AsyncIoTokioAsStd; use hickory_resolver::proto::iocompat::AsyncIoTokioAsStd;
@ -9,6 +10,7 @@ use log::warn;
use std::future::Future; use std::future::Future;
use std::net::{IpAddr, SocketAddr, SocketAddrV4, SocketAddrV6}; use std::net::{IpAddr, SocketAddr, SocketAddrV4, SocketAddrV6};
use std::pin::Pin; use std::pin::Pin;
use std::time::Duration;
use tokio::net::{TcpStream, UdpSocket}; use tokio::net::{TcpStream, UdpSocket};
use url::{Host, Url}; use url::{Host, Url};
@ -147,24 +149,30 @@ impl RuntimeProvider for TokioRuntimeProviderWithSoMark {
#[inline] #[inline]
fn connect_tcp(&self, server_addr: SocketAddr) -> Pin<Box<dyn Send + Future<Output = std::io::Result<Self::Tcp>>>> { fn connect_tcp(&self, server_addr: SocketAddr) -> Pin<Box<dyn Send + Future<Output = std::io::Result<Self::Tcp>>>> {
let socket = TcpStream::connect(server_addr); #[cfg(not(target_os = "linux"))]
let so_mark = None;
#[cfg(target_os = "linux")] #[cfg(target_os = "linux")]
let socket = { let so_mark = self.so_mark;
use socket2::SockRef; let socket = async move {
let host = match server_addr.ip() {
IpAddr::V4(addr) => Host::<String>::Ipv4(addr),
IpAddr::V6(addr) => Host::<String>::Ipv6(addr),
};
socket.map({ tcp::connect(
let so_mark = self.so_mark; &host,
move |sock| { server_addr.port(),
if let (Ok(sock), Some(so_mark)) = (&sock, so_mark) { so_mark,
SockRef::from(sock).set_mark(so_mark)?; Duration::from_secs(10),
} &DnsResolver::System, // not going to be used as host is directly an ip address
sock )
} .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))
}) .map(|s| s.map(AsyncIoTokioAsStd))
.await
}; };
Box::pin(socket.map(|s| s.map(AsyncIoTokioAsStd))) Box::pin(socket)
} }
fn bind_udp( fn bind_udp(