From 216cceeb8e5e3e4015abe19dacd345398fb27303 Mon Sep 17 00:00:00 2001
From: Ayman Bagabas <ayman.bagabas@gmail.com>
Date: Tue, 23 Aug 2022 15:12:31 -0400
Subject: [PATCH] fix: graceful shutdown

---
 server/server.go | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/server/server.go b/server/server.go
index cbf40fd66..fed56f358 100644
--- a/server/server.go
+++ b/server/server.go
@@ -67,12 +67,18 @@ func (srv *Server) Reload() error {
 
 // Start starts the SSH server.
 func (srv *Server) Start() error {
-	return srv.SSHServer.ListenAndServe()
+	if err := srv.SSHServer.ListenAndServe(); err != ssh.ErrServerClosed {
+		return err
+	}
+	return nil
 }
 
 // Serve serves the SSH server using the provided listener.
 func (srv *Server) Serve(l net.Listener) error {
-	return srv.SSHServer.Serve(l)
+	if err := srv.SSHServer.Serve(l); err != ssh.ErrServerClosed {
+		return err
+	}
+	return nil
 }
 
 // Shutdown lets the server gracefully shutdown.