improve style: flatten future

This commit is contained in:
Astro 2018-09-07 00:12:00 +02:00
parent f4007511ea
commit 81191041c4
3 changed files with 144 additions and 113 deletions

View file

@ -61,7 +61,7 @@ impl Client {
done(idna::domain_to_ascii(&jid.domain)) done(idna::domain_to_ascii(&jid.domain))
.map_err(|_| Error::Idna) .map_err(|_| Error::Idna)
.and_then(|domain| .and_then(|domain|
done(Connecter::from_lookup(&domain, "_xmpp-client._tcp", 5222)) done(Connecter::from_lookup(&domain, Some("_xmpp-client._tcp"), 5222))
.map_err(Error::Connection) .map_err(Error::Connection)
) )
.and_then(|connecter| .and_then(|connecter|
@ -75,10 +75,8 @@ impl Client {
} else { } else {
Err(Error::Protocol(ProtocolError::NoTls)) Err(Error::Protocol(ProtocolError::NoTls))
} }
}).and_then(|starttls| }).flatten()
// TODO: flatten? .and_then(|tls_stream|
starttls
).and_then(|tls_stream|
XMPPStream::start(tls_stream, jid2, NS_JABBER_CLIENT.to_owned()) XMPPStream::start(tls_stream, jid2, NS_JABBER_CLIENT.to_owned())
).and_then(move |xmpp_stream| ).and_then(move |xmpp_stream|
done(Self::auth(xmpp_stream, username, password)) done(Self::auth(xmpp_stream, username, password))

View file

@ -53,8 +53,8 @@ impl Component {
fn make_connect(jid: Jid, password: String, server: &str, port: u16) -> impl Future<Item=XMPPStream, Error=Error> { fn make_connect(jid: Jid, password: String, server: &str, port: u16) -> impl Future<Item=XMPPStream, Error=Error> {
let jid1 = jid.clone(); let jid1 = jid.clone();
let password = password; let password = password;
done(Connecter::from_lookup(server, "_xmpp-component._tcp", port)) done(Connecter::from_lookup(server, None, port))
.and_then(|connecter| connecter) .flatten()
.map_err(Error::Connection) .map_err(Error::Connection)
.and_then(move |tcp_stream| { .and_then(move |tcp_stream| {
xmpp_stream::XMPPStream::start(tcp_stream, jid1, NS_JABBER_COMPONENT_ACCEPT.to_owned()) xmpp_stream::XMPPStream::start(tcp_stream, jid1, NS_JABBER_COMPONENT_ACCEPT.to_owned())

View file

@ -1,44 +1,61 @@
use std::mem; use std::mem;
use std::net::{SocketAddr, IpAddr}; use std::net::SocketAddr;
use std::collections::{BTreeMap, btree_map}; use std::collections::BTreeMap;
use std::collections::VecDeque; use std::collections::VecDeque;
use std::cell::RefCell;
use futures::{Future, Poll, Async}; use futures::{Future, Poll, Async};
use tokio::net::{ConnectFuture, TcpStream}; use tokio::net::{ConnectFuture, TcpStream};
use trust_dns_resolver::{IntoName, Name, ResolverFuture, error::ResolveError}; use trust_dns_resolver::{IntoName, Name, ResolverFuture, error::ResolveError};
use trust_dns_resolver::lookup::SrvLookupFuture; use trust_dns_resolver::lookup::SrvLookupFuture;
use trust_dns_resolver::lookup_ip::LookupIpFuture; use trust_dns_resolver::lookup_ip::LookupIpFuture;
use trust_dns_proto::rr::rdata::srv::SRV;
use ConnecterError; use ConnecterError;
enum State {
AwaitResolver(Box<Future<Item = ResolverFuture, Error = ResolveError> + Send>),
ResolveSrv(ResolverFuture, SrvLookupFuture),
ResolveTarget(ResolverFuture, LookupIpFuture, u16),
Connecting(Option<ResolverFuture>, Vec<RefCell<ConnectFuture>>),
Invalid,
}
pub struct Connecter { pub struct Connecter {
fallback_port: u16, fallback_port: u16,
name: Name, srv_domain: Option<Name>,
domain: Name, domain: Name,
resolver_future: Box<Future<Item = ResolverFuture, Error = ResolveError> + Send>, state: State,
resolver_opt: Option<ResolverFuture>, targets: VecDeque<(Name, u16)>,
srv_lookup_opt: Option<SrvLookupFuture>,
srvs_opt: Option<btree_map::IntoIter<u16, SRV>>,
ip_lookup_opt: Option<(u16, LookupIpFuture)>,
ips_opt: Option<(u16, VecDeque<IpAddr>)>,
connect_opt: Option<ConnectFuture>,
} }
impl Connecter { impl Connecter {
pub fn from_lookup(domain: &str, srv: &str, fallback_port: u16) -> Result<Connecter, ConnecterError> { pub fn from_lookup(domain: &str, srv: Option<&str>, fallback_port: u16) -> Result<Connecter, ConnecterError> {
if let Ok(ip) = domain.parse() {
// use specified IP address, not domain name, skip the whole dns part
let connect =
RefCell::new(TcpStream::connect(&SocketAddr::new(ip, fallback_port)));
return Ok(Connecter {
fallback_port,
srv_domain: None,
domain: "nohost".into_name()?,
state: State::Connecting(None, vec![connect]),
targets: VecDeque::new(),
});
}
let resolver_future = ResolverFuture::from_system_conf()?; let resolver_future = ResolverFuture::from_system_conf()?;
let name = format!("{}.{}.", srv, domain).into_name()?; let state = State::AwaitResolver(resolver_future);
let srv_domain = match srv {
Some(srv) =>
Some(format!("{}.{}.", srv, domain).into_name()?),
None =>
None,
};
Ok(Connecter { Ok(Connecter {
fallback_port, fallback_port,
name, srv_domain,
domain: domain.into_name()?, domain: domain.into_name()?,
resolver_future, state,
resolver_opt: None, targets: VecDeque::new(),
srv_lookup_opt: None,
srvs_opt: None,
ip_lookup_opt: None,
ips_opt: None,
connect_opt: None,
}) })
} }
} }
@ -48,102 +65,118 @@ impl Future for Connecter {
type Error = ConnecterError; type Error = ConnecterError;
fn poll(&mut self) -> Poll<Self::Item, Self::Error> { fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
if self.resolver_opt.is_none() { let state = mem::replace(&mut self.state, State::Invalid);
//println!("Poll resolver future"); match state {
match self.resolver_future.poll()? { State::AwaitResolver(mut resolver_future) => {
Async::Ready(resolver) => match resolver_future.poll()? {
self.resolver_opt = Some(resolver), Async::NotReady => {
Async::NotReady => self.state = State::AwaitResolver(resolver_future);
return Ok(Async::NotReady),
}
}
if let Some(ref resolver) = self.resolver_opt {
if self.srvs_opt.is_none() {
if self.srv_lookup_opt.is_none() {
//println!("Lookup srv: {:?}", self.name);
self.srv_lookup_opt = Some(resolver.lookup_srv(&self.name));
}
if let Some(ref mut srv_lookup) = self.srv_lookup_opt {
match srv_lookup.poll() {
Ok(Async::Ready(t)) => {
let mut srvs = BTreeMap::new();
for srv in t.iter() {
srvs.insert(srv.priority(), srv.clone());
}
srvs.insert(65535, SRV::new(65535, 0, self.fallback_port, self.domain.clone()));
self.srvs_opt = Some(srvs.into_iter());
}
Ok(Async::NotReady) => return Ok(Async::NotReady),
Err(_) => {
//println!("Ignore SVR error: {:?}", e);
let mut srvs = BTreeMap::new();
srvs.insert(65535, SRV::new(65535, 0, self.fallback_port, self.domain.clone()));
self.srvs_opt = Some(srvs.into_iter());
},
}
}
}
if self.connect_opt.is_none() {
if self.ips_opt.is_none() {
if self.ip_lookup_opt.is_none() {
if let Some(ref mut srvs) = self.srvs_opt {
if let Some((_, srv)) = srvs.next() {
//println!("Lookup ip: {:?}", srv);
self.ip_lookup_opt = Some((srv.port(), resolver.lookup_ip(srv.target())));
} else {
return Err(ConnecterError::NoSrv);
}
}
}
if let Some((port, mut ip_lookup)) = mem::replace(&mut self.ip_lookup_opt, None) {
match ip_lookup.poll() {
Ok(Async::Ready(t)) => {
let mut ip_deque = VecDeque::new();
ip_deque.extend(t.iter());
//println!("IPs: {:?}", ip_deque);
self.ips_opt = Some((port, ip_deque));
self.ip_lookup_opt = None;
},
Ok(Async::NotReady) => {
self.ip_lookup_opt = Some((port, ip_lookup));
return Ok(Async::NotReady)
},
Err(_) => {
//println!("Ignore lookup error: {:?}", e);
self.ip_lookup_opt = None;
}
}
}
}
if let Some((port, mut ip_deque)) = mem::replace(&mut self.ips_opt, None) {
if let Some(ip) = ip_deque.pop_front() {
//println!("Connect to {:?}:{}", ip, port);
self.connect_opt = Some(TcpStream::connect(&SocketAddr::new(ip, port)));
self.ips_opt = Some((port, ip_deque));
}
}
}
if let Some(mut connect_future) = mem::replace(&mut self.connect_opt, None) {
match connect_future.poll() {
Ok(Async::Ready(t)) => return Ok(Async::Ready(t)),
Ok(Async::NotReady) => {
self.connect_opt = Some(connect_future);
return Ok(Async::NotReady)
}
Err(_) => {
//println!("Ignore connect error: {:?}", e);
},
}
}
}
Ok(Async::NotReady) Ok(Async::NotReady)
} }
Async::Ready(resolver) => {
match &self.srv_domain {
&Some(ref srv_domain) => {
let srv_lookup = resolver.lookup_srv(srv_domain);
self.state = State::ResolveSrv(resolver, srv_lookup);
}
None => {
self.targets =
[(self.domain.clone(), self.fallback_port)].into_iter()
.cloned()
.collect();
self.state = State::Connecting(Some(resolver), vec![]);
}
}
self.poll()
}
}
}
State::ResolveSrv(resolver, mut srv_lookup) => {
match srv_lookup.poll() {
Ok(Async::NotReady) => {
self.state = State::ResolveSrv(resolver, srv_lookup);
Ok(Async::NotReady)
}
Ok(Async::Ready(srv_result)) => {
let mut srv_map: BTreeMap<_, _> =
srv_result.iter()
.map(|srv| (srv.priority(), (srv.target().clone(), srv.port())))
.collect();
let targets =
srv_map.into_iter()
.map(|(_, tp)| tp)
.collect();
self.targets = targets;
self.state = State::Connecting(Some(resolver), vec![]);
self.poll()
}
Err(_) => {
// ignore, fallback
self.targets =
[(self.domain.clone(), self.fallback_port)].into_iter()
.cloned()
.collect();
self.state = State::Connecting(Some(resolver), vec![]);
self.poll()
}
}
}
State::Connecting(resolver, mut connects) => {
if resolver.is_some() &&
connects.len() == 0 &&
self.targets.len() > 0 {
let resolver = resolver.unwrap();
let (host, port) = self.targets.pop_front().unwrap();
let ip_lookup = resolver.lookup_ip(host);
self.state = State::ResolveTarget(resolver, ip_lookup, port);
self.poll()
} else if connects.len() > 0 {
let mut success = None;
connects.retain(|connect| {
match connect.borrow_mut().poll() {
Ok(Async::NotReady) => true,
Ok(Async::Ready(connection)) => {
success = Some(connection);
false
}
Err(_) => false,
}
});
match success {
Some(connection) =>
Ok(Async::Ready(connection)),
None => {
self.state = State::Connecting(resolver, connects);
Ok(Async::NotReady)
},
}
} else {
Err(ConnecterError::AllFailed)
}
}
State::ResolveTarget(resolver, mut ip_lookup, port) => {
match ip_lookup.poll() {
Ok(Async::NotReady) => {
self.state = State::ResolveTarget(resolver, ip_lookup, port);
Ok(Async::NotReady)
}
Ok(Async::Ready(ip_result)) => {
let connects =
ip_result.iter()
.map(|ip| RefCell::new(TcpStream::connect(&SocketAddr::new(ip, port))))
.collect();
self.state = State::Connecting(Some(resolver), connects);
self.poll()
}
Err(_) => {
// ignore, next…
self.state = State::Connecting(Some(resolver), vec![]);
self.poll()
}
}
}
_ => panic!("")
}
}
} }