From 7676e455c5f5cdd0ce1bbef9070df53c937b11d0 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Sun, 7 Mar 2021 19:19:47 +0800
Subject: [PATCH] add a function to tell if a packet is a 0-RTT packet

---
 fuzzing/header/fuzz.go       |  4 ++++
 internal/wire/header.go      | 16 ++++++++++++++++
 internal/wire/header_test.go | 18 ++++++++++++++++++
 3 files changed, 38 insertions(+)

diff --git a/fuzzing/header/fuzz.go b/fuzzing/header/fuzz.go
index 9b444349..46bdfa89 100644
--- a/fuzzing/header/fuzz.go
+++ b/fuzzing/header/fuzz.go
@@ -29,6 +29,7 @@ func Fuzz(data []byte) int {
 	if err != nil {
 		return 0
 	}
+	is0RTTPacket := wire.Is0RTTPacket(data)
 	hdr, _, _, err := wire.ParsePacket(data, connIDLen)
 	if err != nil {
 		return 0
@@ -36,6 +37,9 @@ func Fuzz(data []byte) int {
 	if !hdr.DestConnectionID.Equal(connID) {
 		panic(fmt.Sprintf("Expected connection IDs to match: %s vs %s", hdr.DestConnectionID, connID))
 	}
+	if (hdr.Type == protocol.PacketType0RTT) != is0RTTPacket {
+		panic("inconsistent 0-RTT packet detection")
+	}
 
 	var extHdr *wire.ExtendedHeader
 	// Parse the extended header, if this is not a Retry packet.
diff --git a/internal/wire/header.go b/internal/wire/header.go
index ddd7bb92..07ca9f05 100644
--- a/internal/wire/header.go
+++ b/internal/wire/header.go
@@ -2,6 +2,7 @@ package wire
 
 import (
 	"bytes"
+	"encoding/binary"
 	"errors"
 	"fmt"
 	"io"
@@ -43,6 +44,21 @@ func IsVersionNegotiationPacket(b []byte) bool {
 	return b[0]&0x80 > 0 && b[1] == 0 && b[2] == 0 && b[3] == 0 && b[4] == 0
 }
 
+// Is0RTTPacket says if this is a 0-RTT packet.
+// A packet sent with a version we don't understand can never be a 0-RTT packet.
+func Is0RTTPacket(b []byte) bool {
+	if len(b) < 5 {
+		return false
+	}
+	if b[0]&0x80 == 0 {
+		return false
+	}
+	if !protocol.IsSupportedVersion(protocol.SupportedVersions, protocol.VersionNumber(binary.BigEndian.Uint32(b[1:5]))) {
+		return false
+	}
+	return b[0]&0x30>>4 == 0x1
+}
+
 var ErrUnsupportedVersion = errors.New("unsupported version")
 
 // The Header is the version independent part of the header
diff --git a/internal/wire/header_test.go b/internal/wire/header_test.go
index 5a90947a..782e292d 100644
--- a/internal/wire/header_test.go
+++ b/internal/wire/header_test.go
@@ -93,6 +93,24 @@ var _ = Describe("Header Parsing", func() {
 		})
 	})
 
+	Context("identifying 0-RTT packets", func() {
+		var zeroRTTHeader []byte
+
+		BeforeEach(func() {
+			zeroRTTHeader = make([]byte, 5)
+			zeroRTTHeader[0] = 0x80 | 0x1<<4
+			binary.BigEndian.PutUint32(zeroRTTHeader[1:], uint32(versionIETFFrames))
+		})
+
+		It("recognizes 0-RTT packets", func() {
+			Expect(Is0RTTPacket(zeroRTTHeader[:4])).To(BeFalse())                           // too short
+			Expect(Is0RTTPacket([]byte{zeroRTTHeader[0], 1, 2, 3, 4})).To(BeFalse())        // unknown version
+			Expect(Is0RTTPacket([]byte{zeroRTTHeader[0] | 0x80, 1, 2, 3, 4})).To(BeFalse()) // short header
+			Expect(Is0RTTPacket(zeroRTTHeader)).To(BeTrue())
+			Expect(Is0RTTPacket(append(zeroRTTHeader, []byte("foobar")...))).To(BeTrue())
+		})
+	})
+
 	Context("Identifying Version Negotiation Packets", func() {
 		It("identifies version negotiation packets", func() {
 			Expect(IsVersionNegotiationPacket([]byte{0x80 | 0x56, 0, 0, 0, 0})).To(BeTrue())