quic-go · birneee · Jan 8, 2024 · Jan 15, 2024
diff --git a/integrationtests/self/handshake_test.go b/integrationtests/self/handshake_test.go
@@ -75,9 +75,11 @@ var _ = Describe("Handshake tests", func() {
 			defer GinkgoRecover()
 			defer close(acceptStopped)
 			for {
-				if _, err := server.Accept(context.Background()); err != nil {
+				conn, err := server.Accept(context.Background())
+				if err != nil {
 					return
 				}
+				defer conn.CloseWithError(0, "")
 			}
 		}()
 	}
@@ -368,7 +370,8 @@ var _ = Describe("Handshake tests", func() {
 			time.Sleep(scaleDuration(200 * time.Millisecond))
 
 			// dial again, and expect that this dial succeeds
-			_, err = dial()
+			anotherConn, err := dial()
+			defer anotherConn.CloseWithError(0, "")
 			Expect(err).ToNot(HaveOccurred())
 			time.Sleep(scaleDuration(20 * time.Millisecond)) // wait a bit for the connection to be queued
 

diff --git a/integrationtests/self/self_suite_test.go b/integrationtests/self/self_suite_test.go
@@ -7,6 +7,7 @@ import (
 	"crypto/x509"
 	"flag"
 	"fmt"
+	"github.com/quic-go/quic-go/internal/testutils"
 	"log"
 	"os"
 	"runtime/pprof"
@@ -209,6 +210,10 @@ var _ = BeforeEach(func() {
 	}
 })
 
+var _ = AfterEach(func() {
+	Eventually(testutils.AreConnsRunning).Should(BeFalse())
+})
+
 func areHandshakesRunning() bool {
 	var b bytes.Buffer
 	pprof.Lookup("goroutine").WriteTo(&b, 1)

diff --git a/internal/testutils/testutils.go b/internal/testutils/testutils.go
@@ -1,7 +1,10 @@
 package testutils
 
 import (
+	"bytes"
 	"fmt"
+	"runtime/pprof"
+	"strings"
 
 	"github.com/quic-go/quic-go/internal/handshake"
 	"github.com/quic-go/quic-go/internal/protocol"
@@ -99,3 +102,9 @@ func ComposeRetryPacket(
 	data := writePacket(hdr, nil)
 	return append(data, handshake.GetRetryIntegrityTag(data, origDestConnID, version)[:]...)
 }
+
+func AreConnsRunning() bool {
+	var b bytes.Buffer
+	pprof.Lookup("goroutine").WriteTo(&b, 1)
+	return strings.Contains(b.String(), "quic-go.(*connection).run")
+}
diff --git a/server.go b/server.go
@@ -101,10 +101,11 @@ type baseServer struct {
 		protocol.VersionNumber,
 	) quicConn
 
-	closeOnce sync.Once
-	errorChan chan struct{} // is closed when the server is closed
-	closeErr  error
-	running   chan struct{} // closed as soon as run() returns
+	closeOnce          sync.Once
+	errorChan          chan struct{} // is closed when the server is closed
+	closeErr           error
+	running            chan struct{}  // closed as soon as run() returns
+	handshakeWaitGroup sync.WaitGroup // connections that are started but not passed to connQueue yet
 
 	versionNegotiationQueue chan receivedPacket
 	invalidTokenQueue       chan rejectedPacket
@@ -340,6 +341,19 @@ func (s *baseServer) close(e error, notifyOnClose bool) {
 		close(s.errorChan)
 
 		<-s.running
+		s.handshakeWaitGroup.Wait()
+
+	loop: // drain connQueue
+		for {
+			select {
+			case conn := <-s.connQueue:
+				conn.destroy(&qerr.TransportError{ErrorCode: ConnectionRefused})
+				atomic.AddInt32(&s.connQueueLen, -1)
+			default:
+				break loop
+			}
+		}
+
 		if notifyOnClose {
 			s.onClose()
 		}
@@ -681,6 +695,7 @@ func (s *baseServer) handleInitialImpl(p receivedPacket, hdr *wire.Header) error
 		return nil
 	}
 	go conn.run()
+	s.handshakeWaitGroup.Add(1)
 	go s.handleNewConn(conn)
 	if conn == nil {
 		p.buffer.Release()
@@ -690,6 +705,7 @@ func (s *baseServer) handleInitialImpl(p receivedPacket, hdr *wire.Header) error
 }
 
 func (s *baseServer) handleNewConn(conn quicConn) {
+	defer s.handshakeWaitGroup.Done()
 	connCtx := conn.Context()
 	if s.acceptEarlyConns {
 		// wait until the early connection is ready, the handshake fails, or the server is closed
@@ -720,6 +736,10 @@ func (s *baseServer) handleNewConn(conn quicConn) {
 	case <-connCtx.Done():
 		atomic.AddInt32(&s.connQueueLen, -1)
 		// don't pass connections that were already closed to Accept()
+	case <-s.errorChan:
+		conn.destroy(&qerr.TransportError{ErrorCode: ConnectionRefused})
+		atomic.AddInt32(&s.connQueueLen, -1)
+		// don't pass connections to Accept() if server already closed
 	}
 }
 

diff --git a/server_test.go b/server_test.go
@@ -660,6 +660,7 @@ var _ = Describe("Server", func() {
 					c := make(chan struct{})
 					close(c)
 					conn.EXPECT().HandshakeComplete().Return(c)
+					conn.EXPECT().destroy(gomock.Any())
 					return conn
 				}
 
@@ -1276,6 +1277,7 @@ var _ = Describe("Server", func() {
 				conn.EXPECT().run()
 				conn.EXPECT().earlyConnReady().Return(ready)
 				conn.EXPECT().Context().Return(context.Background())
+				conn.EXPECT().destroy(gomock.Any())
 				return conn
 			}