delay moving to the connected state until we get an ack from the server

this should make adding additional handshake steps during login more
straightforward

3 files changed, 35 insertions, 3 deletions

diff --git a/src/client.rs b/src/client.rs
index 605b68d..b51784e 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -297,9 +297,7 @@ impl<S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + 'static>
 
                     self.reset_reconnect_timer();
 
-                    Ok(crate::component_future::Poll::Event(Event::Connect(
-                        self.size.clone().unwrap(),
-                    )))
+                    Ok(crate::component_future::Poll::DidWork)
                 }
                 Ok(futures::Async::NotReady) => {
                     Ok(crate::component_future::Poll::NotReady)
@@ -338,6 +336,11 @@ impl<S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + 'static>
                     self.last_server_time = std::time::Instant::now();
                     self.rsock = ReadSocket::Connected(s);
                     match msg {
+                        crate::protocol::Message::LoggedIn { .. } => {
+                            Ok(crate::component_future::Poll::Event(
+                                Event::Connect(self.size.clone().unwrap()),
+                            ))
+                        }
                         crate::protocol::Message::Heartbeat => {
                             Ok(crate::component_future::Poll::DidWork)
                         }
diff --git a/src/protocol.rs b/src/protocol.rs
index b47cb37..37fe793 100644
--- a/src/protocol.rs
+++ b/src/protocol.rs
@@ -88,6 +88,9 @@ pub enum Message {
     Resize {
         size: crate::term::Size,
     },
+    LoggedIn {
+        username: String,
+    },
 }
 
 const MSG_LOGIN: u32 = 0;
@@ -100,6 +103,7 @@ const MSG_SESSIONS: u32 = 6;
 const MSG_DISCONNECTED: u32 = 7;
 const MSG_ERROR: u32 = 8;
 const MSG_RESIZE: u32 = 9;
+const MSG_LOGGED_IN: u32 = 10;
 
 impl Message {
     pub fn login_plain(
@@ -159,6 +163,12 @@ impl Message {
         Self::Resize { size: size.clone() }
     }
 
+    pub fn logged_in(username: &str) -> Self {
+        Self::LoggedIn {
+            username: username.to_string(),
+        }
+    }
+
     #[allow(dead_code)]
     pub fn read<R: std::io::Read>(r: R) -> Result<Self> {
         Packet::read(r).and_then(Self::try_from)
@@ -418,6 +428,16 @@ impl From<&Message> for Packet {
                     data,
                 }
             }
+            Message::LoggedIn { username } => {
+                let mut data = vec![];
+
+                write_str(username, &mut data);
+
+                Self {
+                    ty: MSG_LOGGED_IN,
+                    data,
+                }
+            }
         }
     }
 }
@@ -563,6 +583,11 @@ impl std::convert::TryFrom<Packet> for Message {
 
                 (Self::Resize { size }, data)
             }
+            MSG_LOGGED_IN => {
+                let (username, data) = read_str(data)?;
+
+                (Self::LoggedIn { username }, data)
+            }
             _ => return Err(Error::InvalidMessageType { ty: packet.ty }),
         };
 
@@ -701,6 +726,7 @@ mod test {
             Message::disconnected(),
             Message::error("error message"),
             Message::resize(&crate::term::Size { rows: 25, cols: 81 }),
+            Message::logged_in("doy"),
         ]
     }
 
diff --git a/src/server.rs b/src/server.rs
index 0de39ba..9031425 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -301,6 +301,9 @@ impl<S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + 'static>
             crate::protocol::Auth::Plain { username } => {
                 log::info!("{}: login({})", conn.id, username);
                 conn.state.login(&username, term_type, &size);
+                conn.send_message(crate::protocol::Message::logged_in(
+                    &username,
+                ));
             }
         }