Skip to content

Commit bd9a87b

Browse files
ZettaScriptZettaScript
committed
use boolean flag, websocket config, Litep2pConfig
1 parent d275270 commit bd9a87b

File tree

6 files changed

+59
-23
lines changed

6 files changed

+59
-23
lines changed

‎src/config.rs

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,9 @@ pub struct ConfigBuilder {
124124

125125
/// Close the connection if no substreams are open within this time frame.
126126
keep_alive_timeout: Duration,
127+
128+
/// Use system's DNS config.
129+
use_system_dns_config: bool,
127130
}
128131

129132
impl Default for ConfigBuilder {
@@ -157,6 +160,7 @@ impl ConfigBuilder {
157160
known_addresses: Vec::new(),
158161
connection_limits: ConnectionLimitsConfig::default(),
159162
keep_alive_timeout: KEEP_ALIVE_TIMEOUT,
163+
use_system_dns_config: false,
160164
}
161165
}
162166

@@ -278,6 +282,12 @@ impl ConfigBuilder {
278282
self
279283
}
280284

285+
/// Set DNS resolver according to system configuration instead of default (Google).
286+
pub fn with_system_resolver(mut self) -> Self {
287+
self.use_system_dns_config = true;
288+
self
289+
}
290+
281291
/// Build [`Litep2pConfig`].
282292
pub fn build(mut self) -> Litep2pConfig {
283293
let keypair = match self.keypair {
@@ -307,6 +317,7 @@ impl ConfigBuilder {
307317
known_addresses: self.known_addresses,
308318
connection_limits: self.connection_limits,
309319
keep_alive_timeout: self.keep_alive_timeout,
320+
use_system_dns_config: self.use_system_dns_config,
310321
}
311322
}
312323
}
@@ -369,4 +380,7 @@ pub struct Litep2pConfig {
369380

370381
/// Close the connection if no substreams are open within this time frame.
371382
pub(crate) keep_alive_timeout: Duration,
383+
384+
/// Use system's DNS config.
385+
pub(crate) use_system_dns_config: bool,
372386
}

‎src/transport/tcp/config.rs

Lines changed: 3 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -81,11 +81,8 @@ pub struct Config {
8181
/// the substream rejected.
8282
pub substream_open_timeout: std::time::Duration,
8383

84-
/// DNS resolver config.
85-
pub resolver_config: hickory_resolver::config::ResolverConfig,
86-
87-
/// DNS resolver options.
88-
pub resolver_opts: hickory_resolver::config::ResolverOpts,
84+
/// Use system's DNS config.
85+
pub use_system_dns_config: bool,
8986
}
9087

9188
impl Default for Config {
@@ -102,18 +99,7 @@ impl Default for Config {
10299
noise_write_buffer_size: MAX_WRITE_BUFFER_SIZE,
103100
connection_open_timeout: CONNECTION_OPEN_TIMEOUT,
104101
substream_open_timeout: SUBSTREAM_OPEN_TIMEOUT,
105-
resolver_config: Default::default(),
106-
resolver_opts: Default::default(),
102+
use_system_dns_config: false,
107103
}
108104
}
109105
}
110-
111-
impl Config {
112-
/// Set DNS resolver according to system configuration.
113-
pub fn set_system_resolver(&mut self) -> hickory_resolver::error::ResolveResult<()> {
114-
let (resolver_config, resolver_opts) = hickory_resolver::system_conf::read_system_conf()?;
115-
self.resolver_config = resolver_config;
116-
self.resolver_opts = resolver_opts;
117-
Ok(())
118-
}
119-
}

‎src/transport/tcp/mod.rs

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -294,13 +294,16 @@ impl TransportBuilder for TcpTransport {
294294
config.nodelay,
295295
);
296296

297+
let (resolver_config, resolver_opts) = if config.use_system_dns_config {
298+
hickory_resolver::system_conf::read_system_conf()
299+
.expect("TODO failed to read system DNS config")
300+
} else {
301+
(Default::default(), Default::default())
302+
};
303+
297304
Ok((
298305
Self {
299306
listener,
300-
resolver: Arc::new(TokioAsyncResolver::tokio(
301-
config.resolver_config.clone(),
302-
config.resolver_opts.clone(),
303-
)),
304307
config,
305308
context,
306309
dial_addresses,
@@ -311,6 +314,7 @@ impl TransportBuilder for TcpTransport {
311314
pending_connections: FuturesStream::new(),
312315
pending_raw_connections: FuturesStream::new(),
313316
cancel_futures: HashMap::new(),
317+
resolver: Arc::new(TokioAsyncResolver::tokio(resolver_config, resolver_opts)),
314318
},
315319
listen_addresses,
316320
))

‎src/transport/websocket/config.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,9 @@ pub struct Config {
8080
/// How long should litep2p wait for a substream to be opened before considering
8181
/// the substream rejected.
8282
pub substream_open_timeout: std::time::Duration,
83+
84+
/// Use system's DNS config.
85+
pub use_system_dns_config: bool,
8386
}
8487

8588
impl Default for Config {
@@ -96,6 +99,7 @@ impl Default for Config {
9699
noise_write_buffer_size: MAX_WRITE_BUFFER_SIZE,
97100
connection_open_timeout: CONNECTION_OPEN_TIMEOUT,
98101
substream_open_timeout: SUBSTREAM_OPEN_TIMEOUT,
102+
use_system_dns_config: false,
99103
}
100104
}
101105
}

‎src/transport/websocket/connection.rs

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -605,6 +605,10 @@ mod tests {
605605

606606
use super::*;
607607
use futures::AsyncWriteExt;
608+
use hickory_resolver::{
609+
config::{ResolverConfig, ResolverOpts},
610+
TokioAsyncResolver,
611+
};
608612
use tokio::net::TcpListener;
609613

610614
#[tokio::test]
@@ -638,6 +642,7 @@ mod tests {
638642
Default::default(),
639643
Duration::from_secs(10),
640644
false,
645+
TokioAsyncResolver::tokio(ResolverConfig::default(), ResolverOpts::default()),
641646
)
642647
.await
643648
.unwrap();
@@ -753,6 +758,7 @@ mod tests {
753758
Default::default(),
754759
Duration::from_secs(10),
755760
false,
761+
TokioAsyncResolver::tokio(ResolverConfig::default(), ResolverOpts::default()),
756762
)
757763
.await
758764
.unwrap();
@@ -1023,6 +1029,7 @@ mod tests {
10231029
Default::default(),
10241030
Duration::from_secs(10),
10251031
false,
1032+
TokioAsyncResolver::tokio(ResolverConfig::default(), ResolverOpts::default()),
10261033
)
10271034
.await
10281035
.unwrap();
@@ -1188,6 +1195,7 @@ mod tests {
11881195
Default::default(),
11891196
Duration::from_secs(10),
11901197
false,
1198+
TokioAsyncResolver::tokio(ResolverConfig::default(), ResolverOpts::default()),
11911199
)
11921200
.await
11931201
.unwrap();
@@ -1339,6 +1347,7 @@ mod tests {
13391347
Default::default(),
13401348
Duration::from_secs(10),
13411349
false,
1350+
TokioAsyncResolver::tokio(ResolverConfig::default(), ResolverOpts::default()),
13421351
)
13431352
.await
13441353
.unwrap();

‎src/transport/websocket/mod.rs

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,9 +42,10 @@ use futures::{
4242
stream::{AbortHandle, FuturesUnordered},
4343
Stream, StreamExt, TryFutureExt,
4444
};
45+
use hickory_resolver::TokioAsyncResolver;
4546
use multiaddr::{Multiaddr, Protocol};
4647
use socket2::{Domain, Socket, Type};
47-
use std::net::SocketAddr;
48+
use std::{borrow::Borrow, net::SocketAddr, sync::Arc};
4849
use tokio::net::TcpStream;
4950
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
5051

@@ -132,6 +133,9 @@ pub(crate) struct WebSocketTransport {
132133

133134
/// Negotiated connections waiting validation.
134135
pending_open: HashMap<ConnectionId, NegotiatedConnection>,
136+
137+
/// DNS resolver.
138+
resolver: Arc<TokioAsyncResolver>,
135139
}
136140

137141
impl WebSocketTransport {
@@ -224,12 +228,15 @@ impl WebSocketTransport {
224228
dial_addresses: DialAddresses,
225229
connection_open_timeout: Duration,
226230
nodelay: bool,
231+
resolver: impl Borrow<TokioAsyncResolver>,
227232
) -> Result<(Multiaddr, WebSocketStream<MaybeTlsStream<TcpStream>>), DialError> {
228233
let (url, _) = Self::multiaddr_into_url(address.clone())?;
229234

230235
let (socket_address, _) = WebSocketAddress::multiaddr_to_socket_address(&address)?;
231236
let remote_address =
232-
match tokio::time::timeout(connection_open_timeout, socket_address.lookup_ip()).await {
237+
match tokio::time::timeout(connection_open_timeout, socket_address.lookup_ip(resolver))
238+
.await
239+
{
233240
Err(_) => return Err(DialError::Timeout),
234241
Ok(Err(error)) => return Err(error.into()),
235242
Ok(Ok(address)) => address,
@@ -317,6 +324,13 @@ impl TransportBuilder for WebSocketTransport {
317324
config.nodelay,
318325
);
319326

327+
let (resolver_config, resolver_opts) = if config.use_system_dns_config {
328+
hickory_resolver::system_conf::read_system_conf()
329+
.expect("TODO failed to read system DNS config")
330+
} else {
331+
(Default::default(), Default::default())
332+
};
333+
320334
Ok((
321335
Self {
322336
listener,
@@ -330,6 +344,7 @@ impl TransportBuilder for WebSocketTransport {
330344
pending_connections: FuturesStream::new(),
331345
pending_raw_connections: FuturesStream::new(),
332346
cancel_futures: HashMap::new(),
347+
resolver: Arc::new(TokioAsyncResolver::tokio(resolver_config, resolver_opts)),
333348
},
334349
listen_addresses,
335350
))
@@ -347,6 +362,7 @@ impl Transport for WebSocketTransport {
347362
let substream_open_timeout = self.config.substream_open_timeout;
348363
let dial_addresses = self.dial_addresses.clone();
349364
let nodelay = self.config.nodelay;
365+
let resolver = self.resolver.clone();
350366

351367
self.pending_dials.insert(connection_id, address.clone());
352368

@@ -358,6 +374,7 @@ impl Transport for WebSocketTransport {
358374
dial_addresses,
359375
connection_open_timeout,
360376
nodelay,
377+
resolver,
361378
)
362379
.await
363380
.map_err(|error| (connection_id, error))?;
@@ -475,13 +492,15 @@ impl Transport for WebSocketTransport {
475492
let connection_open_timeout = self.config.connection_open_timeout;
476493
let dial_addresses = self.dial_addresses.clone();
477494
let nodelay = self.config.nodelay;
495+
let resolver = self.resolver.clone();
478496

479497
async move {
480498
WebSocketTransport::dial_peer(
481499
address.clone(),
482500
dial_addresses,
483501
connection_open_timeout,
484502
nodelay,
503+
resolver,
485504
)
486505
.await
487506
.map_err(|error| (address, error))

0 commit comments

Comments
 (0)