Improve keepalive and inactivity timers

* Add an analogue of OpenSSH's `ServerAliveCountMax`.
* Use disjunctive futures for cleanly making these timers optional.
* Use the `Session` to pass information back to the main bg loop from
  the plaintext packet reader, so that only nontrivial data transfer
  will reset the inactivity timer. (And so that `ServerAliveCountMax`
  will be judged correctly.)

Author: mmirate

russh/src/client/encrypted.rs

@@ -540,6 +540,7 @@ impl Session {
                 );
                 match req {
                     b"xon-xoff" => {
+                        self.activity = false;
                         r.read_byte().map_err(crate::Error::from)?; // should be 0.
                         let client_can_do = r.read_byte().map_err(crate::Error::from)? != 0;
                         if let Some(chan) = self.channels.get(&channel_num) {
@@ -586,6 +587,7 @@ impl Session {
                             .await
                     }
                     b"keepalive@openssh.com" => {
+                        self.activity = false;
                         let wants_reply = r.read_byte().map_err(crate::Error::from)?;
                         if wants_reply == 1 {
                             if let Some(ref mut enc) = self.common.encrypted {
@@ -605,6 +607,7 @@ impl Session {
                         Ok((client, self))
                     }
                     _ => {
+                        self.activity = false;
                         let wants_reply = r.read_byte().map_err(crate::Error::from)?;
                         if wants_reply == 1 {
                             if let Some(ref mut enc) = self.common.encrypted {
@@ -704,6 +707,7 @@ impl Session {
                         push_packet!(enc.write, enc.write.push(msg::REQUEST_FAILURE))
                     }
                 }
+                self.activity = false;
                 Ok((client, self))
             }
             Some(&msg::CHANNEL_SUCCESS) => {
@@ -816,7 +820,16 @@ impl Session {
                     Err(crate::Error::Inconsistent.into())
                 }
             }
+            Some(&msg::REQUEST_SUCCESS | &msg::REQUEST_FAILURE)
+                if self.server_alive_timeouts > 0 =>
+            {
+                self.activity = false;
+                // TODO what other things might need to happen in response to these two opcodes?
+                self.server_alive_timeouts = 0;
+                Ok((client, self))
+            }
             _ => {
+                self.activity = false;
                 info!("Unhandled packet: {:?}", buf);
                 Ok((client, self))
             }

russh/src/client/mod.rs

@@ -105,7 +105,7 @@ use crate::session::{CommonSession, EncryptedState, Exchange, Kex, KexDhDone, Ke
 use crate::ssh_read::SshRead;
 use crate::sshbuffer::{SSHBuffer, SshId};
 use crate::{
-    auth, msg, negotiation, strict_kex_violation, timeout, ChannelId, ChannelOpenFailure,
+    auth, msg, negotiation, strict_kex_violation, ChannelId, ChannelOpenFailure,
     Disconnect, Limits, Sig,
 };
 
@@ -128,6 +128,8 @@ pub struct Session {
     pending_len: u32,
     inbound_channel_sender: Sender<Msg>,
     inbound_channel_receiver: Receiver<Msg>,
+    server_alive_timeouts: usize,
+    activity: bool,
 }
 
 const STRICT_KEX_MSG_ORDER: &[u8] = &[msg::KEXINIT, msg::KEX_ECDH_REPLY, msg::NEWKEYS];
@@ -728,6 +730,16 @@ async fn start_reading<R: AsyncRead + Unpin>(
     Ok((n, stream_read, buffer, cipher))
 }
 
+fn future_or_pending<F: futures::Future, T>(
+    val: Option<T>,
+    f: impl FnOnce(T) -> F,
+) -> futures::future::Either<futures::future::Pending<<F as futures::Future>::Output>, F> {
+    val.map_or(
+        futures::future::Either::Left(futures::future::pending()),
+        |x| futures::future::Either::Right(f(x)),
+    )
+}
+
 impl Session {
     fn new(
         target_window_size: u32,
@@ -746,6 +758,8 @@ impl Session {
             channels: HashMap::new(),
             pending_reads: Vec::new(),
             pending_len: 0,
+            server_alive_timeouts: 0,
+            activity: false,
         }
     }
 
@@ -774,19 +788,32 @@ impl Session {
         let mut opening_cipher = Box::new(clear::Key) as Box<dyn OpeningKey + Send>;
         std::mem::swap(&mut opening_cipher, &mut self.common.cipher.remote_to_local);
 
-        let time_for_keepalive = tokio::time::sleep_until(self.common.config.keepalive_deadline());
+        let keepalive_timer =
+            future_or_pending(self.common.config.keepalive_interval, tokio::time::sleep);
+        pin!(keepalive_timer);
+
+        let inactivity_timer =
+            future_or_pending(self.common.config.inactivity_timeout, tokio::time::sleep);
+        pin!(inactivity_timer);
+
         let reading = start_reading(stream_read, buffer, opening_cipher);
         pin!(reading);
-        pin!(time_for_keepalive);
-
-        let delay = self.common.config.inactivity_timeout;
 
         #[allow(clippy::panic)] // false positive in select! macro
         while !self.common.disconnected {
+            self.activity = false;
             tokio::select! {
-                () = &mut time_for_keepalive => {
-                    time_for_keepalive.as_mut().reset(self.common.config.keepalive_deadline());
+                () = &mut keepalive_timer => {
                     self.send_keepalive(true);
+                    if self.common.config.keepalive_max != 0 && self.server_alive_timeouts > self.common.config.keepalive_max {
+                        debug!("Timeout, server not responding to keepalives");
+                        break
+                    }
+                    self.server_alive_timeouts = self.server_alive_timeouts.saturating_add(1);
+                }
+                () = &mut inactivity_timer => {
+                    debug!("timeout");
+                    break
                 }
                 r = &mut reading => {
                     let (stream_read, mut buffer, mut opening_cipher) = match r {
@@ -819,6 +846,7 @@ impl Session {
                         if buf[0] == crate::msg::DISCONNECT {
                             break;
                         } else {
+                            self.activity = true;
                             let (h, s) = reply(self, handler, &mut encrypted_signal, &mut buffer.seqn, buf).await?;
                             handler = h;
                             self = s;
@@ -827,7 +855,6 @@ impl Session {
 
                     std::mem::swap(&mut opening_cipher, &mut self.common.cipher.remote_to_local);
                     reading.set(start_reading(stream_read, buffer, opening_cipher));
-                    time_for_keepalive.as_mut().reset(self.common.config.keepalive_deadline());
                 }
                 msg = self.receiver.recv(), if !self.is_rekeying() => {
                     match msg {
@@ -845,7 +872,6 @@ impl Session {
                             Err(_) => break
                         }
                     }
-                    time_for_keepalive.as_mut().reset(self.common.config.keepalive_deadline());
                 }
                 msg = self.inbound_channel_receiver.recv(), if !self.is_rekeying() => {
                     match msg {
@@ -861,17 +887,15 @@ impl Session {
                         }
                     }
                 }
-                _ = timeout(delay) => {
-                    debug!("timeout");
-                    break
-                },
-            }
+            };
+
             self.flush()?;
             if !self.common.write_buffer.buffer.is_empty() {
                 trace!(
                     "writing to stream: {:?} bytes",
                     self.common.write_buffer.buffer.len()
                 );
+                self.activity = true;
                 stream_write
                     .write_all(&self.common.write_buffer.buffer)
                     .await
@@ -885,6 +909,22 @@ impl Session {
                     enc.state = EncryptedState::Authenticated;
                 }
             }
+
+            if let (futures::future::Either::Right(ref mut sleep), Some(d)) = (
+                keepalive_timer.as_mut().as_pin_mut(),
+                self.common.config.keepalive_interval,
+            ) {
+                sleep.as_mut().reset(tokio::time::Instant::now() + d);
+            }
+
+            if self.activity {
+                if let (futures::future::Either::Right(ref mut sleep), Some(d)) = (
+                    inactivity_timer.as_mut().as_pin_mut(),
+                    self.common.config.inactivity_timeout,
+                ) {
+                    sleep.as_mut().reset(tokio::time::Instant::now() + d);
+                }
+            }
         }
         debug!("disconnected");
         self.receiver.close();
@@ -1310,19 +1350,12 @@ pub struct Config {
     pub inactivity_timeout: Option<std::time::Duration>,
     /// If nothing is sent or received for this amount of time, send a keepalive message.
     pub keepalive_interval: Option<std::time::Duration>,
+    /// If this many keepalives have been sent without reply, close the connection.
+    pub keepalive_max: usize,
     /// Whether to expect and wait for an authentication call.
     pub anonymous: bool,
 }
 
-impl Config {
-    fn keepalive_deadline(&self) -> tokio::time::Instant {
-        tokio::time::Instant::now()
-            + self
-                .keepalive_interval
-                .unwrap_or(std::time::Duration::from_secs(86400 * 365))
-    }
-}
-
 impl Default for Config {
     fn default() -> Config {
         Config {
@@ -1337,6 +1370,7 @@ impl Default for Config {
             preferred: Default::default(),
             inactivity_timeout: None,
             keepalive_interval: None,
+            keepalive_max: 3,
             anonymous: false,
         }
     }

russh/src/client/session.rs

@@ -292,7 +292,7 @@ impl Session {
         if let Some(ref mut enc) = self.common.encrypted {
             push_packet!(enc.write, {
                 enc.write.push(msg::GLOBAL_REQUEST);
-                enc.write.extend_ssh_string(b"keepalive@libssh2.org");
+                enc.write.extend_ssh_string(b"keepalive@openssh.org");
                 enc.write.push(want_reply as u8);
             });
         }