From bdf53744acc971c4fcc5dce6a30d864b14ba7982 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Wed, 11 May 2016 10:09:29 +0700
Subject: [PATCH] code cleanup in StreamFrame writing

---
 frames/stream_frame.go      | 42 +++++++++++++------------------------
 frames/stream_frame_test.go | 13 +++++++++++-
 2 files changed, 26 insertions(+), 29 deletions(-)

diff --git a/frames/stream_frame.go b/frames/stream_frame.go
index 55f27796..364fd56c 100644
--- a/frames/stream_frame.go
+++ b/frames/stream_frame.go
@@ -2,6 +2,7 @@ package frames
 
 import (
 	"bytes"
+	"errors"
 	"io/ioutil"
 
 	"github.com/lucas-clemente/quic-go/protocol"
@@ -17,6 +18,11 @@ type StreamFrame struct {
 	Data        []byte
 }
 
+var (
+	errInvalidStreamIDLen = errors.New("StreamFrame: Invalid StreamID length")
+	errInvalidOffsetLen   = errors.New("StreamFrame: Invalid offset length")
+)
+
 // ParseStreamFrame reads a stream frame. The type byte must not have been read yet.
 func ParseStreamFrame(r *bytes.Reader) (*StreamFrame, error) {
 	frame := &StreamFrame{}
@@ -79,38 +85,14 @@ func (f *StreamFrame) Write(b *bytes.Buffer, packetNumber protocol.PacketNumber,
 	typeByte ^= 0x20 // dataLenPresent
 
 	offsetLength := f.getOffsetLength()
-	switch offsetLength {
-	case 0:
-		typeByte ^= 0x0
-	case 2:
-		typeByte ^= 0x1 << 2
-	case 3:
-		typeByte ^= 0x2 << 2
-	case 4:
-		typeByte ^= 0x3 << 2
-	case 5:
-		typeByte ^= 0x4 << 2
-	case 6:
-		typeByte ^= 0x5 << 2
-	case 7:
-		typeByte ^= 0x6 << 2
-	case 8:
-		typeByte ^= 0x7 << 2
+
+	if offsetLength > 0 {
+		typeByte ^= (uint8(offsetLength) - 1) << 2
 	}
 
 	if f.streamIDLen == 0 {
 		f.calculateStreamIDLength()
-	}
-
-	switch f.streamIDLen {
-	case 1:
-		typeByte ^= 0x0
-	case 2:
-		typeByte ^= 0x01
-	case 3:
-		typeByte ^= 0x02
-	case 4:
-		typeByte ^= 0x03
+		typeByte ^= uint8(f.streamIDLen) - 1
 	}
 
 	b.WriteByte(typeByte)
@@ -124,6 +106,8 @@ func (f *StreamFrame) Write(b *bytes.Buffer, packetNumber protocol.PacketNumber,
 		utils.WriteUint24(b, uint32(f.StreamID))
 	case 4:
 		utils.WriteUint32(b, uint32(f.StreamID))
+	default:
+		return errInvalidStreamIDLen
 	}
 
 	switch offsetLength {
@@ -142,6 +126,8 @@ func (f *StreamFrame) Write(b *bytes.Buffer, packetNumber protocol.PacketNumber,
 		utils.WriteUint56(b, uint64(f.Offset))
 	case 8:
 		utils.WriteUint64(b, uint64(f.Offset))
+	default:
+		return errInvalidOffsetLen
 	}
 
 	utils.WriteUint16(b, uint16(len(f.Data)))
diff --git a/frames/stream_frame_test.go b/frames/stream_frame_test.go
index 432a48f8..3bf1686b 100644
--- a/frames/stream_frame_test.go
+++ b/frames/stream_frame_test.go
@@ -153,7 +153,18 @@ var _ = Describe("StreamFrame", func() {
 		})
 
 		Context("lengths of StreamIDs", func() {
-			It("writes a 2 byte StreamID", func() {
+			It("returns an error for a non-valid StreamID length", func() {
+				b := &bytes.Buffer{}
+				err := (&StreamFrame{
+					StreamID:    1,
+					streamIDLen: 13,
+					Data:        []byte("foobar"),
+				}).Write(b, 1, protocol.PacketNumberLen6, 0)
+				Expect(err).To(HaveOccurred())
+				Expect(err).To(Equal(errInvalidStreamIDLen))
+			})
+
+			It("writes a 1 byte StreamID", func() {
 				b := &bytes.Buffer{}
 				(&StreamFrame{
 					StreamID: 13,