From e8ef0de5b4644ce388648e2f126546e766555edb Mon Sep 17 00:00:00 2001
From: Lucas Clemente <lucas@clemente.io>
Date: Fri, 16 Sep 2016 21:32:57 +0200
Subject: [PATCH] exit framer loop early if no space is left in packet

---
 stream_framer.go    | 5 +++++
 streams_map.go      | 2 +-
 streams_map_test.go | 8 ++++----
 3 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/stream_framer.go b/stream_framer.go
index 92ed6e65..f1e40892 100644
--- a/stream_framer.go
+++ b/stream_framer.go
@@ -125,6 +125,11 @@ func (f *streamFramer) maybePopNormalFrames(maxBytes protocol.ByteCount) (res []
 
 		res = append(res, frame)
 		currentLen += frameHeaderBytes + frame.DataLen()
+
+		if currentLen == maxBytes {
+			return false, nil
+		}
+
 		frame = &frames.StreamFrame{DataLenPresent: true}
 		return true, nil
 	}
diff --git a/streams_map.go b/streams_map.go
index bf36f4f5..aebe7c85 100644
--- a/streams_map.go
+++ b/streams_map.go
@@ -140,10 +140,10 @@ func (m *streamsMap) RoundRobinIterate(fn streamLambda) error {
 		if err != nil {
 			return err
 		}
+		m.roundRobinIndex = (m.roundRobinIndex + 1) % numStreams
 		if !cont {
 			break
 		}
-		m.roundRobinIndex = (m.roundRobinIndex + 1) % numStreams
 	}
 	return nil
 }
diff --git a/streams_map_test.go b/streams_map_test.go
index 9edd56b4..6b9d3745 100644
--- a/streams_map_test.go
+++ b/streams_map_test.go
@@ -321,7 +321,7 @@ var _ = Describe("Streams Map", func() {
 			Expect(m.roundRobinIndex).To(Equal(3))
 		})
 
-		It("picks up at the index where it last stopped", func() {
+		It("picks up at the index+1 where it last stopped", func() {
 			fn := func(str *stream) (bool, error) {
 				lambdaCalledForStream = append(lambdaCalledForStream, str.StreamID())
 				numIterations++
@@ -334,7 +334,7 @@ var _ = Describe("Streams Map", func() {
 			Expect(err).ToNot(HaveOccurred())
 			Expect(numIterations).To(Equal(2))
 			Expect(lambdaCalledForStream).To(Equal([]protocol.StreamID{4, 5}))
-			Expect(m.roundRobinIndex).To(Equal(1))
+			Expect(m.roundRobinIndex).To(Equal(2))
 			numIterations = 0
 			lambdaCalledForStream = lambdaCalledForStream[:0]
 			fn2 := func(str *stream) (bool, error) {
@@ -347,8 +347,8 @@ var _ = Describe("Streams Map", func() {
 			}
 			err = m.RoundRobinIterate(fn2)
 			Expect(err).ToNot(HaveOccurred())
-			Expect(numIterations).To(Equal(3))
-			Expect(lambdaCalledForStream).To(Equal([]protocol.StreamID{5, 6, 7}))
+			Expect(numIterations).To(Equal(2))
+			Expect(lambdaCalledForStream).To(Equal([]protocol.StreamID{6, 7}))
 		})
 
 		It("adjust the RoundRobinIndex when deleting an element in front", func() {