Repository: quic-go/quic-go
Pull request #4852 (2025-01-08 06:49:37)

fix memory leak on connection ID rotation when closing connection
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.

Thanks to @MarcoPolo for debugging and reporting this issue.

Unfortunately, I don't see a good way to test this for now, since a fresh connection (as we use in tests), comes with a connection ID that doesn't have a stateless reset token. #4850 might make this easier to test at some point.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/quic-go/quic-go/pull/4852?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 83.55%. Comparing base [(`c42f845`)](https://app.codecov.io/gh/quic-go/quic-go/commit/c42f8456ab2f493d9ea6af982115e204946d2045?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go) to head [(`080d80a`)](https://app.codecov.io/gh/quic-go/quic-go/commit/080d80a5461064c786d4ba7fd56f033a02fd8d6c?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go).

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #4852      +/-   ##
==========================================
- Coverage   83.73%   83.55%   -0.18%     
==========================================
  Files         149      146       -3     
  Lines       16095    15751     -344     
==========================================
- Hits        13476    13160     -316     
+ Misses       2089     2080       -9     
+ Partials      530      511      -19     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/quic-go/quic-go/pull/4852?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go).

marten-seemann: > This doesn't fix the issue. Did you mean to add a `defer`?

Yes 🤦🏻‍♂️

Added in aeea67c4c.

I still need to think about panicking. Not really a fan of that, but I see your point.
MarcoPolo: I think it's also okay if we didn't panic and change the return signatures instead to also return an error. Seems like a more annoying change.
p-shahi: @marten-seemann - we're waiting on this PR to land w/ a new quic-go patch release. Do you know the eta?

Patch:
diff --git a/conn_id_manager.go b/conn_id_manager.go
index 0e65231c..4030913d 100644
--- a/conn_id_manager.go
+++ b/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")
+	}
+}
diff --git a/connection.go b/connection.go
index ffb151ed..1279d1ce 100644
--- a/connection.go
+++ b/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)
 	}


Test patch:
diff --git a/conn_id_manager_test.go b/conn_id_manager_test.go
index 193d74c2..e03d3d63 100644
--- a/conn_id_manager_test.go
+++ b/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{}) })
 }