From 4c25cea9ece839cc39395686cd41f18862e080ce Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Thu, 29 Oct 2020 17:56:14 +0700
Subject: [PATCH] fix deadlock when closing the server and the connection at
 the same time

---
 packet_handler_map.go      | 4 ++++
 packet_handler_map_test.go | 1 +
 server.go                  | 2 +-
 3 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/packet_handler_map.go b/packet_handler_map.go
index 3b327135f..a458ee985 100644
--- a/packet_handler_map.go
+++ b/packet_handler_map.go
@@ -238,6 +238,10 @@ func (h *packetHandlerMap) SetServer(s unknownPacketHandler) {
 
 func (h *packetHandlerMap) CloseServer() {
 	h.mutex.Lock()
+	if h.server == nil {
+		h.mutex.Unlock()
+		return
+	}
 	h.server = nil
 	var wg sync.WaitGroup
 	for _, handler := range h.handlers {
diff --git a/packet_handler_map_test.go b/packet_handler_map_test.go
index 95d9737c7..06cb9b8f8 100644
--- a/packet_handler_map_test.go
+++ b/packet_handler_map_test.go
@@ -251,6 +251,7 @@ var _ = Describe("Packet Handler Map", func() {
 			})
 
 			It("closes all server sessions", func() {
+				handler.SetServer(NewMockUnknownPacketHandler(mockCtrl))
 				clientSess := NewMockPacketHandler(mockCtrl)
 				clientSess.EXPECT().getPerspective().Return(protocol.PerspectiveClient)
 				serverSess := NewMockPacketHandler(mockCtrl)
diff --git a/server.go b/server.go
index 29cbed672..1df61d17b 100644
--- a/server.go
+++ b/server.go
@@ -274,12 +274,12 @@ func (s *baseServer) accept(ctx context.Context) (quicSession, error) {
 
 // Close the server
 func (s *baseServer) Close() error {
+	s.sessionHandler.CloseServer()
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
 	if s.closed {
 		return nil
 	}
-	s.sessionHandler.CloseServer()
 	if s.serverError == nil {
 		s.serverError = errors.New("server closed")
 	}