github.com/MerlinKodo/quic-go@v0.39.2/send_stream_test.go (about)

     1  package quic
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"errors"
     7  	"io"
     8  	mrand "math/rand"
     9  	"runtime"
    10  	"time"
    11  
    12  	"golang.org/x/exp/rand"
    13  
    14  	"github.com/MerlinKodo/quic-go/internal/ackhandler"
    15  	"github.com/MerlinKodo/quic-go/internal/mocks"
    16  	"github.com/MerlinKodo/quic-go/internal/protocol"
    17  	"github.com/MerlinKodo/quic-go/internal/wire"
    18  
    19  	. "github.com/onsi/ginkgo/v2"
    20  	. "github.com/onsi/gomega"
    21  	"github.com/onsi/gomega/gbytes"
    22  	"go.uber.org/mock/gomock"
    23  )
    24  
    25  var _ = Describe("Send Stream", func() {
    26  	const streamID protocol.StreamID = 1337
    27  
    28  	var (
    29  		str            *sendStream
    30  		strWithTimeout io.Writer // str wrapped with gbytes.TimeoutWriter
    31  		mockFC         *mocks.MockStreamFlowController
    32  		mockSender     *MockStreamSender
    33  	)
    34  
    35  	BeforeEach(func() {
    36  		mockSender = NewMockStreamSender(mockCtrl)
    37  		mockFC = mocks.NewMockStreamFlowController(mockCtrl)
    38  		str = newSendStream(streamID, mockSender, mockFC)
    39  
    40  		timeout := scaleDuration(250 * time.Millisecond)
    41  		strWithTimeout = gbytes.TimeoutWriter(str, timeout)
    42  	})
    43  
    44  	expectedFrameHeaderLen := func(offset protocol.ByteCount) protocol.ByteCount {
    45  		return (&wire.StreamFrame{
    46  			StreamID:       streamID,
    47  			Offset:         offset,
    48  			DataLenPresent: true,
    49  		}).Length(protocol.Version1)
    50  	}
    51  
    52  	waitForWrite := func() {
    53  		EventuallyWithOffset(0, func() bool {
    54  			str.mutex.Lock()
    55  			hasData := str.dataForWriting != nil || str.nextFrame != nil
    56  			str.mutex.Unlock()
    57  			return hasData
    58  		}).Should(BeTrue())
    59  	}
    60  
    61  	getDataAtOffset := func(offset, length protocol.ByteCount) []byte {
    62  		b := make([]byte, length)
    63  		for i := protocol.ByteCount(0); i < length; i++ {
    64  			b[i] = uint8(offset + i)
    65  		}
    66  		return b
    67  	}
    68  
    69  	getData := func(length protocol.ByteCount) []byte {
    70  		return getDataAtOffset(0, length)
    71  	}
    72  
    73  	It("gets stream id", func() {
    74  		Expect(str.StreamID()).To(Equal(protocol.StreamID(1337)))
    75  	})
    76  
    77  	Context("writing", func() {
    78  		It("writes and gets all data at once", func() {
    79  			done := make(chan struct{})
    80  			go func() {
    81  				defer GinkgoRecover()
    82  				defer close(done)
    83  				mockSender.EXPECT().onHasStreamData(streamID)
    84  				n, err := strWithTimeout.Write([]byte("foobar"))
    85  				Expect(err).ToNot(HaveOccurred())
    86  				Expect(n).To(Equal(6))
    87  			}()
    88  			waitForWrite()
    89  			mockFC.EXPECT().SendWindowSize().Return(protocol.MaxByteCount)
    90  			mockFC.EXPECT().AddBytesSent(protocol.ByteCount(6))
    91  			frame, ok, _ := str.popStreamFrame(protocol.MaxByteCount, protocol.Version1)
    92  			Expect(ok).To(BeTrue())
    93  			f := frame.Frame
    94  			Expect(f.Data).To(Equal([]byte("foobar")))
    95  			Expect(f.Fin).To(BeFalse())
    96  			Expect(f.Offset).To(BeZero())
    97  			Expect(f.DataLenPresent).To(BeTrue())
    98  			Expect(str.writeOffset).To(Equal(protocol.ByteCount(6)))
    99  			Expect(str.dataForWriting).To(BeNil())
   100  			Eventually(done).Should(BeClosed())
   101  		})
   102  
   103  		It("writes and gets data in two turns", func() {
   104  			done := make(chan struct{})
   105  			go func() {
   106  				defer GinkgoRecover()
   107  				mockSender.EXPECT().onHasStreamData(streamID)
   108  				n, err := strWithTimeout.Write([]byte("foobar"))
   109  				Expect(err).ToNot(HaveOccurred())
   110  				Expect(n).To(Equal(6))
   111  				close(done)
   112  			}()
   113  			waitForWrite()
   114  			mockFC.EXPECT().SendWindowSize().Return(protocol.MaxByteCount).Times(2)
   115  			mockFC.EXPECT().AddBytesSent(protocol.ByteCount(3)).Times(2)
   116  			frame, ok, _ := str.popStreamFrame(expectedFrameHeaderLen(0)+3, protocol.Version1)
   117  			Expect(ok).To(BeTrue())
   118  			f := frame.Frame
   119  			Expect(f.Offset).To(BeZero())
   120  			Expect(f.Fin).To(BeFalse())
   121  			Expect(f.Data).To(Equal([]byte("foo")))
   122  			Expect(f.DataLenPresent).To(BeTrue())
   123  			frame, ok, _ = str.popStreamFrame(protocol.MaxByteCount, protocol.Version1)
   124  			Expect(ok).To(BeTrue())
   125  			f = frame.Frame
   126  			Expect(f.Data).To(Equal([]byte("bar")))
   127  			Expect(f.Fin).To(BeFalse())
   128  			Expect(f.Offset).To(Equal(protocol.ByteCount(3)))
   129  			Expect(f.DataLenPresent).To(BeTrue())
   130  			_, ok, _ = str.popStreamFrame(1000, protocol.Version1)
   131  			Expect(ok).To(BeFalse())
   132  			Eventually(done).Should(BeClosed())
   133  		})
   134  
   135  		It("bundles small writes", func() {
   136  			done := make(chan struct{})
   137  			go func() {
   138  				defer GinkgoRecover()
   139  				mockSender.EXPECT().onHasStreamData(streamID).Times(2)
   140  				n, err := strWithTimeout.Write([]byte("foo"))
   141  				Expect(err).ToNot(HaveOccurred())
   142  				Expect(n).To(Equal(3))
   143  				n, err = strWithTimeout.Write([]byte("bar"))
   144  				Expect(err).ToNot(HaveOccurred())
   145  				Expect(n).To(Equal(3))
   146  				close(done)
   147  			}()
   148  			Eventually(done).Should(BeClosed()) // both Write calls returned without any data having been dequeued yet
   149  			mockFC.EXPECT().SendWindowSize().Return(protocol.MaxByteCount)
   150  			mockFC.EXPECT().AddBytesSent(protocol.ByteCount(6))
   151  			frame, ok, _ := str.popStreamFrame(protocol.MaxByteCount, protocol.Version1)
   152  			Expect(ok).To(BeTrue())
   153  			f := frame.Frame
   154  			Expect(f.Offset).To(BeZero())
   155  			Expect(f.Fin).To(BeFalse())
   156  			Expect(f.Data).To(Equal([]byte("foobar")))
   157  		})
   158  
   159  		It("writes and gets data in multiple turns, for large writes", func() {
   160  			mockFC.EXPECT().SendWindowSize().Return(protocol.MaxByteCount).Times(5)
   161  			var totalBytesSent protocol.ByteCount
   162  			mockFC.EXPECT().AddBytesSent(gomock.Any()).Do(func(l protocol.ByteCount) { totalBytesSent += l }).Times(5)
   163  			done := make(chan struct{})
   164  			go func() {
   165  				defer GinkgoRecover()
   166  				mockSender.EXPECT().onHasStreamData(streamID)
   167  				n, err := strWithTimeout.Write(getData(5000))
   168  				Expect(err).ToNot(HaveOccurred())
   169  				Expect(n).To(Equal(5000))
   170  				close(done)
   171  			}()
   172  			waitForWrite()
   173  			for i := 0; i < 5; i++ {
   174  				frame, ok, _ := str.popStreamFrame(1100, protocol.Version1)
   175  				Expect(ok).To(BeTrue())
   176  				f := frame.Frame
   177  				Expect(f.Offset).To(BeNumerically("~", 1100*i, 10*i))
   178  				Expect(f.Fin).To(BeFalse())
   179  				Expect(f.Data).To(Equal(getDataAtOffset(f.Offset, f.DataLen())))
   180  				Expect(f.DataLenPresent).To(BeTrue())
   181  			}
   182  			Expect(totalBytesSent).To(Equal(protocol.ByteCount(5000)))
   183  			Eventually(done).Should(BeClosed())
   184  		})
   185  
   186  		It("unblocks Write as soon as a STREAM frame can be buffered", func() {
   187  			done := make(chan struct{})
   188  			go func() {
   189  				defer GinkgoRecover()
   190  				defer close(done)
   191  				mockSender.EXPECT().onHasStreamData(streamID)
   192  				_, err := strWithTimeout.Write(getData(protocol.MaxPacketBufferSize + 3))
   193  				Expect(err).ToNot(HaveOccurred())
   194  			}()
   195  			waitForWrite()
   196  			mockFC.EXPECT().SendWindowSize().Return(protocol.MaxByteCount).Times(2)
   197  			mockFC.EXPECT().AddBytesSent(protocol.ByteCount(2))
   198  			frame, ok, hasMoreData := str.popStreamFrame(expectedFrameHeaderLen(0)+2, protocol.Version1)
   199  			Expect(ok).To(BeTrue())
   200  			Expect(hasMoreData).To(BeTrue())
   201  			f := frame.Frame
   202  			Expect(f.DataLen()).To(Equal(protocol.ByteCount(2)))
   203  			Consistently(done).ShouldNot(BeClosed())
   204  			mockFC.EXPECT().AddBytesSent(protocol.ByteCount(1))
   205  			frame, ok, hasMoreData = str.popStreamFrame(expectedFrameHeaderLen(1)+1, protocol.Version1)
   206  			Expect(ok).To(BeTrue())
   207  			Expect(hasMoreData).To(BeTrue())
   208  			f = frame.Frame
   209  			Expect(f.DataLen()).To(Equal(protocol.ByteCount(1)))
   210  			Eventually(done).Should(BeClosed())
   211  		})
   212  
   213  		It("only unblocks Write once a previously buffered STREAM frame has been fully dequeued", func() {
   214  			mockSender.EXPECT().onHasStreamData(streamID)
   215  			_, err := strWithTimeout.Write([]byte("foobar"))
   216  			Expect(err).ToNot(HaveOccurred())
   217  			done := make(chan struct{})
   218  			go func() {
   219  				defer GinkgoRecover()
   220  				defer close(done)
   221  				mockSender.EXPECT().onHasStreamData(streamID)
   222  				_, err := str.Write(getData(protocol.MaxPacketBufferSize))
   223  				Expect(err).ToNot(HaveOccurred())
   224  			}()
   225  			waitForWrite()
   226  			mockFC.EXPECT().SendWindowSize().Return(protocol.MaxByteCount).Times(2)
   227  			mockFC.EXPECT().AddBytesSent(protocol.ByteCount(2))
   228  			frame, ok, hasMoreData := str.popStreamFrame(expectedFrameHeaderLen(0)+2, protocol.Version1)
   229  			Expect(ok).To(BeTrue())
   230  			Expect(hasMoreData).To(BeTrue())
   231  			f := frame.Frame
   232  			Expect(f.Data).To(Equal([]byte("fo")))
   233  			Consistently(done).ShouldNot(BeClosed())
   234  			mockFC.EXPECT().AddBytesSent(protocol.ByteCount(4))
   235  			frame, ok, hasMoreData = str.popStreamFrame(expectedFrameHeaderLen(2)+4, protocol.Version1)
   236  			Expect(ok).To(BeTrue())
   237  			Expect(hasMoreData).To(BeTrue())
   238  			f = frame.Frame
   239  			Expect(f.Data).To(Equal([]byte("obar")))
   240  			Eventually(done).Should(BeClosed())
   241  		})
   242  
   243  		It("popStreamFrame returns nil if no data is available", func() {
   244  			_, ok, hasMoreData := str.popStreamFrame(1000, protocol.Version1)
   245  			Expect(ok).To(BeFalse())
   246  			Expect(hasMoreData).To(BeFalse())
   247  		})
   248  
   249  		It("says if it has more data for writing", func() {
   250  			done := make(chan struct{})
   251  			go func() {
   252  				defer GinkgoRecover()
   253  				defer close(done)
   254  				mockSender.EXPECT().onHasStreamData(streamID)
   255  				n, err := strWithTimeout.Write(bytes.Repeat([]byte{0}, 100))
   256  				Expect(err).ToNot(HaveOccurred())
   257  				Expect(n).To(Equal(100))
   258  			}()
   259  			waitForWrite()
   260  			mockFC.EXPECT().SendWindowSize().Return(protocol.MaxByteCount).Times(2)
   261  			mockFC.EXPECT().AddBytesSent(gomock.Any()).Times(2)
   262  			frame, ok, hasMoreData := str.popStreamFrame(50, protocol.Version1)
   263  			Expect(ok).To(BeTrue())
   264  			Expect(frame.Frame.Fin).To(BeFalse())
   265  			Expect(hasMoreData).To(BeTrue())
   266  			frame, ok, hasMoreData = str.popStreamFrame(protocol.MaxByteCount, protocol.Version1)
   267  			Expect(ok).To(BeTrue())
   268  			Expect(frame.Frame.Fin).To(BeFalse())
   269  			Expect(hasMoreData).To(BeFalse())
   270  			_, ok, _ = str.popStreamFrame(protocol.MaxByteCount, protocol.Version1)
   271  			Expect(ok).To(BeFalse())
   272  			Eventually(done).Should(BeClosed())
   273  		})
   274  
   275  		It("copies the slice while writing", func() {
   276  			frameHeaderSize := protocol.ByteCount(4)
   277  			mockFC.EXPECT().SendWindowSize().Return(protocol.MaxByteCount).Times(2)
   278  			mockFC.EXPECT().AddBytesSent(protocol.ByteCount(1))
   279  			mockFC.EXPECT().AddBytesSent(protocol.ByteCount(2))
   280  			s := []byte("foo")
   281  			done := make(chan struct{})
   282  			go func() {
   283  				defer GinkgoRecover()
   284  				defer close(done)
   285  				mockSender.EXPECT().onHasStreamData(streamID)
   286  				n, err := strWithTimeout.Write(s)
   287  				Expect(err).ToNot(HaveOccurred())
   288  				Expect(n).To(Equal(3))
   289  			}()
   290  			waitForWrite()
   291  			frame, ok, _ := str.popStreamFrame(frameHeaderSize+1, protocol.Version1)
   292  			Expect(ok).To(BeTrue())
   293  			f := frame.Frame
   294  			Expect(f.Data).To(Equal([]byte("f")))
   295  			frame, ok, _ = str.popStreamFrame(100, protocol.Version1)
   296  			Expect(ok).To(BeTrue())
   297  			Expect(frame).ToNot(BeNil())
   298  			f = frame.Frame
   299  			Expect(f.Data).To(Equal([]byte("oo")))
   300  			s[1] = 'e'
   301  			Expect(f.Data).To(Equal([]byte("oo")))
   302  			Eventually(done).Should(BeClosed())
   303  		})
   304  
   305  		It("returns when given a nil input", func() {
   306  			n, err := strWithTimeout.Write(nil)
   307  			Expect(n).To(BeZero())
   308  			Expect(err).ToNot(HaveOccurred())
   309  		})
   310  
   311  		It("returns when given an empty slice", func() {
   312  			n, err := strWithTimeout.Write([]byte(""))
   313  			Expect(n).To(BeZero())
   314  			Expect(err).ToNot(HaveOccurred())
   315  		})
   316  
   317  		It("cancels the context when Close is called", func() {
   318  			mockSender.EXPECT().onHasStreamData(streamID)
   319  			Expect(str.Context().Done()).ToNot(BeClosed())
   320  			Expect(str.Close()).To(Succeed())
   321  			Expect(str.Context().Done()).To(BeClosed())
   322  			Expect(context.Cause(str.Context())).To(MatchError(context.Canceled))
   323  		})
   324  
   325  		Context("flow control blocking", func() {
   326  			It("queues a BLOCKED frame if the stream is flow control blocked", func() {
   327  				mockFC.EXPECT().SendWindowSize().Return(protocol.ByteCount(0))
   328  				mockFC.EXPECT().IsNewlyBlocked().Return(true, protocol.ByteCount(12))
   329  				mockSender.EXPECT().queueControlFrame(&wire.StreamDataBlockedFrame{
   330  					StreamID:          streamID,
   331  					MaximumStreamData: 12,
   332  				})
   333  				done := make(chan struct{})
   334  				go func() {
   335  					defer GinkgoRecover()
   336  					defer close(done)
   337  					mockSender.EXPECT().onHasStreamData(streamID)
   338  					_, err := str.Write([]byte("foobar"))
   339  					Expect(err).ToNot(HaveOccurred())
   340  				}()
   341  				waitForWrite()
   342  				_, ok, hasMoreData := str.popStreamFrame(1000, protocol.Version1)
   343  				Expect(ok).To(BeFalse())
   344  				Expect(hasMoreData).To(BeFalse())
   345  				// make the Write go routine return
   346  				str.closeForShutdown(nil)
   347  				Eventually(done).Should(BeClosed())
   348  			})
   349  
   350  			It("says that it doesn't have any more data, when it is flow control blocked", func() {
   351  				done := make(chan struct{})
   352  				go func() {
   353  					defer GinkgoRecover()
   354  					defer close(done)
   355  					mockSender.EXPECT().onHasStreamData(streamID)
   356  					_, err := str.Write([]byte("foobar"))
   357  					Expect(err).ToNot(HaveOccurred())
   358  				}()
   359  				waitForWrite()
   360  
   361  				// first pop a STREAM frame of the maximum size allowed by flow control
   362  				mockFC.EXPECT().SendWindowSize().Return(protocol.ByteCount(3))
   363  				mockFC.EXPECT().AddBytesSent(protocol.ByteCount(3))
   364  				f, ok, hasMoreData := str.popStreamFrame(expectedFrameHeaderLen(0)+3, protocol.Version1)
   365  				Expect(ok).To(BeTrue())
   366  				Expect(f).ToNot(BeNil())
   367  				Expect(hasMoreData).To(BeTrue())
   368  
   369  				// try to pop again, this time noticing that we're blocked
   370  				mockFC.EXPECT().SendWindowSize()
   371  				// don't use offset 3 here, to make sure the BLOCKED frame contains the number returned by the flow controller
   372  				mockFC.EXPECT().IsNewlyBlocked().Return(true, protocol.ByteCount(10))
   373  				mockSender.EXPECT().queueControlFrame(&wire.StreamDataBlockedFrame{
   374  					StreamID:          streamID,
   375  					MaximumStreamData: 10,
   376  				})
   377  				_, ok, hasMoreData = str.popStreamFrame(1000, protocol.Version1)
   378  				Expect(ok).To(BeFalse())
   379  				Expect(hasMoreData).To(BeFalse())
   380  				// make the Write go routine return
   381  				str.closeForShutdown(nil)
   382  				Eventually(done).Should(BeClosed())
   383  			})
   384  		})
   385  
   386  		Context("deadlines", func() {
   387  			It("returns an error when Write is called after the deadline", func() {
   388  				str.SetWriteDeadline(time.Now().Add(-time.Second))
   389  				n, err := strWithTimeout.Write([]byte("foobar"))
   390  				Expect(err).To(MatchError(errDeadline))
   391  				Expect(n).To(BeZero())
   392  			})
   393  
   394  			It("unblocks after the deadline", func() {
   395  				mockSender.EXPECT().onHasStreamData(streamID)
   396  				deadline := time.Now().Add(scaleDuration(50 * time.Millisecond))
   397  				str.SetWriteDeadline(deadline)
   398  				n, err := strWithTimeout.Write(getData(5000))
   399  				Expect(err).To(MatchError(errDeadline))
   400  				Expect(n).To(BeZero())
   401  				Expect(time.Now()).To(BeTemporally("~", deadline, scaleDuration(20*time.Millisecond)))
   402  			})
   403  
   404  			It("unblocks when the deadline is changed to the past", func() {
   405  				mockSender.EXPECT().onHasStreamData(streamID)
   406  				str.SetWriteDeadline(time.Now().Add(time.Hour))
   407  				done := make(chan struct{})
   408  				go func() {
   409  					defer GinkgoRecover()
   410  					_, err := str.Write(getData(5000))
   411  					Expect(err).To(MatchError(errDeadline))
   412  					close(done)
   413  				}()
   414  				Consistently(done).ShouldNot(BeClosed())
   415  				str.SetWriteDeadline(time.Now().Add(-time.Hour))
   416  				Eventually(done).Should(BeClosed())
   417  			})
   418  
   419  			It("returns the number of bytes written, when the deadline expires", func() {
   420  				mockFC.EXPECT().SendWindowSize().Return(protocol.MaxByteCount).AnyTimes()
   421  				mockFC.EXPECT().AddBytesSent(gomock.Any())
   422  				deadline := time.Now().Add(scaleDuration(50 * time.Millisecond))
   423  				str.SetWriteDeadline(deadline)
   424  				var n int
   425  				writeReturned := make(chan struct{})
   426  				go func() {
   427  					defer GinkgoRecover()
   428  					defer close(writeReturned)
   429  					mockSender.EXPECT().onHasStreamData(streamID)
   430  					var err error
   431  					n, err = strWithTimeout.Write(getData(5000))
   432  					Expect(err).To(MatchError(errDeadline))
   433  					Expect(time.Now()).To(BeTemporally("~", deadline, scaleDuration(20*time.Millisecond)))
   434  				}()
   435  				waitForWrite()
   436  				frame, ok, hasMoreData := str.popStreamFrame(50, protocol.Version1)
   437  				Expect(ok).To(BeTrue())
   438  				Expect(frame).ToNot(BeNil())
   439  				Expect(hasMoreData).To(BeTrue())
   440  				Eventually(writeReturned, scaleDuration(80*time.Millisecond)).Should(BeClosed())
   441  				Expect(n).To(BeEquivalentTo(frame.Frame.DataLen()))
   442  			})
   443  
   444  			It("doesn't pop any data after the deadline expired", func() {
   445  				mockFC.EXPECT().SendWindowSize().Return(protocol.MaxByteCount).AnyTimes()
   446  				mockFC.EXPECT().AddBytesSent(gomock.Any())
   447  				deadline := time.Now().Add(scaleDuration(50 * time.Millisecond))
   448  				str.SetWriteDeadline(deadline)
   449  				writeReturned := make(chan struct{})
   450  				go func() {
   451  					defer GinkgoRecover()
   452  					defer close(writeReturned)
   453  					mockSender.EXPECT().onHasStreamData(streamID)
   454  					_, err := strWithTimeout.Write(getData(5000))
   455  					Expect(err).To(MatchError(errDeadline))
   456  				}()
   457  				waitForWrite()
   458  				frame, ok, hasMoreData := str.popStreamFrame(50, protocol.Version1)
   459  				Expect(ok).To(BeTrue())
   460  				Expect(frame).ToNot(BeNil())
   461  				Expect(hasMoreData).To(BeTrue())
   462  				Eventually(writeReturned, scaleDuration(80*time.Millisecond)).Should(BeClosed())
   463  				_, ok, hasMoreData = str.popStreamFrame(50, protocol.Version1)
   464  				Expect(ok).To(BeFalse())
   465  				Expect(hasMoreData).To(BeFalse())
   466  			})
   467  
   468  			It("doesn't unblock if the deadline is changed before the first one expires", func() {
   469  				mockSender.EXPECT().onHasStreamData(streamID)
   470  				deadline1 := time.Now().Add(scaleDuration(50 * time.Millisecond))
   471  				deadline2 := time.Now().Add(scaleDuration(100 * time.Millisecond))
   472  				str.SetWriteDeadline(deadline1)
   473  				done := make(chan struct{})
   474  				go func() {
   475  					defer GinkgoRecover()
   476  					time.Sleep(scaleDuration(20 * time.Millisecond))
   477  					str.SetWriteDeadline(deadline2)
   478  					// make sure that this was actually execute before the deadline expires
   479  					Expect(time.Now()).To(BeTemporally("<", deadline1))
   480  					close(done)
   481  				}()
   482  				runtime.Gosched()
   483  				n, err := strWithTimeout.Write(getData(5000))
   484  				Expect(err).To(MatchError(errDeadline))
   485  				Expect(n).To(BeZero())
   486  				Expect(time.Now()).To(BeTemporally("~", deadline2, scaleDuration(20*time.Millisecond)))
   487  				Eventually(done).Should(BeClosed())
   488  			})
   489  
   490  			It("unblocks earlier, when a new deadline is set", func() {
   491  				mockSender.EXPECT().onHasStreamData(streamID)
   492  				deadline1 := time.Now().Add(scaleDuration(200 * time.Millisecond))
   493  				deadline2 := time.Now().Add(scaleDuration(50 * time.Millisecond))
   494  				done := make(chan struct{})
   495  				go func() {
   496  					defer GinkgoRecover()
   497  					time.Sleep(scaleDuration(10 * time.Millisecond))
   498  					str.SetWriteDeadline(deadline2)
   499  					// make sure that this was actually execute before the deadline expires
   500  					Expect(time.Now()).To(BeTemporally("<", deadline2))
   501  					close(done)
   502  				}()
   503  				str.SetWriteDeadline(deadline1)
   504  				runtime.Gosched()
   505  				_, err := strWithTimeout.Write(getData(5000))
   506  				Expect(err).To(MatchError(errDeadline))
   507  				Expect(time.Now()).To(BeTemporally("~", deadline2, scaleDuration(20*time.Millisecond)))
   508  				Eventually(done).Should(BeClosed())
   509  			})
   510  
   511  			It("doesn't unblock if the deadline is removed", func() {
   512  				mockSender.EXPECT().onHasStreamData(streamID)
   513  				deadline := time.Now().Add(scaleDuration(50 * time.Millisecond))
   514  				str.SetWriteDeadline(deadline)
   515  				deadlineUnset := make(chan struct{})
   516  				go func() {
   517  					defer GinkgoRecover()
   518  					time.Sleep(scaleDuration(20 * time.Millisecond))
   519  					str.SetWriteDeadline(time.Time{})
   520  					// make sure that this was actually execute before the deadline expires
   521  					Expect(time.Now()).To(BeTemporally("<", deadline))
   522  					close(deadlineUnset)
   523  				}()
   524  				done := make(chan struct{})
   525  				go func() {
   526  					defer GinkgoRecover()
   527  					_, err := strWithTimeout.Write(getData(5000))
   528  					Expect(err).To(MatchError("test done"))
   529  					close(done)
   530  				}()
   531  				runtime.Gosched()
   532  				Eventually(deadlineUnset).Should(BeClosed())
   533  				Consistently(done, scaleDuration(100*time.Millisecond)).ShouldNot(BeClosed())
   534  				// make the go routine return
   535  				str.closeForShutdown(errors.New("test done"))
   536  				Eventually(done).Should(BeClosed())
   537  			})
   538  		})
   539  
   540  		Context("closing", func() {
   541  			It("doesn't allow writes after it has been closed", func() {
   542  				mockSender.EXPECT().onHasStreamData(streamID)
   543  				str.Close()
   544  				_, err := strWithTimeout.Write([]byte("foobar"))
   545  				Expect(err).To(MatchError("write on closed stream 1337"))
   546  			})
   547  
   548  			It("allows FIN", func() {
   549  				mockSender.EXPECT().onHasStreamData(streamID)
   550  				str.Close()
   551  				frame, ok, hasMoreData := str.popStreamFrame(1000, protocol.Version1)
   552  				Expect(ok).To(BeTrue())
   553  				Expect(frame).ToNot(BeNil())
   554  				f := frame.Frame
   555  				Expect(f.Data).To(BeEmpty())
   556  				Expect(f.Fin).To(BeTrue())
   557  				Expect(f.DataLenPresent).To(BeTrue())
   558  				Expect(hasMoreData).To(BeFalse())
   559  			})
   560  
   561  			It("doesn't send a FIN when there's still data", func() {
   562  				const frameHeaderLen protocol.ByteCount = 4
   563  				mockSender.EXPECT().onHasStreamData(streamID).Times(2)
   564  				_, err := strWithTimeout.Write([]byte("foobar"))
   565  				Expect(err).ToNot(HaveOccurred())
   566  				Expect(str.Close()).To(Succeed())
   567  				mockFC.EXPECT().SendWindowSize().Return(protocol.MaxByteCount).Times(2)
   568  				mockFC.EXPECT().AddBytesSent(gomock.Any()).Times(2)
   569  				frame, ok, _ := str.popStreamFrame(3+frameHeaderLen, protocol.Version1)
   570  				Expect(ok).To(BeTrue())
   571  				Expect(frame).ToNot(BeNil())
   572  				f := frame.Frame
   573  				Expect(f.Data).To(Equal([]byte("foo")))
   574  				Expect(f.Fin).To(BeFalse())
   575  				frame, ok, _ = str.popStreamFrame(protocol.MaxByteCount, protocol.Version1)
   576  				Expect(ok).To(BeTrue())
   577  				f = frame.Frame
   578  				Expect(f.Data).To(Equal([]byte("bar")))
   579  				Expect(f.Fin).To(BeTrue())
   580  			})
   581  
   582  			It("doesn't send a FIN when there's still data, for long writes", func() {
   583  				done := make(chan struct{})
   584  				go func() {
   585  					defer GinkgoRecover()
   586  					defer close(done)
   587  					mockSender.EXPECT().onHasStreamData(streamID)
   588  					_, err := strWithTimeout.Write(getData(5000))
   589  					Expect(err).ToNot(HaveOccurred())
   590  					mockSender.EXPECT().onHasStreamData(streamID)
   591  					Expect(str.Close()).To(Succeed())
   592  				}()
   593  				waitForWrite()
   594  				for i := 1; i <= 5; i++ {
   595  					mockFC.EXPECT().SendWindowSize().Return(protocol.MaxByteCount)
   596  					mockFC.EXPECT().AddBytesSent(gomock.Any())
   597  					if i == 5 {
   598  						Eventually(done).Should(BeClosed())
   599  					}
   600  					frame, ok, _ := str.popStreamFrame(1100, protocol.Version1)
   601  					Expect(ok).To(BeTrue())
   602  					Expect(frame).ToNot(BeNil())
   603  					f := frame.Frame
   604  					Expect(f.Data).To(Equal(getDataAtOffset(f.Offset, f.DataLen())))
   605  					Expect(f.Fin).To(Equal(i == 5)) // the last frame should have the FIN bit set
   606  				}
   607  			})
   608  
   609  			It("doesn't allow FIN after it is closed for shutdown", func() {
   610  				str.closeForShutdown(errors.New("test"))
   611  				_, ok, hasMoreData := str.popStreamFrame(1000, protocol.Version1)
   612  				Expect(ok).To(BeFalse())
   613  				Expect(hasMoreData).To(BeFalse())
   614  
   615  				Expect(str.Close()).To(Succeed())
   616  				_, ok, hasMoreData = str.popStreamFrame(1000, protocol.Version1)
   617  				Expect(ok).To(BeFalse())
   618  				Expect(hasMoreData).To(BeFalse())
   619  			})
   620  
   621  			It("doesn't allow FIN twice", func() {
   622  				mockSender.EXPECT().onHasStreamData(streamID)
   623  				str.Close()
   624  				frame, ok, _ := str.popStreamFrame(1000, protocol.Version1)
   625  				Expect(ok).To(BeTrue())
   626  				Expect(frame).ToNot(BeNil())
   627  				f := frame.Frame
   628  				Expect(f.Data).To(BeEmpty())
   629  				Expect(f.Fin).To(BeTrue())
   630  				_, ok, hasMoreData := str.popStreamFrame(1000, protocol.Version1)
   631  				Expect(ok).To(BeFalse())
   632  				Expect(hasMoreData).To(BeFalse())
   633  			})
   634  		})
   635  
   636  		Context("closing for shutdown", func() {
   637  			testErr := errors.New("test")
   638  
   639  			It("returns errors when the stream is cancelled", func() {
   640  				str.closeForShutdown(testErr)
   641  				n, err := strWithTimeout.Write([]byte("foo"))
   642  				Expect(n).To(BeZero())
   643  				Expect(err).To(MatchError(testErr))
   644  			})
   645  
   646  			It("doesn't get data for writing if an error occurred", func() {
   647  				mockFC.EXPECT().SendWindowSize().Return(protocol.MaxByteCount)
   648  				mockFC.EXPECT().AddBytesSent(gomock.Any())
   649  				mockSender.EXPECT().onHasStreamData(streamID)
   650  				done := make(chan struct{})
   651  				go func() {
   652  					defer GinkgoRecover()
   653  					_, err := strWithTimeout.Write(getData(5000))
   654  					Expect(err).To(MatchError(testErr))
   655  					close(done)
   656  				}()
   657  				waitForWrite()
   658  				frame, ok, hasMoreData := str.popStreamFrame(50, protocol.Version1) // get a STREAM frame containing some data, but not all
   659  				Expect(ok).To(BeTrue())
   660  				Expect(frame).ToNot(BeNil())
   661  				Expect(hasMoreData).To(BeTrue())
   662  				str.closeForShutdown(testErr)
   663  				_, ok, hasMoreData = str.popStreamFrame(1000, protocol.Version1)
   664  				Expect(ok).To(BeFalse())
   665  				Expect(hasMoreData).To(BeFalse())
   666  				Eventually(done).Should(BeClosed())
   667  			})
   668  
   669  			It("cancels the context", func() {
   670  				Expect(str.Context().Done()).ToNot(BeClosed())
   671  				str.closeForShutdown(testErr)
   672  				Expect(str.Context().Done()).To(BeClosed())
   673  				Expect(context.Cause(str.Context())).To(MatchError(testErr))
   674  			})
   675  		})
   676  	})
   677  
   678  	Context("handling MAX_STREAM_DATA frames", func() {
   679  		It("informs the flow controller", func() {
   680  			mockFC.EXPECT().UpdateSendWindow(protocol.ByteCount(0x1337))
   681  			str.updateSendWindow(0x1337)
   682  		})
   683  
   684  		It("says when it has data for sending", func() {
   685  			mockFC.EXPECT().UpdateSendWindow(gomock.Any())
   686  			mockSender.EXPECT().onHasStreamData(streamID)
   687  			done := make(chan struct{})
   688  			go func() {
   689  				defer GinkgoRecover()
   690  				_, err := str.Write([]byte("foobar"))
   691  				Expect(err).ToNot(HaveOccurred())
   692  				close(done)
   693  			}()
   694  			waitForWrite()
   695  			mockSender.EXPECT().onHasStreamData(streamID)
   696  			str.updateSendWindow(42)
   697  			// make sure the Write go routine returns
   698  			str.closeForShutdown(nil)
   699  			Eventually(done).Should(BeClosed())
   700  		})
   701  	})
   702  
   703  	Context("stream cancellations", func() {
   704  		Context("canceling writing", func() {
   705  			It("queues a RESET_STREAM frame", func() {
   706  				gomock.InOrder(
   707  					mockSender.EXPECT().queueControlFrame(&wire.ResetStreamFrame{
   708  						StreamID:  streamID,
   709  						FinalSize: 1234,
   710  						ErrorCode: 9876,
   711  					}),
   712  					mockSender.EXPECT().onStreamCompleted(streamID),
   713  				)
   714  				str.writeOffset = 1234
   715  				str.CancelWrite(9876)
   716  			})
   717  
   718  			// This test is inherently racy, as it tests a concurrent call to Write() and CancelRead().
   719  			// A single successful run of this test therefore doesn't mean a lot,
   720  			// for reliable results it has to be run many times.
   721  			It("returns a nil error when the whole slice has been sent out", func() {
   722  				mockSender.EXPECT().queueControlFrame(gomock.Any()).MaxTimes(1)
   723  				mockSender.EXPECT().onHasStreamData(streamID).MaxTimes(1)
   724  				mockSender.EXPECT().onStreamCompleted(streamID).MaxTimes(1)
   725  				mockFC.EXPECT().SendWindowSize().Return(protocol.MaxByteCount).MaxTimes(1)
   726  				mockFC.EXPECT().AddBytesSent(gomock.Any()).MaxTimes(1)
   727  				errChan := make(chan error)
   728  				go func() {
   729  					defer GinkgoRecover()
   730  					n, err := strWithTimeout.Write(getData(100))
   731  					if n == 0 {
   732  						errChan <- nil
   733  						return
   734  					}
   735  					errChan <- err
   736  				}()
   737  
   738  				runtime.Gosched()
   739  				go str.popStreamFrame(protocol.MaxByteCount, protocol.Version1)
   740  				go str.CancelWrite(1234)
   741  				Eventually(errChan).Should(Receive(Not(HaveOccurred())))
   742  			})
   743  
   744  			It("unblocks Write", func() {
   745  				mockSender.EXPECT().queueControlFrame(gomock.Any())
   746  				mockSender.EXPECT().onHasStreamData(streamID)
   747  				mockFC.EXPECT().SendWindowSize().Return(protocol.MaxByteCount)
   748  				mockFC.EXPECT().AddBytesSent(gomock.Any())
   749  				writeReturned := make(chan struct{})
   750  				var n int
   751  				go func() {
   752  					defer GinkgoRecover()
   753  					var err error
   754  					n, err = strWithTimeout.Write(getData(5000))
   755  					Expect(err).To(Equal(&StreamError{
   756  						StreamID:  streamID,
   757  						ErrorCode: 1234,
   758  						Remote:    false,
   759  					}))
   760  					close(writeReturned)
   761  				}()
   762  				waitForWrite()
   763  				frame, ok, _ := str.popStreamFrame(50, protocol.Version1)
   764  				Expect(ok).To(BeTrue())
   765  				Expect(frame).ToNot(BeNil())
   766  				mockSender.EXPECT().onStreamCompleted(streamID)
   767  				str.CancelWrite(1234)
   768  				Eventually(writeReturned).Should(BeClosed())
   769  				Expect(n).To(BeEquivalentTo(frame.Frame.DataLen()))
   770  			})
   771  
   772  			It("doesn't pop STREAM frames after being canceled", func() {
   773  				mockSender.EXPECT().queueControlFrame(gomock.Any())
   774  				mockSender.EXPECT().onHasStreamData(streamID)
   775  				mockFC.EXPECT().SendWindowSize().Return(protocol.MaxByteCount)
   776  				mockFC.EXPECT().AddBytesSent(gomock.Any())
   777  				writeReturned := make(chan struct{})
   778  				go func() {
   779  					defer GinkgoRecover()
   780  					strWithTimeout.Write(getData(100))
   781  					close(writeReturned)
   782  				}()
   783  				waitForWrite()
   784  				frame, ok, hasMoreData := str.popStreamFrame(50, protocol.Version1)
   785  				Expect(ok).To(BeTrue())
   786  				Expect(hasMoreData).To(BeTrue())
   787  				Expect(frame).ToNot(BeNil())
   788  				mockSender.EXPECT().onStreamCompleted(streamID)
   789  				str.CancelWrite(1234)
   790  				_, ok, hasMoreData = str.popStreamFrame(10, protocol.Version1)
   791  				Expect(ok).To(BeFalse())
   792  				Expect(hasMoreData).To(BeFalse())
   793  				Eventually(writeReturned).Should(BeClosed())
   794  			})
   795  
   796  			It("doesn't pop STREAM frames after being canceled, for large writes", func() {
   797  				mockSender.EXPECT().queueControlFrame(gomock.Any())
   798  				mockSender.EXPECT().onHasStreamData(streamID)
   799  				mockFC.EXPECT().SendWindowSize().Return(protocol.MaxByteCount)
   800  				mockFC.EXPECT().AddBytesSent(gomock.Any())
   801  				writeReturned := make(chan struct{})
   802  				go func() {
   803  					defer GinkgoRecover()
   804  					_, err := strWithTimeout.Write(getData(5000))
   805  					Expect(err).To(Equal(&StreamError{
   806  						StreamID:  streamID,
   807  						ErrorCode: 1234,
   808  						Remote:    false,
   809  					}))
   810  					close(writeReturned)
   811  				}()
   812  				waitForWrite()
   813  				frame, ok, hasMoreData := str.popStreamFrame(50, protocol.Version1)
   814  				Expect(ok).To(BeTrue())
   815  				Expect(hasMoreData).To(BeTrue())
   816  				Expect(frame).ToNot(BeNil())
   817  				mockSender.EXPECT().onStreamCompleted(streamID)
   818  				str.CancelWrite(1234)
   819  				_, ok, hasMoreData = str.popStreamFrame(10, protocol.Version1)
   820  				Expect(ok).To(BeFalse())
   821  				Expect(hasMoreData).To(BeFalse())
   822  				Eventually(writeReturned).Should(BeClosed())
   823  			})
   824  
   825  			It("ignores acknowledgements for STREAM frames after it was cancelled", func() {
   826  				mockSender.EXPECT().queueControlFrame(gomock.Any())
   827  				mockSender.EXPECT().onHasStreamData(streamID)
   828  				mockFC.EXPECT().SendWindowSize().Return(protocol.MaxByteCount)
   829  				mockFC.EXPECT().AddBytesSent(gomock.Any())
   830  				writeReturned := make(chan struct{})
   831  				go func() {
   832  					defer GinkgoRecover()
   833  					strWithTimeout.Write(getData(100))
   834  					close(writeReturned)
   835  				}()
   836  				waitForWrite()
   837  				frame, ok, hasMoreData := str.popStreamFrame(50, protocol.Version1)
   838  				Expect(ok).To(BeTrue())
   839  				Expect(hasMoreData).To(BeTrue())
   840  				Expect(frame).ToNot(BeNil())
   841  				mockSender.EXPECT().onStreamCompleted(streamID)
   842  				str.CancelWrite(1234)
   843  				frame.Handler.OnAcked(frame.Frame)
   844  			})
   845  
   846  			It("cancels the context", func() {
   847  				mockSender.EXPECT().queueControlFrame(gomock.Any())
   848  				mockSender.EXPECT().onStreamCompleted(gomock.Any())
   849  				Expect(str.Context().Done()).ToNot(BeClosed())
   850  				str.CancelWrite(1234)
   851  				Expect(str.Context().Done()).To(BeClosed())
   852  				Expect(context.Cause(str.Context())).To(BeAssignableToTypeOf(&StreamError{}))
   853  				Expect(context.Cause(str.Context()).(*StreamError).ErrorCode).To(Equal(StreamErrorCode(1234)))
   854  			})
   855  
   856  			It("doesn't allow further calls to Write", func() {
   857  				mockSender.EXPECT().queueControlFrame(gomock.Any())
   858  				mockSender.EXPECT().onStreamCompleted(gomock.Any())
   859  				str.CancelWrite(1234)
   860  				_, err := strWithTimeout.Write([]byte("foobar"))
   861  				Expect(err).To(MatchError(&StreamError{
   862  					StreamID:  streamID,
   863  					ErrorCode: 1234,
   864  					Remote:    false,
   865  				}))
   866  			})
   867  
   868  			It("only cancels once", func() {
   869  				mockSender.EXPECT().queueControlFrame(&wire.ResetStreamFrame{StreamID: streamID, ErrorCode: 1234})
   870  				mockSender.EXPECT().onStreamCompleted(gomock.Any())
   871  				str.CancelWrite(1234)
   872  				str.CancelWrite(4321)
   873  			})
   874  
   875  			It("queues a RESET_STREAM frame, even if the stream was already closed", func() {
   876  				mockSender.EXPECT().onHasStreamData(streamID)
   877  				mockSender.EXPECT().queueControlFrame(gomock.Any()).Do(func(f wire.Frame) {
   878  					Expect(f).To(BeAssignableToTypeOf(&wire.ResetStreamFrame{}))
   879  				})
   880  				mockSender.EXPECT().onStreamCompleted(gomock.Any())
   881  				Expect(str.Close()).To(Succeed())
   882  				// don't EXPECT any calls to queueControlFrame
   883  				str.CancelWrite(123)
   884  			})
   885  		})
   886  
   887  		Context("receiving STOP_SENDING frames", func() {
   888  			It("queues a RESET_STREAM frames, and copies the error code from the STOP_SENDING frame", func() {
   889  				mockSender.EXPECT().queueControlFrame(&wire.ResetStreamFrame{
   890  					StreamID:  streamID,
   891  					ErrorCode: 101,
   892  				})
   893  				mockSender.EXPECT().onStreamCompleted(gomock.Any())
   894  
   895  				str.handleStopSendingFrame(&wire.StopSendingFrame{
   896  					StreamID:  streamID,
   897  					ErrorCode: 101,
   898  				})
   899  			})
   900  
   901  			It("unblocks Write", func() {
   902  				mockSender.EXPECT().onHasStreamData(streamID)
   903  				mockSender.EXPECT().queueControlFrame(gomock.Any())
   904  				mockSender.EXPECT().onStreamCompleted(gomock.Any())
   905  				done := make(chan struct{})
   906  				go func() {
   907  					defer GinkgoRecover()
   908  					_, err := str.Write(getData(5000))
   909  					Expect(err).To(Equal(&StreamError{
   910  						StreamID:  streamID,
   911  						ErrorCode: 123,
   912  						Remote:    true,
   913  					}))
   914  					close(done)
   915  				}()
   916  				waitForWrite()
   917  				str.handleStopSendingFrame(&wire.StopSendingFrame{
   918  					StreamID:  streamID,
   919  					ErrorCode: 123,
   920  				})
   921  				Eventually(done).Should(BeClosed())
   922  			})
   923  
   924  			It("doesn't allow further calls to Write", func() {
   925  				mockSender.EXPECT().queueControlFrame(gomock.Any())
   926  				mockSender.EXPECT().onStreamCompleted(gomock.Any())
   927  				str.handleStopSendingFrame(&wire.StopSendingFrame{
   928  					StreamID:  streamID,
   929  					ErrorCode: 123,
   930  				})
   931  				_, err := str.Write([]byte("foobar"))
   932  				Expect(err).To(Equal(&StreamError{
   933  					StreamID:  streamID,
   934  					ErrorCode: 123,
   935  					Remote:    true,
   936  				}))
   937  			})
   938  		})
   939  	})
   940  
   941  	Context("retransmissions", func() {
   942  		It("queues and retrieves frames", func() {
   943  			str.numOutstandingFrames = 1
   944  			f := &wire.StreamFrame{
   945  				Data:           []byte("foobar"),
   946  				Offset:         0x42,
   947  				DataLenPresent: false,
   948  			}
   949  			mockSender.EXPECT().onHasStreamData(streamID)
   950  			(*sendStreamAckHandler)(str).OnLost(f)
   951  			frame, ok, _ := str.popStreamFrame(protocol.MaxByteCount, protocol.Version1)
   952  			Expect(ok).To(BeTrue())
   953  			Expect(frame).ToNot(BeNil())
   954  			f = frame.Frame
   955  			Expect(f.Offset).To(Equal(protocol.ByteCount(0x42)))
   956  			Expect(f.Data).To(Equal([]byte("foobar")))
   957  			Expect(f.DataLenPresent).To(BeTrue())
   958  		})
   959  
   960  		It("splits a retransmission", func() {
   961  			str.numOutstandingFrames = 1
   962  			sf := &wire.StreamFrame{
   963  				Data:           []byte("foobar"),
   964  				Offset:         0x42,
   965  				DataLenPresent: false,
   966  			}
   967  			mockSender.EXPECT().onHasStreamData(streamID)
   968  			(*sendStreamAckHandler)(str).OnLost(sf)
   969  			frame, ok, hasMoreData := str.popStreamFrame(sf.Length(protocol.Version1)-3, protocol.Version1)
   970  			Expect(ok).To(BeTrue())
   971  			Expect(frame).ToNot(BeNil())
   972  			f := frame.Frame
   973  			Expect(hasMoreData).To(BeTrue())
   974  			Expect(f.Offset).To(Equal(protocol.ByteCount(0x42)))
   975  			Expect(f.Data).To(Equal([]byte("foo")))
   976  			Expect(f.DataLenPresent).To(BeTrue())
   977  			frame, ok, _ = str.popStreamFrame(protocol.MaxByteCount, protocol.Version1)
   978  			Expect(ok).To(BeTrue())
   979  			Expect(frame).ToNot(BeNil())
   980  			f = frame.Frame
   981  			Expect(f.Offset).To(Equal(protocol.ByteCount(0x45)))
   982  			Expect(f.Data).To(Equal([]byte("bar")))
   983  			Expect(f.DataLenPresent).To(BeTrue())
   984  		})
   985  
   986  		It("returns nil if the size is too small", func() {
   987  			str.numOutstandingFrames = 1
   988  			f := &wire.StreamFrame{
   989  				Data:           []byte("foobar"),
   990  				Offset:         0x42,
   991  				DataLenPresent: false,
   992  			}
   993  			mockSender.EXPECT().onHasStreamData(streamID)
   994  			(*sendStreamAckHandler)(str).OnLost(f)
   995  			_, ok, hasMoreData := str.popStreamFrame(2, protocol.Version1)
   996  			Expect(ok).To(BeFalse())
   997  			Expect(hasMoreData).To(BeTrue())
   998  		})
   999  
  1000  		It("queues lost STREAM frames", func() {
  1001  			mockSender.EXPECT().onHasStreamData(streamID)
  1002  			mockFC.EXPECT().SendWindowSize().Return(protocol.ByteCount(9999))
  1003  			mockFC.EXPECT().AddBytesSent(protocol.ByteCount(6))
  1004  			done := make(chan struct{})
  1005  			go func() {
  1006  				defer GinkgoRecover()
  1007  				_, err := strWithTimeout.Write([]byte("foobar"))
  1008  				Expect(err).ToNot(HaveOccurred())
  1009  				close(done)
  1010  			}()
  1011  			waitForWrite()
  1012  			frame, ok, _ := str.popStreamFrame(protocol.MaxByteCount, protocol.Version1)
  1013  			Expect(ok).To(BeTrue())
  1014  			Eventually(done).Should(BeClosed())
  1015  			Expect(frame).ToNot(BeNil())
  1016  			Expect(frame.Frame.Data).To(Equal([]byte("foobar")))
  1017  
  1018  			// now lose the frame
  1019  			mockSender.EXPECT().onHasStreamData(streamID)
  1020  			frame.Handler.OnLost(frame.Frame)
  1021  			newFrame, ok, _ := str.popStreamFrame(protocol.MaxByteCount, protocol.Version1)
  1022  			Expect(ok).To(BeTrue())
  1023  			Expect(newFrame).ToNot(BeNil())
  1024  			Expect(newFrame.Frame.Data).To(Equal([]byte("foobar")))
  1025  		})
  1026  
  1027  		It("doesn't queue retransmissions for a stream that was canceled", func() {
  1028  			mockSender.EXPECT().onHasStreamData(streamID)
  1029  			mockFC.EXPECT().SendWindowSize().Return(protocol.MaxByteCount)
  1030  			mockFC.EXPECT().AddBytesSent(protocol.ByteCount(6))
  1031  			done := make(chan struct{})
  1032  			go func() {
  1033  				defer GinkgoRecover()
  1034  				_, err := str.Write([]byte("foobar"))
  1035  				Expect(err).ToNot(HaveOccurred())
  1036  				close(done)
  1037  			}()
  1038  			waitForWrite()
  1039  			f, ok, _ := str.popStreamFrame(100, protocol.Version1)
  1040  			Expect(ok).To(BeTrue())
  1041  			Eventually(done).Should(BeClosed())
  1042  			Expect(f).ToNot(BeNil())
  1043  			gomock.InOrder(
  1044  				mockSender.EXPECT().queueControlFrame(gomock.Any()),
  1045  				mockSender.EXPECT().onStreamCompleted(streamID),
  1046  			)
  1047  			str.CancelWrite(9876)
  1048  			// don't EXPECT any calls to onHasStreamData
  1049  			f.Handler.OnLost(f.Frame)
  1050  			Expect(str.retransmissionQueue).To(BeEmpty())
  1051  		})
  1052  	})
  1053  
  1054  	Context("determining when a stream is completed", func() {
  1055  		BeforeEach(func() {
  1056  			mockFC.EXPECT().SendWindowSize().Return(protocol.MaxByteCount).AnyTimes()
  1057  			mockFC.EXPECT().AddBytesSent(gomock.Any()).AnyTimes()
  1058  		})
  1059  
  1060  		It("says when a stream is completed", func() {
  1061  			mockSender.EXPECT().onHasStreamData(streamID)
  1062  			done := make(chan struct{})
  1063  			go func() {
  1064  				defer GinkgoRecover()
  1065  				_, err := strWithTimeout.Write(make([]byte, 100))
  1066  				Expect(err).ToNot(HaveOccurred())
  1067  				close(done)
  1068  			}()
  1069  			waitForWrite()
  1070  
  1071  			// get a bunch of small frames (max. 20 bytes)
  1072  			var frames []ackhandler.StreamFrame
  1073  			for {
  1074  				frame, ok, hasMoreData := str.popStreamFrame(20, protocol.Version1)
  1075  				if !ok {
  1076  					continue
  1077  				}
  1078  				frames = append(frames, frame)
  1079  				if !hasMoreData {
  1080  					break
  1081  				}
  1082  			}
  1083  			Eventually(done).Should(BeClosed())
  1084  
  1085  			// Acknowledge all frames.
  1086  			// We don't expect the stream to be completed, since we still need to send the FIN.
  1087  			for _, f := range frames {
  1088  				f.Handler.OnAcked(f.Frame)
  1089  			}
  1090  
  1091  			// Now close the stream and acknowledge the FIN.
  1092  			mockSender.EXPECT().onHasStreamData(streamID)
  1093  			Expect(str.Close()).To(Succeed())
  1094  			frame, ok, _ := str.popStreamFrame(protocol.MaxByteCount, protocol.Version1)
  1095  			Expect(ok).To(BeTrue())
  1096  			Expect(frame).ToNot(BeNil())
  1097  			mockSender.EXPECT().onStreamCompleted(streamID)
  1098  			frame.Handler.OnAcked(frame.Frame)
  1099  		})
  1100  
  1101  		It("says when a stream is completed, if Close() is called before popping the frame", func() {
  1102  			mockSender.EXPECT().onHasStreamData(streamID).Times(2)
  1103  			done := make(chan struct{})
  1104  			go func() {
  1105  				defer GinkgoRecover()
  1106  				_, err := strWithTimeout.Write(make([]byte, 100))
  1107  				Expect(err).ToNot(HaveOccurred())
  1108  				close(done)
  1109  			}()
  1110  			waitForWrite()
  1111  			Eventually(done).Should(BeClosed())
  1112  			Expect(str.Close()).To(Succeed())
  1113  
  1114  			frame, ok, hasMoreData := str.popStreamFrame(protocol.MaxByteCount, protocol.Version1)
  1115  			Expect(ok).To(BeTrue())
  1116  			Expect(hasMoreData).To(BeFalse())
  1117  			Expect(frame).ToNot(BeNil())
  1118  			Expect(frame.Frame.Fin).To(BeTrue())
  1119  
  1120  			mockSender.EXPECT().onStreamCompleted(streamID)
  1121  			frame.Handler.OnAcked(frame.Frame)
  1122  		})
  1123  
  1124  		It("doesn't say it's completed when there are frames waiting to be retransmitted", func() {
  1125  			mockSender.EXPECT().onHasStreamData(streamID)
  1126  			done := make(chan struct{})
  1127  			go func() {
  1128  				defer GinkgoRecover()
  1129  				_, err := strWithTimeout.Write(getData(100))
  1130  				Expect(err).ToNot(HaveOccurred())
  1131  				mockSender.EXPECT().onHasStreamData(streamID)
  1132  				Expect(str.Close()).To(Succeed())
  1133  				close(done)
  1134  			}()
  1135  			waitForWrite()
  1136  
  1137  			// get a bunch of small frames (max. 20 bytes)
  1138  			var frames []ackhandler.StreamFrame
  1139  			for {
  1140  				frame, ok, _ := str.popStreamFrame(20, protocol.Version1)
  1141  				if !ok {
  1142  					continue
  1143  				}
  1144  				frames = append(frames, frame)
  1145  				if frame.Frame.Fin {
  1146  					break
  1147  				}
  1148  			}
  1149  			Eventually(done).Should(BeClosed())
  1150  
  1151  			// lose the first frame, acknowledge all others
  1152  			for _, f := range frames[1:] {
  1153  				f.Handler.OnAcked(f.Frame)
  1154  			}
  1155  			mockSender.EXPECT().onHasStreamData(streamID)
  1156  			frames[0].Handler.OnLost(frames[0].Frame)
  1157  
  1158  			// get the retransmission and acknowledge it
  1159  			ret, ok, _ := str.popStreamFrame(protocol.MaxByteCount, protocol.Version1)
  1160  			Expect(ok).To(BeTrue())
  1161  			Expect(ret).ToNot(BeNil())
  1162  			mockSender.EXPECT().onStreamCompleted(streamID)
  1163  			ret.Handler.OnAcked(ret.Frame)
  1164  		})
  1165  
  1166  		// This test is kind of an integration test.
  1167  		// It writes 4 MB of data, and pops STREAM frames that sometimes are and sometimes aren't limited by flow control.
  1168  		// Half of these STREAM frames are then received and their content saved, while the other half is reported lost
  1169  		// and has to be retransmitted.
  1170  		It("retransmits data until everything has been acknowledged", func() {
  1171  			const dataLen = 1 << 22 // 4 MB
  1172  			mockSender.EXPECT().onHasStreamData(streamID).AnyTimes()
  1173  			mockFC.EXPECT().SendWindowSize().DoAndReturn(func() protocol.ByteCount {
  1174  				return protocol.ByteCount(mrand.Intn(500)) + 50
  1175  			}).AnyTimes()
  1176  			mockFC.EXPECT().AddBytesSent(gomock.Any()).AnyTimes()
  1177  
  1178  			data := make([]byte, dataLen)
  1179  			_, err := rand.Read(data)
  1180  			Expect(err).ToNot(HaveOccurred())
  1181  			done := make(chan struct{})
  1182  			go func() {
  1183  				defer GinkgoRecover()
  1184  				defer close(done)
  1185  				_, err := str.Write(data)
  1186  				Expect(err).ToNot(HaveOccurred())
  1187  				str.Close()
  1188  			}()
  1189  
  1190  			var completed bool
  1191  			mockSender.EXPECT().onStreamCompleted(streamID).Do(func(protocol.StreamID) { completed = true })
  1192  
  1193  			received := make([]byte, dataLen)
  1194  			for {
  1195  				if completed {
  1196  					break
  1197  				}
  1198  				f, ok, _ := str.popStreamFrame(protocol.ByteCount(mrand.Intn(300)+100), protocol.Version1)
  1199  				if !ok {
  1200  					continue
  1201  				}
  1202  				sf := f.Frame
  1203  				// 50%: acknowledge the frame and save the data
  1204  				// 50%: lose the frame
  1205  				if mrand.Intn(100) < 50 {
  1206  					copy(received[sf.Offset:sf.Offset+sf.DataLen()], sf.Data)
  1207  					f.Handler.OnAcked(f.Frame)
  1208  				} else {
  1209  					f.Handler.OnLost(f.Frame)
  1210  				}
  1211  			}
  1212  			Expect(received).To(Equal(data))
  1213  		})
  1214  	})
  1215  })