never write too much data in Stream.Write()

ref #145

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

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)