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.