use std::{fmt, io, vec};
use std::error::Error;
use std::net::{
IpAddr, Ipv4Addr, Ipv6Addr,
SocketAddr, ToSocketAddrs,
SocketAddrV4, SocketAddrV6,
};
use std::str::FromStr;
use tokio_sync::{mpsc, oneshot};
use crate::common::{Future, Never, Pin, Poll, Unpin, task};
pub trait Resolve: Unpin {
type Addrs: Iterator<Item=IpAddr>;
type Future: Future<Output=Result<Self::Addrs, io::Error>> + Unpin;
fn resolve(&self, name: Name) -> Self::Future;
}
#[derive(Clone, Hash, Eq, PartialEq)]
pub struct Name {
host: String,
}
#[derive(Clone)]
pub struct GaiResolver {
_priv: (),
}
#[derive(Clone)]
struct ThreadPoolKeepAlive(mpsc::Sender<Never>);
pub struct GaiAddrs {
inner: IpAddrs,
}
pub struct GaiFuture {
inner: tokio_executor::blocking::Blocking<Result<IpAddrs, io::Error>>,
}
impl Name {
pub(super) fn new(host: String) -> Name {
Name {
host,
}
}
pub fn as_str(&self) -> &str {
&self.host
}
}
impl fmt::Debug for Name {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Debug::fmt(&self.host, f)
}
}
impl fmt::Display for Name {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Display::fmt(&self.host, f)
}
}
impl FromStr for Name {
type Err = InvalidNameError;
fn from_str(host: &str) -> Result<Self, Self::Err> {
Ok(Name::new(host.to_owned()))
}
}
#[derive(Debug)]
pub struct InvalidNameError(());
impl fmt::Display for InvalidNameError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("Not a valid domain name")
}
}
impl Error for InvalidNameError {}
impl GaiResolver {
pub fn new() -> Self {
GaiResolver {
_priv: (),
}
}
}
impl Resolve for GaiResolver {
type Addrs = GaiAddrs;
type Future = GaiFuture;
fn resolve(&self, name: Name) -> Self::Future {
let blocking = tokio_executor::blocking::run(move || {
debug!("resolving host={:?}", name.host);
(&*name.host, 0).to_socket_addrs()
.map(|i| IpAddrs { iter: i })
});
GaiFuture {
inner: blocking,
}
}
}
impl fmt::Debug for GaiResolver {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.pad("GaiResolver")
}
}
impl Future for GaiFuture {
type Output = Result<GaiAddrs, io::Error>;
fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
Pin::new(&mut self.inner).poll(cx).map(|res| match res {
Ok(addrs) => Ok(GaiAddrs { inner: addrs }),
Err(err) => Err(err),
})
}
}
impl fmt::Debug for GaiFuture {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.pad("GaiFuture")
}
}
impl Iterator for GaiAddrs {
type Item = IpAddr;
fn next(&mut self) -> Option<Self::Item> {
self.inner.next().map(|sa| sa.ip())
}
}
impl fmt::Debug for GaiAddrs {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.pad("GaiAddrs")
}
}
pub(super) struct GaiBlocking {
host: String,
tx: Option<oneshot::Sender<io::Result<IpAddrs>>>,
}
impl GaiBlocking {
fn block(&self) -> io::Result<IpAddrs> {
debug!("resolving host={:?}", self.host);
(&*self.host, 0).to_socket_addrs()
.map(|i| IpAddrs { iter: i })
}
}
impl Future for GaiBlocking {
type Output = ();
fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
if self.tx.as_mut().expect("polled after complete").poll_closed(cx).is_ready() {
trace!("resolve future canceled for {:?}", self.host);
return Poll::Ready(());
}
let res = self.block();
let tx = self.tx.take().expect("polled after complete");
let _ = tx.send(res);
Poll::Ready(())
}
}
pub(super) struct IpAddrs {
iter: vec::IntoIter<SocketAddr>,
}
impl IpAddrs {
pub(super) fn new(addrs: Vec<SocketAddr>) -> Self {
IpAddrs { iter: addrs.into_iter() }
}
pub(super) fn try_parse(host: &str, port: u16) -> Option<IpAddrs> {
if let Ok(addr) = host.parse::<Ipv4Addr>() {
let addr = SocketAddrV4::new(addr, port);
return Some(IpAddrs { iter: vec![SocketAddr::V4(addr)].into_iter() })
}
let host = host.trim_start_matches('[').trim_end_matches(']');
if let Ok(addr) = host.parse::<Ipv6Addr>() {
let addr = SocketAddrV6::new(addr, port, 0, 0);
return Some(IpAddrs { iter: vec![SocketAddr::V6(addr)].into_iter() })
}
None
}
pub(super) fn split_by_preference(self) -> (IpAddrs, IpAddrs) {
let preferring_v6 = self.iter
.as_slice()
.first()
.map(SocketAddr::is_ipv6)
.unwrap_or(false);
let (preferred, fallback) = self.iter
.partition::<Vec<_>, _>(|addr| addr.is_ipv6() == preferring_v6);
(IpAddrs::new(preferred), IpAddrs::new(fallback))
}
pub(super) fn is_empty(&self) -> bool {
self.iter.as_slice().is_empty()
}
}
impl Iterator for IpAddrs {
type Item = SocketAddr;
#[inline]
fn next(&mut self) -> Option<SocketAddr> {
self.iter.next()
}
}
#[cfg(feature = "runtime")]
#[derive(Clone, Debug)]
pub struct TokioThreadpoolGaiResolver(());
#[cfg(feature = "runtime")]
#[derive(Debug)]
pub struct TokioThreadpoolGaiFuture {
name: Name,
}
#[cfg(feature = "runtime")]
impl TokioThreadpoolGaiResolver {
pub fn new() -> Self {
TokioThreadpoolGaiResolver(())
}
}
#[cfg(feature = "runtime")]
impl Resolve for TokioThreadpoolGaiResolver {
type Addrs = GaiAddrs;
type Future = TokioThreadpoolGaiFuture;
fn resolve(&self, name: Name) -> TokioThreadpoolGaiFuture {
TokioThreadpoolGaiFuture { name }
}
}
#[cfg(feature = "runtime")]
impl Future for TokioThreadpoolGaiFuture {
type Output = Result<GaiAddrs, io::Error>;
fn poll(self: Pin<&mut Self>, _cx: &mut task::Context<'_>) -> Poll<Self::Output> {
match ready!(tokio_executor::threadpool::blocking(|| (self.name.as_str(), 0).to_socket_addrs())) {
Ok(Ok(iter)) => Poll::Ready(Ok(GaiAddrs { inner: IpAddrs { iter } })),
Ok(Err(e)) => Poll::Ready(Err(e)),
Err(e) => Poll::Ready(Err(io::Error::new(io::ErrorKind::Other, e))),
}
}
}
#[cfg(test)]
mod tests {
use std::net::{Ipv4Addr, Ipv6Addr};
use super::*;
#[test]
fn test_ip_addrs_split_by_preference() {
let v4_addr = (Ipv4Addr::new(127, 0, 0, 1), 80).into();
let v6_addr = (Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1), 80).into();
let (mut preferred, mut fallback) =
IpAddrs { iter: vec![v4_addr, v6_addr].into_iter() }.split_by_preference();
assert!(preferred.next().unwrap().is_ipv4());
assert!(fallback.next().unwrap().is_ipv6());
let (mut preferred, mut fallback) =
IpAddrs { iter: vec![v6_addr, v4_addr].into_iter() }.split_by_preference();
assert!(preferred.next().unwrap().is_ipv6());
assert!(fallback.next().unwrap().is_ipv4());
}
#[test]
fn test_name_from_str() {
const DOMAIN: &str = "test.example.com";
let name = Name::from_str(DOMAIN).expect("Should be a valid domain");
assert_eq!(name.as_str(), DOMAIN);
assert_eq!(name.to_string(), DOMAIN);
}
#[test]
fn ip_addrs_try_parse_v6() {
let uri = ::http::Uri::from_static("http://[::1]:8080/");
let dst = super::super::Destination { uri };
let mut addrs = IpAddrs::try_parse(
dst.host(),
dst.port().expect("port")
).expect("try_parse");
let expected = "[::1]:8080".parse::<SocketAddr>().expect("expected");
assert_eq!(addrs.next(), Some(expected));
}
}