Implemented more graceful shutdown

src/ssh/mod.rs

@@ -11,7 +11,9 @@ use russh::MethodKind;
 use russh::keys::PrivateKey;
 use russh::server::Server as _;
 use tokio::net::ToSocketAddrs;
-use tracing::{debug, warn};
+use tokio::select;
+use tokio_util::sync::CancellationToken;
+use tracing::{debug, error, warn};
 
 use crate::ldap::Ldap;
 use crate::tunnel::Registry;
@@ -19,18 +21,30 @@ use crate::tunnel::Registry;
 pub struct Server {
     ldap: Ldap,
     registry: Registry,
+    token: CancellationToken,
+}
+
+async fn graceful_shutdown(token: CancellationToken) {
+    token.cancelled().await;
+    let duration = 1;
+    // All pty sessions will close once the token is cancelled, but to properly allow the sessions
+    // to close the ssh server still needs to be driven, so we let it run a little bit longer.
+    // TODO: Figure out a way to wait for all connections to be closed, would require also closing
+    // non-pty sessions somehow
+    debug!("Waiting for {duration}s before stopping");
+    tokio::time::sleep(Duration::from_secs(duration)).await;
 }
 
 impl Server {
-    pub fn new(ldap: Ldap, registry: Registry) -> Self {
-        Server { ldap, registry }
+    pub fn new(ldap: Ldap, registry: Registry, token: CancellationToken) -> Self {
+        Server {
+            ldap,
+            registry,
+            token,
+        }
     }
 
-    pub fn run(
-        &mut self,
-        key: PrivateKey,
-        addr: impl ToSocketAddrs + Send + std::fmt::Debug,
-    ) -> impl Future<Output = Result<(), std::io::Error>> + Send {
+    pub async fn run(mut self, key: PrivateKey, addr: impl ToSocketAddrs + Send + std::fmt::Debug) {
         let config = russh::server::Config {
             inactivity_timeout: Some(Duration::from_secs(3600)),
             auth_rejection_time: Duration::from_secs(1),
@@ -47,7 +61,17 @@ impl Server {
 
         debug!(?addr, "Running ssh");
 
-        async move { self.run_on_address(config, addr).await }
+        let token = self.token.clone();
+        select! {
+            res = self.run_on_address(config, addr) => {
+                if let Err(err) = res {
+                    error!("SSH Server error: {err}");
+                }
+            }
+            _ = graceful_shutdown(token) => {
+                debug!("Graceful shutdown");
+            }
+        }
     }
 }
 
@@ -55,7 +79,7 @@ impl russh::server::Server for Server {
     type Handler = Handler;
 
     fn new_client(&mut self, _peer_addr: Option<SocketAddr>) -> Self::Handler {
-        Handler::new(self.ldap.clone(), self.registry.clone())
+        Handler::new(self.ldap.clone(), self.registry.clone(), self.token.clone())
     }
 
     fn handle_session_error(&mut self, error: <Self::Handler as russh::server::Handler>::Error) {