From a321f9faa65899d36b520cdfad444bbde6bde633 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Sun, 15 Sep 2019 14:33:09 +0700
Subject: [PATCH] save connection IDs sent in NEW_CONNECTION_ID frames

---
 session.go      | 9 ++++++++-
 session_test.go | 8 ++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/session.go b/session.go
index c4cd22fe..09af228d 100644
--- a/session.go
+++ b/session.go
@@ -116,7 +116,8 @@ type session struct {
 	conn      connection
 	sendQueue *sendQueue
 
-	streamsMap streamManager
+	streamsMap    streamManager
+	connIDManager *connIDManager
 
 	rttStats *congestion.RTTStats
 
@@ -332,6 +333,7 @@ func (s *session) preSetup() {
 	s.sendQueue = newSendQueue(s.conn)
 	s.retransmissionQueue = newRetransmissionQueue(s.version)
 	s.frameParser = wire.NewFrameParser(s.version)
+	s.connIDManager = newConnIDManager(s.queueControlFrame)
 	s.rttStats = &congestion.RTTStats{}
 	s.receivedPacketHandler = ackhandler.NewReceivedPacketHandler(s.rttStats, s.logger, s.version)
 	s.connFlowController = flowcontrol.NewConnectionFlowController(
@@ -774,6 +776,7 @@ func (s *session) handleFrame(f wire.Frame, pn protocol.PacketNumber, encLevel p
 	case *wire.NewTokenFrame:
 		err = s.handleNewTokenFrame(frame)
 	case *wire.NewConnectionIDFrame:
+		err = s.handleNewConnectionIDFrame(frame)
 	case *wire.RetireConnectionIDFrame:
 		// since we don't send new connection IDs, we don't expect retirements
 		err = errors.New("unexpected RETIRE_CONNECTION_ID frame")
@@ -888,6 +891,10 @@ func (s *session) handleNewTokenFrame(frame *wire.NewTokenFrame) error {
 	return nil
 }
 
+func (s *session) handleNewConnectionIDFrame(f *wire.NewConnectionIDFrame) error {
+	return s.connIDManager.Add(f)
+}
+
 func (s *session) handleAckFrame(frame *wire.AckFrame, pn protocol.PacketNumber, encLevel protocol.EncryptionLevel) error {
 	if err := s.sentPacketHandler.ReceivedAck(frame, pn, encLevel, s.lastPacketReceivedTime); err != nil {
 		return err
diff --git a/session_test.go b/session_test.go
index d36c478f..2d216a3d 100644
--- a/session_test.go
+++ b/session_test.go
@@ -302,6 +302,14 @@ var _ = Describe("Session", func() {
 			})
 		})
 
+		It("handles NEW_CONNECTION_ID frames", func() {
+			Expect(sess.handleFrame(&wire.NewConnectionIDFrame{
+				SequenceNumber: 10,
+				ConnectionID:   protocol.ConnectionID{1, 2, 3, 4},
+			}, 1, protocol.Encryption1RTT)).To(Succeed())
+			Expect(sess.connIDManager.queue.Front().Value.ConnectionID).To(Equal(protocol.ConnectionID{1, 2, 3, 4}))
+		})
+
 		It("handles PING frames", func() {
 			err := sess.handleFrame(&wire.PingFrame{}, 0, protocol.EncryptionUnspecified)
 			Expect(err).NotTo(HaveOccurred())