implement connection ID handling for path probe packets (#4935)

conn_id_manager.go

@@ -10,6 +10,8 @@ import (
 	"github.com/quic-go/quic-go/internal/wire"
 )
 
+type pathID int64
+
 type newConnID struct {
 	SequenceNumber      uint64
 	ConnectionID        protocol.ConnectionID
@@ -19,6 +21,9 @@ type newConnID struct {
 type connIDManager struct {
 	queue list.List[newConnID]
 
+	highestProbingID uint64
+	pathProbing      map[pathID]newConnID // initialized lazily
+
 	handshakeComplete         bool
 	activeSequenceNumber      uint64
 	highestRetired            uint64
@@ -76,13 +81,23 @@ func (h *connIDManager) add(f *wire.NewConnectionIDFrame) error {
 	}
 	// If the NEW_CONNECTION_ID frame is reordered, such that its sequence number is smaller than the currently active
 	// connection ID or if it was already retired, send the RETIRE_CONNECTION_ID frame immediately.
-	if f.SequenceNumber < h.activeSequenceNumber || f.SequenceNumber < h.highestRetired {
+	if f.SequenceNumber < max(h.activeSequenceNumber, h.highestProbingID) || f.SequenceNumber < h.highestRetired {
 		h.queueControlFrame(&wire.RetireConnectionIDFrame{
 			SequenceNumber: f.SequenceNumber,
 		})
 		return nil
 	}
 
+	if f.RetirePriorTo != 0 && h.pathProbing != nil {
+		for id, entry := range h.pathProbing {
+			if entry.SequenceNumber < f.RetirePriorTo {
+				h.queueControlFrame(&wire.RetireConnectionIDFrame{
+					SequenceNumber: entry.SequenceNumber,
+				})
+				delete(h.pathProbing, id)
+			}
+		}
+	}
 	// Retire elements in the queue.
 	// Doesn't retire the active connection ID.
 	if f.RetirePriorTo > h.highestRetired {
@@ -225,6 +240,50 @@ func (h *connIDManager) SetHandshakeComplete() {
 	h.handshakeComplete = true
 }
 
+// GetConnIDForPath retrieves a connection ID for a new path (i.e. not the active one).
+// Once a connection ID is allocated for a path, it cannot be used for a different path.
+// When called with the same pathID, it will return the same connection ID,
+// unless the peer requested that this connection ID be retired.
+func (h *connIDManager) GetConnIDForPath(id pathID) (protocol.ConnectionID, bool) {
+	h.assertNotClosed()
+	// if we're using zero-length connection IDs, we don't need to change the connection ID
+	if h.activeConnectionID.Len() == 0 {
+		return protocol.ConnectionID{}, true
+	}
+
+	if h.pathProbing == nil {
+		h.pathProbing = make(map[pathID]newConnID)
+	}
+	entry, ok := h.pathProbing[id]
+	if ok {
+		return entry.ConnectionID, true
+	}
+	if h.queue.Len() == 0 {
+		return protocol.ConnectionID{}, false
+	}
+	front := h.queue.Remove(h.queue.Front())
+	h.pathProbing[id] = front
+	h.highestProbingID = front.SequenceNumber
+	return front.ConnectionID, true
+}
+
+func (h *connIDManager) RetireConnIDForPath(pathID pathID) {
+	h.assertNotClosed()
+	// if we're using zero-length connection IDs, we don't need to change the connection ID
+	if h.activeConnectionID.Len() == 0 {
+		return
+	}
+
+	entry, ok := h.pathProbing[pathID]
+	if !ok {
+		return
+	}
+	h.queueControlFrame(&wire.RetireConnectionIDFrame{
+		SequenceNumber: entry.SequenceNumber,
+	})
+	delete(h.pathProbing, pathID)
+}
+
 // 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.

conn_id_manager_test.go

@@ -209,6 +209,78 @@ func TestConnIDManagerConnIDRotation(t *testing.T) {
 	require.Equal(t, []wire.Frame{&wire.RetireConnectionIDFrame{SequenceNumber: 2}}, frameQueue)
 }
 
+func TestConnIDManagerPathMigration(t *testing.T) {
+	var frameQueue []wire.Frame
+	m := newConnIDManager(
+		protocol.ParseConnectionID([]byte{1, 2, 3, 4}),
+		func(protocol.StatelessResetToken) {},
+		func(protocol.StatelessResetToken) {},
+		func(f wire.Frame) { frameQueue = append(frameQueue, f) },
+	)
+
+	// no connection ID available yet
+	_, ok := m.GetConnIDForPath(1)
+	require.False(t, ok)
+
+	// add a connection ID
+	require.NoError(t, m.Add(&wire.NewConnectionIDFrame{
+		SequenceNumber:      1,
+		ConnectionID:        protocol.ParseConnectionID([]byte{4, 3, 2, 1}),
+		StatelessResetToken: protocol.StatelessResetToken{16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1},
+	}))
+	require.NoError(t, m.Add(&wire.NewConnectionIDFrame{
+		SequenceNumber:      2,
+		ConnectionID:        protocol.ParseConnectionID([]byte{5, 4, 3, 2}),
+		StatelessResetToken: protocol.StatelessResetToken{16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1},
+	}))
+	connID, ok := m.GetConnIDForPath(1)
+	require.True(t, ok)
+	require.Equal(t, protocol.ParseConnectionID([]byte{4, 3, 2, 1}), connID)
+	connID, ok = m.GetConnIDForPath(2)
+	require.True(t, ok)
+	require.Equal(t, protocol.ParseConnectionID([]byte{5, 4, 3, 2}), connID)
+	// asking for the connection for path 1 again returns the same connection ID
+	connID, ok = m.GetConnIDForPath(1)
+	require.True(t, ok)
+	require.Equal(t, protocol.ParseConnectionID([]byte{4, 3, 2, 1}), connID)
+
+	// if the connection ID is retired, the path will use another connection ID
+	require.NoError(t, m.Add(&wire.NewConnectionIDFrame{
+		SequenceNumber:      3,
+		RetirePriorTo:       2,
+		ConnectionID:        protocol.ParseConnectionID([]byte{6, 5, 4, 3}),
+		StatelessResetToken: protocol.StatelessResetToken{16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1},
+	}))
+	require.Len(t, frameQueue, 2)
+	frameQueue = nil
+
+	require.Equal(t, protocol.ParseConnectionID([]byte{6, 5, 4, 3}), m.Get())
+	// the connection ID is not used for new paths
+	_, ok = m.GetConnIDForPath(3)
+	require.False(t, ok)
+
+	// Manually retiring the connection ID does nothing.
+	// Path 1 doesn't have a connection ID anymore.
+	m.RetireConnIDForPath(1)
+	require.Empty(t, frameQueue)
+	_, ok = m.GetConnIDForPath(1)
+	require.False(t, ok)
+
+	// only after a new connection ID is added, it will be used for path 1
+	require.NoError(t, m.Add(&wire.NewConnectionIDFrame{
+		SequenceNumber:      4,
+		ConnectionID:        protocol.ParseConnectionID([]byte{7, 6, 5, 4}),
+		StatelessResetToken: protocol.StatelessResetToken{16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1},
+	}))
+	connID, ok = m.GetConnIDForPath(1)
+	require.True(t, ok)
+	require.Equal(t, protocol.ParseConnectionID([]byte{7, 6, 5, 4}), connID)
+
+	// a RETIRE_CONNECTION_ID frame for path 1 is queued when retiring the connection ID
+	m.RetireConnIDForPath(1)
+	require.Equal(t, []wire.Frame{&wire.RetireConnectionIDFrame{SequenceNumber: 4}}, frameQueue)
+}
+
 func TestConnIDManagerZeroLengthConnectionID(t *testing.T) {
 	m := newConnIDManager(
 		protocol.ConnectionID{},
@@ -222,6 +294,17 @@ func TestConnIDManagerZeroLengthConnectionID(t *testing.T) {
 		require.Equal(t, protocol.ConnectionID{}, m.Get())
 	}
 
+	// for path probing, we don't need to change the connection ID
+	for id := pathID(1); id < 10; id++ {
+		connID, ok := m.GetConnIDForPath(id)
+		require.True(t, ok)
+		require.Equal(t, protocol.ConnectionID{}, connID)
+	}
+	// retiring a connection ID for a path is also a no-op
+	for id := pathID(1); id < 20; id++ {
+		m.RetireConnIDForPath(id)
+	}
+
 	require.ErrorIs(t, m.Add(&wire.NewConnectionIDFrame{
 		SequenceNumber:      1,
 		ConnectionID:        protocol.ConnectionID{},