From be34c9e3f4b95ed898cb2e46b56c6f2bf2530782 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Thu, 26 May 2016 09:13:14 +0700
Subject: [PATCH] never write too much data in Stream.Write()

ref #145
---
 stream.go      |  6 +++---
 stream_test.go | 16 ++++++++++++++++
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/stream.go b/stream.go
index fdc51ed2..08a354ad 100644
--- a/stream.go
+++ b/stream.go
@@ -182,13 +182,13 @@ func (s *stream) Write(p []byte) (int, error) {
 
 	for dataWritten < len(p) {
 		s.mutex.Lock()
-		remainingBytesInWindow := s.flowController.SendWindowSize()
+		remainingBytesInWindow := utils.MinByteCount(s.flowController.SendWindowSize(), protocol.ByteCount(len(p)-dataWritten))
 		if s.contributesToConnectionFlowControl {
 			remainingBytesInWindow = utils.MinByteCount(remainingBytesInWindow, s.connectionFlowController.SendWindowSize())
 		}
 		for remainingBytesInWindow == 0 && s.err == nil {
 			s.windowUpdateOrErrCond.Wait()
-			remainingBytesInWindow = s.flowController.SendWindowSize()
+			remainingBytesInWindow = utils.MinByteCount(s.flowController.SendWindowSize(), protocol.ByteCount(len(p)-dataWritten))
 			if s.contributesToConnectionFlowControl {
 				remainingBytesInWindow = utils.MinByteCount(remainingBytesInWindow, s.connectionFlowController.SendWindowSize())
 			}
@@ -223,7 +223,7 @@ func (s *stream) Write(p []byte) (int, error) {
 		s.maybeTriggerBlocked()
 	}
 
-	return len(p), nil
+	return dataWritten, nil
 }
 
 // Close implements io.Closer
diff --git a/stream_test.go b/stream_test.go
index bff4bea3..12b6c0a1 100644
--- a/stream_test.go
+++ b/stream_test.go
@@ -384,6 +384,22 @@ var _ = Describe("Stream", func() {
 				Expect(err).ToNot(HaveOccurred())
 			})
 
+			It("does not write too much data after receiving a window update", func() {
+				var b bool
+				updated := str.flowController.UpdateSendWindow(1)
+				Expect(updated).To(BeTrue())
+
+				go func() {
+					time.Sleep(2 * time.Millisecond)
+					b = true
+					str.UpdateSendFlowControlWindow(5)
+				}()
+				n, err := str.Write([]byte{0x13, 0x37})
+				Expect(b).To(BeTrue())
+				Expect(n).To(Equal(2))
+				Expect(err).ToNot(HaveOccurred())
+			})
+
 			It("waits for a connection flow control window update", func() {
 				var b bool
 				updated := str.flowController.UpdateSendWindow(1000)