From 10f8bd4e067d2d8f674bd5be872587bedd6bbe70 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Sun, 23 Jun 2024 01:13:56 +0800
Subject: [PATCH] http3: fix race condition between Server.Serve and
 Server.Close (#4572)

---
 http3/server.go | 61 ++++++++++++++++++++++++++++++++++---------------
 1 file changed, 42 insertions(+), 19 deletions(-)

diff --git a/http3/server.go b/http3/server.go
index 7e6ffd54..cf52c174 100644
--- a/http3/server.go
+++ b/http3/server.go
@@ -216,7 +216,13 @@ type Server struct {
 //
 // If s.Addr is blank, ":https" is used.
 func (s *Server) ListenAndServe() error {
-	return s.serveConn(s.TLSConfig, nil)
+	ln, err := s.setupListenerForConn(s.TLSConfig, nil)
+	if err != nil {
+		return err
+	}
+	defer s.removeListener(&ln)
+
+	return s.serveListener(ln)
 }
 
 // ListenAndServeTLS listens on the UDP address s.Addr and calls s.Handler to handle HTTP/3 requests on incoming connections.
@@ -231,17 +237,26 @@ func (s *Server) ListenAndServeTLS(certFile, keyFile string) error {
 	}
 	// We currently only use the cert-related stuff from tls.Config,
 	// so we don't need to make a full copy.
-	config := &tls.Config{
-		Certificates: certs,
+	ln, err := s.setupListenerForConn(&tls.Config{Certificates: certs}, nil)
+	if err != nil {
+		return err
 	}
-	return s.serveConn(config, nil)
+	defer s.removeListener(&ln)
+
+	return s.serveListener(ln)
 }
 
 // Serve an existing UDP connection.
 // It is possible to reuse the same connection for outgoing connections.
 // Closing the server does not close the connection.
 func (s *Server) Serve(conn net.PacketConn) error {
-	return s.serveConn(s.TLSConfig, conn)
+	ln, err := s.setupListenerForConn(s.TLSConfig, conn)
+	if err != nil {
+		return err
+	}
+	defer s.removeListener(&ln)
+
+	return s.serveListener(ln)
 }
 
 // ServeQUICConn serves a single QUIC connection.
@@ -255,10 +270,18 @@ func (s *Server) ServeQUICConn(conn quic.Connection) error {
 // Closing the server does close the listener.
 // ServeListener always returns a non-nil error. After Shutdown or Close, the returned error is http.ErrServerClosed.
 func (s *Server) ServeListener(ln QUICEarlyListener) error {
+	s.mutex.Lock()
 	if err := s.addListener(&ln); err != nil {
+		s.mutex.Unlock()
 		return err
 	}
+	s.mutex.Unlock()
 	defer s.removeListener(&ln)
+
+	return s.serveListener(ln)
+}
+
+func (s *Server) serveListener(ln QUICEarlyListener) error {
 	for {
 		conn, err := ln.Accept(context.Background())
 		if err == quic.ErrServerClosed {
@@ -279,16 +302,9 @@ func (s *Server) ServeListener(ln QUICEarlyListener) error {
 
 var errServerWithoutTLSConfig = errors.New("use of http3.Server without TLSConfig")
 
-func (s *Server) serveConn(tlsConf *tls.Config, conn net.PacketConn) error {
+func (s *Server) setupListenerForConn(tlsConf *tls.Config, conn net.PacketConn) (QUICEarlyListener, error) {
 	if tlsConf == nil {
-		return errServerWithoutTLSConfig
-	}
-
-	s.mutex.Lock()
-	closed := s.closed
-	s.mutex.Unlock()
-	if closed {
-		return http.ErrServerClosed
+		return nil, errServerWithoutTLSConfig
 	}
 
 	baseConf := ConfigureTLSConfig(tlsConf)
@@ -302,6 +318,13 @@ func (s *Server) serveConn(tlsConf *tls.Config, conn net.PacketConn) error {
 		quicConf.EnableDatagrams = true
 	}
 
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+	closed := s.closed
+	if closed {
+		return nil, http.ErrServerClosed
+	}
+
 	var ln QUICEarlyListener
 	var err error
 	if conn == nil {
@@ -314,9 +337,12 @@ func (s *Server) serveConn(tlsConf *tls.Config, conn net.PacketConn) error {
 		ln, err = quicListen(conn, baseConf, quicConf)
 	}
 	if err != nil {
-		return err
+		return nil, err
 	}
-	return s.ServeListener(ln)
+	if err := s.addListener(&ln); err != nil {
+		return nil, err
+	}
+	return ln, nil
 }
 
 func extractPort(addr string) (int, error) {
@@ -392,9 +418,6 @@ func (s *Server) generateAltSvcHeader() {
 // call trackListener via Serve and can track+defer untrack the same pointer to
 // local variable there. We never need to compare a Listener from another caller.
 func (s *Server) addListener(l *QUICEarlyListener) error {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-
 	if s.closed {
 		return http.ErrServerClosed
 	}