From a58bcb747b0360e6a6489cf626432cd44c3d0d98 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Mon, 27 Jan 2020 12:59:16 +0700
Subject: [PATCH] qlog received Retry packets

---
 qlog/event.go              | 13 +++++++++++++
 qlog/packet_header.go      |  9 +++++++--
 qlog/packet_header_test.go |  2 +-
 qlog/qlog.go               | 14 ++++++++++++--
 qlog/qlog_test.go          | 21 +++++++++++++++++++++
 session.go                 |  3 +++
 6 files changed, 57 insertions(+), 5 deletions(-)

diff --git a/qlog/event.go b/qlog/event.go
index 3f481e7d8..d63c323a2 100644
--- a/qlog/event.go
+++ b/qlog/event.go
@@ -83,3 +83,16 @@ func (e eventPacketReceived) MarshalJSONObject(enc *gojay.Encoder) {
 	enc.BoolKeyOmitEmpty("is_coalesced", e.IsCoalesced)
 	enc.StringKeyOmitEmpty("trigger", e.Trigger)
 }
+
+type eventRetryReceived struct {
+	Header packetHeader
+}
+
+func (e eventRetryReceived) Category() category { return categoryTransport }
+func (e eventRetryReceived) Name() string       { return "packet_received" }
+func (e eventRetryReceived) IsNil() bool        { return false }
+
+func (e eventRetryReceived) MarshalJSONObject(enc *gojay.Encoder) {
+	enc.StringKey("packet_type", packetTypeRetry.String())
+	enc.ObjectKey("header", e.Header)
+}
diff --git a/qlog/packet_header.go b/qlog/packet_header.go
index e0c3139a8..f749143e0 100644
--- a/qlog/packet_header.go
+++ b/qlog/packet_header.go
@@ -28,9 +28,8 @@ func getPacketType(hdr *wire.ExtendedHeader) packetType {
 	}
 }
 
-func transformHeader(hdr *wire.ExtendedHeader) *packetHeader {
+func transformHeader(hdr *wire.Header) *packetHeader {
 	return &packetHeader{
-		PacketNumber:     hdr.PacketNumber,
 		PayloadLength:    hdr.Length,
 		SrcConnectionID:  hdr.SrcConnectionID,
 		DestConnectionID: hdr.DestConnectionID,
@@ -38,6 +37,12 @@ func transformHeader(hdr *wire.ExtendedHeader) *packetHeader {
 	}
 }
 
+func transformExtendedHeader(hdr *wire.ExtendedHeader) *packetHeader {
+	h := transformHeader(&hdr.Header)
+	h.PacketNumber = hdr.PacketNumber
+	return h
+}
+
 type packetHeader struct {
 	PacketNumber  protocol.PacketNumber
 	PacketSize    protocol.ByteCount
diff --git a/qlog/packet_header_test.go b/qlog/packet_header_test.go
index 247cf2a01..a91aa48a4 100644
--- a/qlog/packet_header_test.go
+++ b/qlog/packet_header_test.go
@@ -72,7 +72,7 @@ var _ = Describe("Packet Header", func() {
 		check := func(hdr *wire.ExtendedHeader, expected map[string]interface{}) {
 			buf := &bytes.Buffer{}
 			enc := gojay.NewEncoder(buf)
-			ExpectWithOffset(1, enc.Encode(transformHeader(hdr))).To(Succeed())
+			ExpectWithOffset(1, enc.Encode(transformExtendedHeader(hdr))).To(Succeed())
 			data := buf.Bytes()
 			ExpectWithOffset(1, json.Valid(data)).To(BeTrue())
 			checkEncoding(data, expected)
diff --git a/qlog/qlog.go b/qlog/qlog.go
index 5723e6988..982ee5af7 100644
--- a/qlog/qlog.go
+++ b/qlog/qlog.go
@@ -13,6 +13,7 @@ import (
 type Tracer interface {
 	Export() error
 	SentPacket(time.Time, *wire.ExtendedHeader, *wire.AckFrame, []wire.Frame)
+	ReceivedRetry(time.Time, *wire.Header)
 	ReceivedPacket(time.Time, *wire.ExtendedHeader, []wire.Frame)
 }
 
@@ -71,7 +72,7 @@ func (t *tracer) SentPacket(time time.Time, hdr *wire.ExtendedHeader, ack *wire.
 		Time: time,
 		eventDetails: eventPacketSent{
 			PacketType: getPacketType(hdr),
-			Header:     *transformHeader(hdr),
+			Header:     *transformExtendedHeader(hdr),
 			Frames:     fs,
 		},
 	})
@@ -86,8 +87,17 @@ func (t *tracer) ReceivedPacket(time time.Time, hdr *wire.ExtendedHeader, frames
 		Time: time,
 		eventDetails: eventPacketReceived{
 			PacketType: getPacketType(hdr),
-			Header:     *transformHeader(hdr),
+			Header:     *transformExtendedHeader(hdr),
 			Frames:     fs,
 		},
 	})
 }
+
+func (t *tracer) ReceivedRetry(time time.Time, hdr *wire.Header) {
+	t.events = append(t.events, event{
+		Time: time,
+		eventDetails: eventRetryReceived{
+			Header: *transformHeader(hdr),
+		},
+	})
+}
diff --git a/qlog/qlog_test.go b/qlog/qlog_test.go
index 16832cb22..a23291624 100644
--- a/qlog/qlog_test.go
+++ b/qlog/qlog_test.go
@@ -159,5 +159,26 @@ var _ = Describe("Tracer", func() {
 			Expect(ev).To(HaveKey("frames"))
 			Expect(ev["frames"].([]interface{})).To(HaveLen(2))
 		})
+
+		It("records a received Retry packet", func() {
+			now := time.Now()
+			tracer.ReceivedRetry(
+				now,
+				&wire.Header{
+					IsLongHeader:     true,
+					Type:             protocol.PacketTypeRetry,
+					DestConnectionID: protocol.ConnectionID{1, 2, 3, 4, 5, 6, 7, 8},
+					SrcConnectionID:  protocol.ConnectionID{4, 3, 2, 1},
+					Version:          protocol.VersionTLS,
+				},
+			)
+			t, category, eventName, ev := exportAndParse()
+			Expect(t).To(BeTemporally("~", now, time.Millisecond))
+			Expect(category).To(Equal("transport"))
+			Expect(eventName).To(Equal("packet_received"))
+			Expect(ev).To(HaveKeyWithValue("packet_type", "retry"))
+			Expect(ev).To(HaveKey("header"))
+			Expect(ev).ToNot(HaveKey("frames"))
+		})
 	})
 })
diff --git a/session.go b/session.go
index be4c37035..e0bdcdfa3 100644
--- a/session.go
+++ b/session.go
@@ -791,6 +791,9 @@ func (s *session) handleRetryPacket(hdr *wire.Header, data []byte) bool /* was t
 	}
 	s.logger.Debugf("<- Received Retry")
 	s.logger.Debugf("Switching destination connection ID to: %s", hdr.SrcConnectionID)
+	if s.qlogger != nil {
+		s.qlogger.ReceivedRetry(time.Now(), hdr)
+	}
 	s.origDestConnID = s.handshakeDestConnID
 	newDestConnID := hdr.SrcConnectionID
 	s.receivedRetry = true