Close any remaining connections once the tui exits

src/io/stats.rs

@@ -19,15 +19,15 @@ pub struct Stats {
 
 impl Stats {
     pub fn add_connection(&self) {
-        self.connections.store(1, Ordering::Relaxed);
+        self.connections.fetch_add(1, Ordering::Relaxed);
     }
 
     pub fn add_rx_bytes(&self, n: usize) {
-        self.rx.store(n, Ordering::Relaxed);
+        self.rx.fetch_add(n, Ordering::Relaxed);
     }
 
     pub fn add_tx_bytes(&self, n: usize) {
-        self.tx.store(n, Ordering::Relaxed);
+        self.tx.fetch_add(n, Ordering::Relaxed);
     }
 
     pub fn connections(&self) -> usize {
@@ -39,7 +39,7 @@ impl Stats {
     }
 
     pub fn set_failed(&self, failed: bool) {
-        self.failed.store(failed, Ordering::Relaxed);
+        self.failed.fetch_and(failed, Ordering::Relaxed);
     }
 
     pub fn rx(&self) -> Unit {

src/main.rs

@@ -20,11 +20,28 @@ use tracing_subscriber::EnvFilter;
 use tracing_subscriber::layer::SubscriberExt;
 use tracing_subscriber::util::SubscriberInitExt;
 
+#[cfg(unix)]
+async fn sigterm() {
+    use tokio::signal::unix::SignalKind;
+
+    let mut sigterm =
+        tokio::signal::unix::signal(SignalKind::terminate()).expect("should be able to initialize");
+    sigterm.recv().await;
+}
+
+#[cfg(not(unix))]
+async fn sigterm() {
+    std::future::pending::<()>().await;
+}
+
 async fn shutdown_task(token: CancellationToken) {
     select! {
         _ = tokio::signal::ctrl_c() => {
             debug!("Received SIGINT");
         }
+        _ = sigterm() => {
+            debug!("Received SIGTERM");
+        }
         _ = token.cancelled() => {
             debug!("Application called for graceful shutdown");
         }

src/ssh/handler.rs

@@ -330,6 +330,26 @@ impl russh::server::Handler for Handler {
         Ok(session.channel_success(channel)?)
     }
 
+    async fn channel_close(
+        &mut self,
+        channel: ChannelId,
+        session: &mut Session,
+    ) -> Result<(), Self::Error> {
+        if let Some(pty_channel) = self.pty_channel
+            && pty_channel == channel
+        {
+            debug!("Pty channel closed");
+
+            session.disconnect(
+                russh::Disconnect::ByApplication,
+                "Remaining active connections have been closed",
+                "EN",
+            )?;
+        }
+
+        Ok(())
+    }
+
     async fn tcpip_forward(
         &mut self,
         address: &str,