From 476acc2966ced947ddb2416e43ce614b1b1f3da2 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Mon, 9 May 2016 22:27:55 +0700
Subject: [PATCH] fix PublicHeader of VersionNegotiation and PublicReset
 packets

fixes #71
---
 public_header.go      | 5 ++++-
 public_header_test.go | 4 ++++
 server_test.go        | 2 +-
 3 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/public_header.go b/public_header.go
index 1da588e0..66a72bb4 100644
--- a/public_header.go
+++ b/public_header.go
@@ -50,7 +50,10 @@ func (h *PublicHeader) WritePublicHeader(b *bytes.Buffer) error {
 		utils.WriteUint64(b, uint64(h.ConnectionID))
 	}
 
-	utils.WriteUint48(b, uint64(h.PacketNumber)) // TODO: Send shorter packet number if possible
+	if !h.ResetFlag && !h.VersionFlag {
+		utils.WriteUint48(b, uint64(h.PacketNumber)) // TODO: Send shorter packet number if possible
+	}
+
 	return nil
 }
 
diff --git a/public_header_test.go b/public_header_test.go
index 86d7b150..fc146f36 100644
--- a/public_header_test.go
+++ b/public_header_test.go
@@ -97,6 +97,8 @@ var _ = Describe("Public Header", func() {
 				PacketNumber: 2,
 			}
 			publicHeader.WritePublicHeader(b)
+			// must be the first assertion
+			Expect(b.Len()).To(Equal(1 + 8)) // 1 FlagByte + 8 ConnectionID
 			firstByte, _ := b.ReadByte()
 			Expect(firstByte & 0x01).To(Equal(uint8(1)))
 		})
@@ -109,6 +111,8 @@ var _ = Describe("Public Header", func() {
 				PacketNumber: 2,
 			}
 			publicHeader.WritePublicHeader(b)
+			// must be the first assertion
+			Expect(b.Len()).To(Equal(1 + 8)) // 1 FlagByte + 8 ConnectionID
 			firstByte, _ := b.ReadByte()
 			Expect((firstByte & 0x02) >> 1).To(Equal(uint8(1)))
 		})
diff --git a/server_test.go b/server_test.go
index 5ebd521a..426fbdbf 100644
--- a/server_test.go
+++ b/server_test.go
@@ -47,7 +47,7 @@ var _ = Describe("Server", func() {
 
 		It("composes version negotiation packets", func() {
 			expected := append(
-				[]byte{0x3d, 0x1, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1, 0x0, 0x0, 0x0, 0x0, 0x0},
+				[]byte{0x3d, 0x1, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0},
 				protocol.SupportedVersionsAsTags...,
 			)
 			Expect(composeVersionNegotiation(1)).To(Equal(expected))