diff --git a/examples/client.rs b/examples/client.rs index ba94d9118..3a6b295b4 100644 --- a/examples/client.rs +++ b/examples/client.rs @@ -9,6 +9,11 @@ fn main() { let server = env::args().nth(1).expect("need server name as an argument"); println!( "RESULT: {:?}", - quinn::Client::connect(&server, 4433).wait() + quinn::Client::connect(&server, 4433) + .and_then(|_| { + println!("client is connected"); + futures::future::ok(()) + }) + .wait(), ); } diff --git a/src/client.rs b/src/client.rs index 228775f1c..a1113fe62 100644 --- a/src/client.rs +++ b/src/client.rs @@ -6,6 +6,7 @@ use tls; use types::Side; use std::io; +use std::mem; use std::net::ToSocketAddrs; use tokio::net::UdpSocket; @@ -26,60 +27,69 @@ impl Client { let addr = (server, port).to_socket_addrs().unwrap().next().unwrap(); let socket = UdpSocket::bind(&"0.0.0.0:0".parse().unwrap()).unwrap(); socket.connect(&addr).unwrap(); - ConnectFuture { - client: Client { + ConnectFuture::Waiting( + Client { endpoint, socket, buf, }, - state: ConnectionState::Sending, - } + ConnectionState::Sending, + ) } } #[must_use = "futures do nothing unless polled"] -pub struct ConnectFuture { - client: Client, - state: ConnectionState, +pub enum ConnectFuture { + Waiting(Client, ConnectionState), + Empty, } impl Future for ConnectFuture { - type Item = (); + type Item = Client; type Error = io::Error; fn poll(&mut self) -> Poll { let mut waiting; + let mut done = false; loop { waiting = true; - if self.state == ConnectionState::Sending { - let len = try_ready!(self.client.socket.poll_send(&self.client.buf)); - debug_assert_eq!(len, self.client.buf.len()); - let size = self.client.buf.capacity(); - self.client.buf.resize(size, 0); - self.state = ConnectionState::Receiving; - waiting = false; - } + if let ConnectFuture::Waiting(ref mut client, ref mut state) = *self { + if let ConnectionState::Sending = *state { + let len = try_ready!(client.socket.poll_send(&client.buf)); + debug_assert_eq!(len, client.buf.len()); + let size = client.buf.capacity(); + client.buf.resize(size, 0); - if self.state == ConnectionState::Receiving { - let len = try_ready!(self.client.socket.poll_recv(&mut self.client.buf)); - let packet = { - let partial = Packet::start_decode(&mut self.client.buf[..len]); - let key = self.client.endpoint.decode_key(&partial.header); - partial.finish(&key) - }; - - let req = match self.client.endpoint.handle_handshake(&packet) { - Some(p) => p, - None => { - return Err(io::Error::new( - io::ErrorKind::Other, - "no response to packet", - )); + if !client.endpoint.is_handshaking() { + done = true; + break; } - }; - req.encode(&self.client.endpoint.encode_key(&req.header), &mut self.client.buf); - self.state = ConnectionState::Sending; - waiting = false; + *state = ConnectionState::Receiving; + waiting = false; + } + + if let ConnectionState::Receiving = *state { + let len = try_ready!(client.socket.poll_recv(&mut client.buf)); + let packet = { + let partial = Packet::start_decode(&mut client.buf[..len]); + let key = client.endpoint.decode_key(&partial.header); + partial.finish(&key) + }; + + let req = match client.endpoint.handle_handshake(&packet) { + Some(p) => p, + None => { + return Err(io::Error::new( + io::ErrorKind::Other, + "no response to packet", + )); + } + }; + + req.encode(&client.endpoint.encode_key(&req.header), &mut client.buf); + *state = ConnectionState::Sending; + waiting = false; + } } if waiting { @@ -87,12 +97,18 @@ impl Future for ConnectFuture { } } - Ok(Async::NotReady) + if done { + match mem::replace(self, ConnectFuture::Empty) { + ConnectFuture::Waiting(client, _) => Ok(Async::Ready(client)), + _ => panic!("invalid future state"), + } + } else { + Ok(Async::NotReady) + } } } -#[derive(PartialEq)] -enum ConnectionState { +pub enum ConnectionState { Sending, Receiving, } diff --git a/src/endpoint.rs b/src/endpoint.rs index 723f445a3..9bf9ea0bb 100644 --- a/src/endpoint.rs +++ b/src/endpoint.rs @@ -49,6 +49,13 @@ where } } + pub fn is_handshaking(&self) -> bool { + match self.secret { + Secret::For1Rtt(_, _, _) => false, + Secret::Handshake(_) => true, + } + } + pub fn update_src_cid(&mut self) { self.src_cid = thread_rng().gen(); }