Skip to content

Commit cbd458b

Browse files
author
Stjepan Glavina
committed
Cleanup ToSocketAddrs, add more net reexports
1 parent e75c3a9 commit cbd458b

File tree

2 files changed

+151
-45
lines changed

2 files changed

+151
-45
lines changed

‎src/net/addr.rs‎

Lines changed: 146 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use std::marker::PhantomData;
1+
use std::mem;
22
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
33
use std::net::{SocketAddr, SocketAddrV4, SocketAddrV6};
44
use std::pin::Pin;
@@ -12,24 +12,38 @@ use crate::task::{blocking, Context, JoinHandle, Poll};
1212
cfg_if! {
1313
if #[cfg(feature = "docs")] {
1414
#[doc(hidden)]
15-
pub struct ImplFuture<'a, T>(std::marker::PhantomData<&'a T>);
15+
pub struct ImplFuture<T>(std::marker::PhantomData<T>);
1616

1717
macro_rules! ret {
18-
($a:lifetime, $f:tt, $i:ty) => (ImplFuture<$a, io::Result<$i>>);
18+
(impl Future<Output = $out:ty>, $fut:ty) => (ImplFuture<$out>);
1919
}
2020
} else {
2121
macro_rules! ret {
22-
($a:lifetime, $f:tt, $i:ty) => ($f<$a, $i>);
22+
(impl Future<Output = $out:ty>, $fut:ty) => ($fut);
2323
}
2424
}
2525
}
2626

27-
/// A trait for objects which can be converted or resolved to one or more [`SocketAddr`] values.
27+
/// Converts or resolves addresses to [`SocketAddr`] values.
2828
///
2929
/// This trait is an async version of [`std::net::ToSocketAddrs`].
3030
///
3131
/// [`std::net::ToSocketAddrs`]: https://doc.rust-lang.org/std/net/trait.ToSocketAddrs.html
32-
/// [`SocketAddr`]: https://doc.rust-lang.org/std/net/enum.SocketAddr.html
32+
/// [`SocketAddr`]: enum.SocketAddr.html
33+
///
34+
/// # Examples
35+
///
36+
/// ```
37+
/// # fn main() -> std::io::Result<()> { async_std::task::block_on(async {
38+
/// #
39+
/// use async_std::net::ToSocketAddrs;
40+
///
41+
/// let addr1 = "localhost:8080".to_socket_addrs().await?.next().unwrap();
42+
/// let addr2 = ("127.0.0.1", 8080).to_socket_addrs().await?.next().unwrap();
43+
/// assert_eq!(addr1, addr2);
44+
/// #
45+
/// # Ok(()) }) }
46+
/// ```
3347
pub trait ToSocketAddrs {
3448
/// Returned iterator over socket addresses which this type may correspond to.
3549
type Iter: Iterator<Item = SocketAddr>;
@@ -40,124 +54,211 @@ pub trait ToSocketAddrs {
4054
/// resolution performed.
4155
///
4256
/// Note that this function may block a backend thread while resolution is performed.
43-
fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter);
57+
fn to_socket_addrs(
58+
&self,
59+
) -> ret!(
60+
impl Future<Output = Self::Iter>,
61+
ToSocketAddrsFuture<Self::Iter>
62+
);
4463
}
4564

4665
#[doc(hidden)]
4766
#[allow(missing_debug_implementations)]
48-
pub enum ToSocketAddrsFuture<'a, I> {
49-
Phantom(PhantomData<&'a ()>),
50-
Join(JoinHandle<io::Result<I>>),
51-
Ready(Option<io::Result<I>>),
67+
pub enum ToSocketAddrsFuture<I> {
68+
Resolving(JoinHandle<io::Result<I>>),
69+
Ready(io::Result<I>),
70+
Done,
5271
}
5372

54-
impl<I: Iterator<Item = SocketAddr>> Future for ToSocketAddrsFuture<'_, I> {
73+
impl<I: Iterator<Item = SocketAddr>> Future for ToSocketAddrsFuture<I> {
5574
type Output = io::Result<I>;
5675

5776
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
58-
match unsafe { self.get_unchecked_mut() } {
59-
ToSocketAddrsFuture::Join(f) => Pin::new(&mut *f).poll(cx),
60-
ToSocketAddrsFuture::Ready(res) => {
61-
let res = res.take().expect("polled a completed future");
62-
Poll::Ready(res)
77+
let this = unsafe { self.get_unchecked_mut() };
78+
let state = mem::replace(this, ToSocketAddrsFuture::Done);
79+
80+
match state {
81+
ToSocketAddrsFuture::Resolving(mut task) => {
82+
let poll = Pin::new(&mut task).poll(cx);
83+
if poll.is_pending() {
84+
*this = ToSocketAddrsFuture::Resolving(task);
85+
}
86+
poll
6387
}
64-
_ => unreachable!(),
88+
ToSocketAddrsFuture::Ready(res) => Poll::Ready(res),
89+
ToSocketAddrsFuture::Done => panic!("polled a completed future"),
6590
}
6691
}
6792
}
6893

6994
impl ToSocketAddrs for SocketAddr {
7095
type Iter = std::option::IntoIter<SocketAddr>;
7196

72-
fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter) {
73-
ToSocketAddrsFuture::Ready(Some(std::net::ToSocketAddrs::to_socket_addrs(self)))
97+
fn to_socket_addrs(
98+
&self,
99+
) -> ret!(
100+
impl Future<Output = Self::Iter>,
101+
ToSocketAddrsFuture<Self::Iter>
102+
) {
103+
ToSocketAddrsFuture::Ready(Ok(Some(*self).into_iter()))
74104
}
75105
}
76106

77107
impl ToSocketAddrs for SocketAddrV4 {
78108
type Iter = std::option::IntoIter<SocketAddr>;
79109

80-
fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter) {
81-
ToSocketAddrsFuture::Ready(Some(std::net::ToSocketAddrs::to_socket_addrs(self)))
110+
fn to_socket_addrs(
111+
&self,
112+
) -> ret!(
113+
impl Future<Output = Self::Iter>,
114+
ToSocketAddrsFuture<Self::Iter>
115+
) {
116+
SocketAddr::V4(*self).to_socket_addrs()
82117
}
83118
}
84119

85120
impl ToSocketAddrs for SocketAddrV6 {
86121
type Iter = std::option::IntoIter<SocketAddr>;
87122

88-
fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter) {
89-
ToSocketAddrsFuture::Ready(Some(std::net::ToSocketAddrs::to_socket_addrs(self)))
123+
fn to_socket_addrs(
124+
&self,
125+
) -> ret!(
126+
impl Future<Output = Self::Iter>,
127+
ToSocketAddrsFuture<Self::Iter>
128+
) {
129+
SocketAddr::V6(*self).to_socket_addrs()
90130
}
91131
}
92132

93133
impl ToSocketAddrs for (IpAddr, u16) {
94134
type Iter = std::option::IntoIter<SocketAddr>;
95135

96-
fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter) {
97-
ToSocketAddrsFuture::Ready(Some(std::net::ToSocketAddrs::to_socket_addrs(self)))
136+
fn to_socket_addrs(
137+
&self,
138+
) -> ret!(
139+
impl Future<Output = Self::Iter>,
140+
ToSocketAddrsFuture<Self::Iter>
141+
) {
142+
let (ip, port) = *self;
143+
match ip {
144+
IpAddr::V4(a) => (a, port).to_socket_addrs(),
145+
IpAddr::V6(a) => (a, port).to_socket_addrs(),
146+
}
98147
}
99148
}
100149

101150
impl ToSocketAddrs for (Ipv4Addr, u16) {
102151
type Iter = std::option::IntoIter<SocketAddr>;
103152

104-
fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter) {
105-
ToSocketAddrsFuture::Ready(Some(std::net::ToSocketAddrs::to_socket_addrs(self)))
153+
fn to_socket_addrs(
154+
&self,
155+
) -> ret!(
156+
impl Future<Output = Self::Iter>,
157+
ToSocketAddrsFuture<Self::Iter>
158+
) {
159+
let (ip, port) = *self;
160+
SocketAddrV4::new(ip, port).to_socket_addrs()
106161
}
107162
}
108163

109164
impl ToSocketAddrs for (Ipv6Addr, u16) {
110165
type Iter = std::option::IntoIter<SocketAddr>;
111166

112-
fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter) {
113-
ToSocketAddrsFuture::Ready(Some(std::net::ToSocketAddrs::to_socket_addrs(self)))
167+
fn to_socket_addrs(
168+
&self,
169+
) -> ret!(
170+
impl Future<Output = Self::Iter>,
171+
ToSocketAddrsFuture<Self::Iter>
172+
) {
173+
let (ip, port) = *self;
174+
SocketAddrV6::new(ip, port, 0, 0).to_socket_addrs()
114175
}
115176
}
116177

117178
impl ToSocketAddrs for (&str, u16) {
118179
type Iter = std::vec::IntoIter<SocketAddr>;
119180

120-
fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter) {
121-
let host = self.0.to_string();
122-
let port = self.1;
123-
let join = blocking::spawn(async move {
181+
fn to_socket_addrs(
182+
&self,
183+
) -> ret!(
184+
impl Future<Output = Self::Iter>,
185+
ToSocketAddrsFuture<Self::Iter>
186+
) {
187+
let (host, port) = *self;
188+
189+
if let Ok(addr) = host.parse::<Ipv4Addr>() {
190+
let addr = SocketAddrV4::new(addr, port);
191+
return ToSocketAddrsFuture::Ready(Ok(vec![SocketAddr::V4(addr)].into_iter()));
192+
}
193+
194+
if let Ok(addr) = host.parse::<Ipv6Addr>() {
195+
let addr = SocketAddrV6::new(addr, port, 0, 0);
196+
return ToSocketAddrsFuture::Ready(Ok(vec![SocketAddr::V6(addr)].into_iter()));
197+
}
198+
199+
let host = host.to_string();
200+
let task = blocking::spawn(async move {
124201
std::net::ToSocketAddrs::to_socket_addrs(&(host.as_str(), port))
125202
});
126-
ToSocketAddrsFuture::Join(join)
203+
ToSocketAddrsFuture::Resolving(task)
127204
}
128205
}
129206

130207
impl ToSocketAddrs for str {
131208
type Iter = std::vec::IntoIter<SocketAddr>;
132209

133-
fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter) {
134-
let socket_addrs = self.to_string();
135-
let join =
136-
blocking::spawn(async move { std::net::ToSocketAddrs::to_socket_addrs(&socket_addrs) });
137-
ToSocketAddrsFuture::Join(join)
210+
fn to_socket_addrs(
211+
&self,
212+
) -> ret!(
213+
impl Future<Output = Self::Iter>,
214+
ToSocketAddrsFuture<Self::Iter>
215+
) {
216+
if let Some(addr) = self.parse().ok() {
217+
return ToSocketAddrsFuture::Ready(Ok(vec![addr].into_iter()));
218+
}
219+
220+
let addr = self.to_string();
221+
let task =
222+
blocking::spawn(async move { std::net::ToSocketAddrs::to_socket_addrs(addr.as_str()) });
223+
ToSocketAddrsFuture::Resolving(task)
138224
}
139225
}
140226

141227
impl<'a> ToSocketAddrs for &'a [SocketAddr] {
142228
type Iter = std::iter::Cloned<std::slice::Iter<'a, SocketAddr>>;
143229

144-
fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter) {
145-
ToSocketAddrsFuture::Ready(Some(std::net::ToSocketAddrs::to_socket_addrs(self)))
230+
fn to_socket_addrs(
231+
&self,
232+
) -> ret!(
233+
impl Future<Output = Self::Iter>,
234+
ToSocketAddrsFuture<Self::Iter>
235+
) {
236+
ToSocketAddrsFuture::Ready(Ok(self.iter().cloned()))
146237
}
147238
}
148239

149240
impl<T: ToSocketAddrs + ?Sized> ToSocketAddrs for &T {
150241
type Iter = T::Iter;
151242

152-
fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter) {
243+
fn to_socket_addrs(
244+
&self,
245+
) -> ret!(
246+
impl Future<Output = Self::Iter>,
247+
ToSocketAddrsFuture<Self::Iter>
248+
) {
153249
(**self).to_socket_addrs()
154250
}
155251
}
156252

157253
impl ToSocketAddrs for String {
158254
type Iter = std::vec::IntoIter<SocketAddr>;
159255

160-
fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter) {
161-
ToSocketAddrs::to_socket_addrs(self.as_str())
256+
fn to_socket_addrs(
257+
&self,
258+
) -> ret!(
259+
impl Future<Output = Self::Iter>,
260+
ToSocketAddrsFuture<Self::Iter>
261+
) {
262+
(&**self).to_socket_addrs()
162263
}
163264
}

‎src/net/mod.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,11 @@
2828
//! # }) }
2929
//! ```
3030
31+
pub use std::net::AddrParseError;
32+
pub use std::net::Shutdown;
33+
pub use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
34+
pub use std::net::{SocketAddr, SocketAddrV4, SocketAddrV6};
35+
3136
pub use addr::ToSocketAddrs;
3237
pub use tcp::{Incoming, TcpListener, TcpStream};
3338
pub use udp::UdpSocket;

0 commit comments

Comments
 (0)