From 4bf50901e8f453a50020525626c1e913c127a124 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Tue, 30 Aug 2022 18:58:37 +0300
Subject: [PATCH] add a function to determine the length of a short header

---
 internal/wire/short_header.go      | 4 ++++
 internal/wire/short_header_test.go | 5 +++++
 qlog/qlog.go                       | 3 +--
 3 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/internal/wire/short_header.go b/internal/wire/short_header.go
index 9639b5d4..1ddb9be9 100644
--- a/internal/wire/short_header.go
+++ b/internal/wire/short_header.go
@@ -50,6 +50,10 @@ func ParseShortHeader(data []byte, connIDLen int) (length int, _ protocol.Packet
 	return 1 + connIDLen + int(pnLen), pn, pnLen, kp, err
 }
 
+func ShortHeaderLen(dest protocol.ConnectionID, pnLen protocol.PacketNumberLen) protocol.ByteCount {
+	return 1 + protocol.ByteCount(dest.Len()) + protocol.ByteCount(pnLen)
+}
+
 func LogShortHeader(logger utils.Logger, dest protocol.ConnectionID, pn protocol.PacketNumber, pnLen protocol.PacketNumberLen, kp protocol.KeyPhaseBit) {
 	logger.Debugf("\tShort Header{DestConnectionID: %s, PacketNumber: %d, PacketNumberLen: %d, KeyPhase: %s}", dest, pn, pnLen, kp)
 }
diff --git a/internal/wire/short_header_test.go b/internal/wire/short_header_test.go
index 0977fc8b..a475a288 100644
--- a/internal/wire/short_header_test.go
+++ b/internal/wire/short_header_test.go
@@ -70,6 +70,11 @@ var _ = Describe("Short Header", func() {
 		})
 	})
 
+	It("determines the length", func() {
+		Expect(ShortHeaderLen(protocol.ParseConnectionID([]byte{1, 2, 3, 4}), protocol.PacketNumberLen3)).To(BeEquivalentTo(8))
+		Expect(ShortHeaderLen(protocol.ParseConnectionID([]byte{}), protocol.PacketNumberLen1)).To(BeEquivalentTo(2))
+	})
+
 	Context("logging", func() {
 		var (
 			buf    *bytes.Buffer
diff --git a/qlog/qlog.go b/qlog/qlog.go
index 3db0727e..0c4bc117 100644
--- a/qlog/qlog.go
+++ b/qlog/qlog.go
@@ -326,12 +326,11 @@ func (t *connectionTracer) ReceivedShortHeaderPacket(hdr *logging.ShortHeader, p
 		fs[i] = frame{Frame: f}
 	}
 	header := *transformShortHeader(hdr)
-	hdrLen := 1 + hdr.DestConnectionID.Len() + int(hdr.PacketNumberLen)
 	t.mutex.Lock()
 	t.recordEvent(time.Now(), &eventPacketReceived{
 		Header:        header,
 		Length:        packetSize,
-		PayloadLength: packetSize - protocol.ByteCount(hdrLen),
+		PayloadLength: packetSize - wire.ShortHeaderLen(hdr.DestConnectionID, hdr.PacketNumberLen),
 		Frames:        fs,
 	})
 	t.mutex.Unlock()