From 152a5cab87b1eb08be9014c67fd9d4b396c2a0e1 Mon Sep 17 00:00:00 2001
From: Naatan <n.rijksen@gmail.com>
Date: Thu, 18 Dec 2025 10:36:33 -0800
Subject: [PATCH] Fix shutdown not actually shutting down

---
 smtpd.go      | 17 ++++++++---------
 smtpd_test.go | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/smtpd.go b/smtpd.go
index 2cace97..02bdfb8 100644
--- a/smtpd.go
+++ b/smtpd.go
@@ -197,21 +197,20 @@ func (srv *Server) Serve(ln net.Listener) error {
 		return ErrServerClosed
 	}
 
-	defer ln.Close()
-	for {
-
-		// if we are shutting down, don't accept new connections
-		select {
-		case <-srv.getShutdownChan():
-			return ErrServerClosed
-		default:
-		}
+	go func() {
+		<-srv.getShutdownChan()
+		ln.Close()
+	}()
 
+	for {
 		conn, err := ln.Accept()
 		if err != nil {
 			if netErr, ok := err.(net.Error); ok && netErr.Temporary() {
 				continue
 			}
+			if errors.Is(err, net.ErrClosed) {
+				return ErrServerClosed
+			}
 			return err
 		}
 
diff --git a/smtpd_test.go b/smtpd_test.go
index 6fbdc70..9cf4e7d 100644
--- a/smtpd_test.go
+++ b/smtpd_test.go
@@ -1602,3 +1602,35 @@ func TestCmdShutdown(t *testing.T) {
 
 	conn.Close()
 }
+
+func TestGracefulShutdown(t *testing.T) {
+	// Run server
+	srv := &Server{Addr: ":1234"}
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		err := srv.ListenAndServe()
+		if err != nil && !errors.Is(err, ErrServerClosed) {
+			t.Errorf("Error serving: %v\n", err)
+		}
+	}()
+
+	// Shutdown server
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	if err := srv.Shutdown(ctx); err != nil {
+		t.Errorf("Error shutting down server: %v\n", err)
+	}
+
+	// Verify that port is no longer listening
+	_, err := net.Listen("tcp", srv.Addr)
+	if err != nil {
+		t.Errorf("Port is still listening: %v\n", err)
+	}
+
+	select {
+	case <-done:
+	case <-time.After(5 * time.Second):
+		t.Errorf("Timeout waiting for server to finish")
+	}
+}