From b52455c3f7483a9e45e149bfe46bade1d0720218 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Tue, 19 Apr 2016 10:20:48 +0700
Subject: [PATCH] add method to Write ConnectionClose frames

---
 frames/connection_close_frame.go      | 27 +++++++++---
 frames/connection_close_frame_test.go | 59 ++++++++++++++++++++++++++-
 2 files changed, 79 insertions(+), 7 deletions(-)

diff --git a/frames/connection_close_frame.go b/frames/connection_close_frame.go
index 9d484714..107b053d 100644
--- a/frames/connection_close_frame.go
+++ b/frames/connection_close_frame.go
@@ -3,6 +3,7 @@ package frames
 import (
 	"bytes"
 	"errors"
+	"math"
 
 	"github.com/lucas-clemente/quic-go/utils"
 )
@@ -13,12 +14,7 @@ type ConnectionCloseFrame struct {
 	ReasonPhrase string
 }
 
-// Write writes an ACK frame.
-func (f *ConnectionCloseFrame) Write(b *bytes.Buffer) error {
-	return errors.New("ConnectionCloseFrame: Write not yet implemented")
-}
-
-// ParseConnectionCloseFrame reads an ACK frame
+// ParseConnectionCloseFrame reads a CONNECTION_CLOSE frame
 func ParseConnectionCloseFrame(r *bytes.Reader) (*ConnectionCloseFrame, error) {
 	frame := &ConnectionCloseFrame{}
 
@@ -49,3 +45,22 @@ func ParseConnectionCloseFrame(r *bytes.Reader) (*ConnectionCloseFrame, error) {
 
 	return frame, nil
 }
+
+// Write writes an CONNECTION_CLOSE frame.
+func (f *ConnectionCloseFrame) Write(b *bytes.Buffer) error {
+	b.WriteByte(0x02)
+	utils.WriteUint32(b, f.ErrorCode)
+
+	if len(f.ReasonPhrase) > math.MaxUint16 {
+		return errors.New("ConnectionFrame: ReasonPhrase too long")
+	}
+
+	reasonPhraseLen := uint16(len(f.ReasonPhrase))
+	utils.WriteUint16(b, reasonPhraseLen)
+
+	for i := 0; i < int(reasonPhraseLen); i++ {
+		b.WriteByte(uint8(f.ReasonPhrase[i]))
+	}
+
+	return nil
+}
diff --git a/frames/connection_close_frame_test.go b/frames/connection_close_frame_test.go
index 5e50e04c..f4b5ddbd 100644
--- a/frames/connection_close_frame_test.go
+++ b/frames/connection_close_frame_test.go
@@ -7,7 +7,7 @@ import (
 	. "github.com/onsi/gomega"
 )
 
-var _ = Describe("AckFrame", func() {
+var _ = Describe("ConnectionCloseFrame", func() {
 	Context("when parsing", func() {
 		It("accepts sample frame", func() {
 			b := bytes.NewReader([]byte{0x40, 0x19, 0x00, 0x00, 0x00, 0x1B, 0x00, 0x4e, 0x6f, 0x20, 0x72, 0x65, 0x63, 0x65, 0x6e, 0x74, 0x20, 0x6e, 0x65, 0x74, 0x77, 0x6f, 0x72, 0x6b, 0x20, 0x61, 0x63, 0x74, 0x69, 0x76, 0x69, 0x74, 0x79, 0x2e})
@@ -27,4 +27,61 @@ var _ = Describe("AckFrame", func() {
 			Expect(b.Len()).To(Equal(0))
 		})
 	})
+
+	Context("when writing", func() {
+		It("writes a frame without a ReasonPhrase", func() {
+			b := &bytes.Buffer{}
+			frame := &ConnectionCloseFrame{
+				ErrorCode: 0xDEADBEEF,
+			}
+			err := frame.Write(b)
+			Expect(err).ToNot(HaveOccurred())
+			Expect(b.Len()).To(Equal(1 + 2 + 4))
+			Expect(b.Bytes()).To(Equal([]byte{0x02, 0xEF, 0xBE, 0xAD, 0xDE, 0x00, 0x00}))
+		})
+
+		It("writes a frame with a ReasonPhrase", func() {
+			b := &bytes.Buffer{}
+			frame := &ConnectionCloseFrame{
+				ErrorCode:    0xDEADBEEF,
+				ReasonPhrase: "foobar",
+			}
+			err := frame.Write(b)
+			Expect(err).ToNot(HaveOccurred())
+			Expect(b.Len()).To(Equal(1 + 2 + 4 + len(frame.ReasonPhrase)))
+			Expect(b.Bytes()[:5]).To(Equal([]byte{0x02, 0xEF, 0xBE, 0xAD, 0xDE}))
+			Expect(b.Bytes()[5:7]).To(Equal([]byte{0x06, 0x00}))
+			Expect(b.Bytes()[7:]).To(Equal([]byte{'f', 'o', 'o', 'b', 'a', 'r'}))
+		})
+
+		It("rejects ReasonPhrases that are too long", func() {
+			b := &bytes.Buffer{}
+
+			var reasonPhrase string
+			for i := 0; i < int(0xFFFF+0x11); i++ {
+				reasonPhrase += "a"
+			}
+
+			frame := &ConnectionCloseFrame{
+				ErrorCode:    0xDEADBEEF,
+				ReasonPhrase: reasonPhrase,
+			}
+			err := frame.Write(b)
+			Expect(err).To(HaveOccurred())
+		})
+	})
+
+	It("is self-consistent", func() {
+		b := &bytes.Buffer{}
+		frame := &ConnectionCloseFrame{
+			ErrorCode:    0xDEADBEEF,
+			ReasonPhrase: "Lorem ipsum dolor sit amet.",
+		}
+		err := frame.Write(b)
+		Expect(err).ToNot(HaveOccurred())
+		readframe, err := ParseConnectionCloseFrame(bytes.NewReader(b.Bytes()))
+		Expect(err).ToNot(HaveOccurred())
+		Expect(readframe.ErrorCode).To(Equal(frame.ErrorCode))
+		Expect(readframe.ReasonPhrase).To(Equal(frame.ReasonPhrase))
+	})
 })