From 0d30e43c3248264f6766c41ebf27fdbc9a2305cb Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Tue, 27 Nov 2018 08:29:31 +0700
Subject: [PATCH] implement the new format of the Retry packet

---
 internal/wire/extended_header.go      | 30 ++++++++++++---------------
 internal/wire/extended_header_test.go | 16 +++++++-------
 internal/wire/header.go               |  6 +-----
 internal/wire/header_test.go          |  3 +--
 4 files changed, 23 insertions(+), 32 deletions(-)

diff --git a/internal/wire/extended_header.go b/internal/wire/extended_header.go
index 14ec7ea8..a08dd487 100644
--- a/internal/wire/extended_header.go
+++ b/internal/wire/extended_header.go
@@ -2,7 +2,6 @@ package wire
 
 import (
 	"bytes"
-	"crypto/rand"
 	"errors"
 	"fmt"
 	"io"
@@ -94,9 +93,16 @@ func (h *ExtendedHeader) writeLongHeader(b *bytes.Buffer, v protocol.VersionNumb
 		packetType = 0x3
 	}
 	firstByte := 0xc0 | packetType<<4
-	if h.Type != protocol.PacketTypeRetry { // Retry packets don't have a packet number
+	if h.Type == protocol.PacketTypeRetry {
+		odcil, err := encodeSingleConnIDLen(h.OrigDestConnectionID)
+		if err != nil {
+			return err
+		}
+		firstByte |= odcil
+	} else { // Retry packets don't have a packet number
 		firstByte |= uint8(h.PacketNumberLen - 1)
 	}
+
 	b.WriteByte(firstByte)
 	utils.BigEndian.WriteUint32(b, uint32(h.Version))
 	connIDLen, err := encodeConnIDLen(h.DestConnectionID, h.SrcConnectionID)
@@ -107,24 +113,14 @@ func (h *ExtendedHeader) writeLongHeader(b *bytes.Buffer, v protocol.VersionNumb
 	b.Write(h.DestConnectionID.Bytes())
 	b.Write(h.SrcConnectionID.Bytes())
 
-	if h.Type == protocol.PacketTypeInitial {
-		utils.WriteVarInt(b, uint64(len(h.Token)))
-		b.Write(h.Token)
-	}
-
-	if h.Type == protocol.PacketTypeRetry {
-		odcil, err := encodeSingleConnIDLen(h.OrigDestConnectionID)
-		if err != nil {
-			return err
-		}
-		// randomize the first 4 bits
-		odcilByte := make([]byte, 1)
-		_, _ = rand.Read(odcilByte) // it's safe to ignore the error here
-		odcilByte[0] = (odcilByte[0] & 0xf0) | odcil
-		b.Write(odcilByte)
+	switch h.Type {
+	case protocol.PacketTypeRetry:
 		b.Write(h.OrigDestConnectionID.Bytes())
 		b.Write(h.Token)
 		return nil
+	case protocol.PacketTypeInitial:
+		utils.WriteVarInt(b, uint64(len(h.Token)))
+		b.Write(h.Token)
 	}
 
 	utils.WriteVarInt(b, uint64(h.Length))
diff --git a/internal/wire/extended_header_test.go b/internal/wire/extended_header_test.go
index 8c276d82..e669fd09 100644
--- a/internal/wire/extended_header_test.go
+++ b/internal/wire/extended_header_test.go
@@ -120,14 +120,14 @@ var _ = Describe("Header", func() {
 					Token:                token,
 					OrigDestConnectionID: protocol.ConnectionID{1, 2, 3, 4, 5, 6, 7, 8, 9},
 				}}).Write(buf, versionIETFHeader)).To(Succeed())
-				Expect(buf.Bytes()[:6]).To(Equal([]byte{
-					0xc0 | 0x3<<4,
-					0x1, 0x2, 0x3, 0x4, // version number
-					0x0, // connection ID lengths))
-				}))
-				Expect(buf.Bytes()[6] & 0xf).To(Equal(uint8(6)))
-				Expect(buf.Bytes()[7 : 7+9]).To(Equal([]byte{1, 2, 3, 4, 5, 6, 7, 8, 9})) // Orig Dest Connection ID
-				Expect(buf.Bytes()[7+9:]).To(Equal(token))
+				expected := []byte{
+					0xc0 | 0x3<<4 | 9 - 3, /* orig dest connection ID length */
+					0x1, 0x2, 0x3, 0x4,    // version number
+					0x0,                       // connection ID lengths))
+					1, 2, 3, 4, 5, 6, 7, 8, 9, // Orig Dest Connection ID
+				}
+				expected = append(expected, token...)
+				Expect(buf.Bytes()).To(Equal(expected))
 			})
 
 			It("refuses to write a Retry packet with an invalid Orig Destination Connection ID length", func() {
diff --git a/internal/wire/header.go b/internal/wire/header.go
index 547b4e28..c40d40b2 100644
--- a/internal/wire/header.go
+++ b/internal/wire/header.go
@@ -117,11 +117,7 @@ func (h *Header) parseLongHeader(b *bytes.Reader) error {
 	}
 
 	if h.Type == protocol.PacketTypeRetry {
-		odcilByte, err := b.ReadByte()
-		if err != nil {
-			return err
-		}
-		odcil := decodeSingleConnIDLen(odcilByte & 0xf)
+		odcil := decodeSingleConnIDLen(h.typeByte & 0xf)
 		h.OrigDestConnectionID, err = protocol.ReadConnectionID(b, odcil)
 		if err != nil {
 			return err
diff --git a/internal/wire/header_test.go b/internal/wire/header_test.go
index 383355c0..ff56b74a 100644
--- a/internal/wire/header_test.go
+++ b/internal/wire/header_test.go
@@ -171,10 +171,9 @@ var _ = Describe("Header Parsing", func() {
 		})
 
 		It("parses a Retry packet", func() {
-			data := []byte{0xc0 ^ 0x3<<4}
+			data := []byte{0xc0 | 0x3<<4 | (10 - 3) /* connection ID length */}
 			data = appendVersion(data, versionIETFFrames)
 			data = append(data, 0x0)                                      // connection ID lengths
-			data = append(data, 0x97)                                     // Orig Destination Connection ID length
 			data = append(data, []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}...) // source connection ID
 			data = append(data, []byte{'f', 'o', 'o', 'b', 'a', 'r'}...)  // token
 			b := bytes.NewReader(data)