fix memory leak on connection ID rotation when closing connection (#4852)

* fix memory leak on connection ID rotation during CONNECTION_CLOSE

In rare instances, the connection ID manager might switch to a new
connection ID when sending the packet containing the CONNECTION_CLOSE
frame. The connection ID manager removes the active stateless reset
token from the packet handler map when it is closed, so we need to make
sure that this happens last, otherwise the packet handler will keep a
pointer to the closed connection indefinitely.

* defer

* panic on use of connIDManager after Close

conn_id_manager.go

@@ -35,6 +35,8 @@ type connIDManager struct {
 	addStatelessResetToken    func(protocol.StatelessResetToken)
 	removeStatelessResetToken func(protocol.StatelessResetToken)
 	queueControlFrame         func(wire.Frame)
+
+	closed bool
 }
 
 func newConnIDManager(
@@ -148,6 +150,7 @@ func (h *connIDManager) addConnectionID(seq uint64, connID protocol.ConnectionID
 }
 
 func (h *connIDManager) updateConnectionID() {
+	h.assertNotClosed()
 	h.queueControlFrame(&wire.RetireConnectionIDFrame{
 		SequenceNumber: h.activeSequenceNumber,
 	})
@@ -166,6 +169,7 @@ func (h *connIDManager) updateConnectionID() {
 }
 
 func (h *connIDManager) Close() {
+	h.closed = true
 	if h.activeStatelessResetToken != nil {
 		h.removeStatelessResetToken(*h.activeStatelessResetToken)
 	}
@@ -182,6 +186,7 @@ func (h *connIDManager) ChangeInitialConnID(newConnID protocol.ConnectionID) {
 
 // is called when the server provides a stateless reset token in the transport parameters
 func (h *connIDManager) SetStatelessResetToken(token protocol.StatelessResetToken) {
+	h.assertNotClosed()
 	if h.activeSequenceNumber != 0 {
 		panic("expected first connection ID to have sequence number 0")
 	}
@@ -209,6 +214,7 @@ func (h *connIDManager) shouldUpdateConnID() bool {
 }
 
 func (h *connIDManager) Get() protocol.ConnectionID {
+	h.assertNotClosed()
 	if h.shouldUpdateConnID() {
 		h.updateConnectionID()
 	}
@@ -218,3 +224,13 @@ func (h *connIDManager) Get() protocol.ConnectionID {
 func (h *connIDManager) SetHandshakeComplete() {
 	h.handshakeComplete = true
 }
+
+// Using the connIDManager after it has been closed can have disastrous effects:
+// If the connection ID is rotated, a new entry would be inserted into the packet handler map,
+// leading to a memory leak of the connection struct.
+// See https://github.com/quic-go/quic-go/pull/4852 for more details.
+func (h *connIDManager) assertNotClosed() {
+	if h.closed {
+		panic("connection ID manager is closed")
+	}
+}

conn_id_manager_test.go

@@ -242,4 +242,7 @@ func TestConnIDManagerClose(t *testing.T) {
 	require.Empty(t, removedTokens)
 	m.Close()
 	require.Equal(t, []protocol.StatelessResetToken{{16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1}}, removedTokens)
+
+	require.Panics(t, func() { m.Get() })
+	require.Panics(t, func() { m.SetStatelessResetToken(protocol.StatelessResetToken{}) })
 }

connection.go

@@ -1667,11 +1667,16 @@ func (s *connection) handleCloseError(closeErr *closeError) {
 	}
 
 	s.streamsMap.CloseWithError(e)
-	s.connIDManager.Close()
 	if s.datagramQueue != nil {
 		s.datagramQueue.CloseWithError(e)
 	}
 
+	// In rare instances, the connection ID manager might switch to a new connection ID
+	// when sending the CONNECTION_CLOSE frame.
+	// The connection ID manager removes the active stateless reset token from the packet
+	// handler map when it is closed, so we need to make sure that this happens last.
+	defer s.connIDManager.Close()
+
 	if s.tracer != nil && s.tracer.ClosedConnection != nil && !errors.As(e, &recreateErr) {
 		s.tracer.ClosedConnection(e)
 	}