From f43c4c7f1a950eb74032bcc5d3e8cbb6e3bf71b1 Mon Sep 17 00:00:00 2001
From: Lucas Clemente <lucas@clemente.io>
Date: Fri, 8 Jul 2016 13:57:43 +0200
Subject: [PATCH] respect flow control in streamFramer estimated length and
 HasData()

ref #83
---
 stream_framer.go      | 50 +++++++++++++++++++++++++++----------------
 stream_framer_test.go | 13 +++++++++++
 2 files changed, 45 insertions(+), 18 deletions(-)

diff --git a/stream_framer.go b/stream_framer.go
index e74a3f270..da64d9d60 100644
--- a/stream_framer.go
+++ b/stream_framer.go
@@ -6,6 +6,7 @@ import (
 	"github.com/lucas-clemente/quic-go/flowcontrol"
 	"github.com/lucas-clemente/quic-go/frames"
 	"github.com/lucas-clemente/quic-go/protocol"
+	"github.com/lucas-clemente/quic-go/utils"
 )
 
 type streamFramer struct {
@@ -36,6 +37,11 @@ func (f *streamFramer) HasData() bool {
 		if s == nil {
 			continue
 		}
+		// An error should never happen, and needlessly complicates the return values
+		fcLimit, _ := f.getFCAllowanceForStream(s)
+		if fcLimit == 0 {
+			continue
+		}
 		if s.lenOfDataForWriting() > 0 || s.shouldSendFin() {
 			return true
 		}
@@ -68,7 +74,9 @@ func (f *streamFramer) EstimatedDataLen() protocol.ByteCount {
 	defer f.streamsMutex.RUnlock()
 	for _, s := range *f.streams {
 		if s != nil {
-			l += s.lenOfDataForWriting()
+			// An error should never happen, and needlessly complicates the return values
+			fcLimit, _ := f.getFCAllowanceForStream(s)
+			l += utils.MinByteCount(s.lenOfDataForWriting(), fcLimit)
 			if s.shouldSendFin() {
 				l += estimatedLenOfFinFrame
 			}
@@ -127,26 +135,11 @@ func (f *streamFramer) maybePopNormalFrame(maxBytes protocol.ByteCount) (*frames
 		}
 		maxLen := maxBytes - frameHeaderBytes
 
-		flowControlWindow, err := f.flowControlManager.SendWindowSize(s.streamID)
+		fcAllowance, err := f.getFCAllowanceForStream(s)
 		if err != nil {
 			return nil, err
 		}
-		flowControlWindow -= s.writeOffset
-
-		if flowControlWindow < maxLen {
-			maxLen = flowControlWindow
-		}
-
-		contributes, err := f.flowControlManager.StreamContributesToConnectionFlowControl(s.StreamID())
-		if err != nil {
-			return nil, err
-		}
-		if contributes {
-			connectionWindow := f.flowControlManager.RemainingConnectionWindowSize()
-			if connectionWindow < maxLen {
-				maxLen = connectionWindow
-			}
-		}
+		maxLen = utils.MinByteCount(maxLen, fcAllowance)
 
 		if maxLen == 0 {
 			continue
@@ -171,6 +164,27 @@ func (f *streamFramer) maybePopNormalFrame(maxBytes protocol.ByteCount) (*frames
 	return nil, nil
 }
 
+func (f *streamFramer) getFCAllowanceForStream(s *stream) (protocol.ByteCount, error) {
+	flowControlWindow, err := f.flowControlManager.SendWindowSize(s.streamID)
+	if err != nil {
+		return 0, err
+	}
+	flowControlWindow -= s.writeOffset
+	if flowControlWindow == 0 {
+		return 0, nil
+	}
+
+	contributes, err := f.flowControlManager.StreamContributesToConnectionFlowControl(s.StreamID())
+	if err != nil {
+		return 0, err
+	}
+	connectionWindow := protocol.ByteCount(protocol.MaxByteCount)
+	if contributes {
+		connectionWindow = f.flowControlManager.RemainingConnectionWindowSize()
+	}
+	return utils.MinByteCount(flowControlWindow, connectionWindow), nil
+}
+
 // maybeSplitOffFrame removes the first n bytes and returns them as a separate frame. If n >= len(frame), nil is returned and nothing is modified.
 func maybeSplitOffFrame(frame *frames.StreamFrame, n protocol.ByteCount) *frames.StreamFrame {
 	if n >= frame.DataLen() {
diff --git a/stream_framer_test.go b/stream_framer_test.go
index a4a8d1d74..8306e1833 100644
--- a/stream_framer_test.go
+++ b/stream_framer_test.go
@@ -79,6 +79,12 @@ var _ = Describe("Stream Framer", func() {
 			Expect(framer.HasData()).To(BeTrue())
 		})
 
+		It("has no data when FC blocked", func() {
+			stream1.dataForWriting = []byte("foobar")
+			Expect(framer.HasData()).To(BeTrue())
+			fcm.sendWindowSizes[stream1.StreamID()] = 0
+			Expect(framer.HasData()).To(BeFalse())
+		})
 	})
 
 	Context("Framer estimated data length", func() {
@@ -109,6 +115,13 @@ var _ = Describe("Stream Framer", func() {
 			Expect(framer.EstimatedDataLen()).To(Equal(protocol.ByteCount(5)))
 		})
 
+		It("is zero when FC blocked", func() {
+			stream1.dataForWriting = []byte("foobar")
+			Expect(framer.EstimatedDataLen()).To(Equal(protocol.ByteCount(6)))
+			fcm.sendWindowSizes[stream1.StreamID()] = 0
+			Expect(framer.EstimatedDataLen()).To(BeZero())
+		})
+
 		It("caps the length", func() {
 			stream1.dataForWriting = bytes.Repeat([]byte{'a'}, int(protocol.MaxPacketSize)+10)
 			Expect(framer.EstimatedDataLen()).To(Equal(protocol.MaxFrameAndPublicHeaderSize))