use channels to properly implement the Stream.Read function

stream.go

@@ -1,60 +1,54 @@
 package quic
 
 import (
-	"errors"
-	"io"
-
 	"github.com/lucas-clemente/quic-go/frames"
+	"github.com/lucas-clemente/quic-go/utils"
 )
 
 // A Stream assembles the data from StreamFrames and provides a super-convenient Read-Interface
 type Stream struct {
-	StreamFrames   []*frames.StreamFrame
-	DataLen        uint64
-	ReadPosFrameNo int
+	StreamFrames   chan *frames.StreamFrame
+	CurrentFrame   *frames.StreamFrame
 	ReadPosInFrame int
 }
 
 // NewStream creates a new Stream
 func NewStream() *Stream {
-	return &Stream{}
+	return &Stream{
+		StreamFrames: make(chan *frames.StreamFrame, 8), // ToDo: add config option for this number
 	}
-
-func (s *Stream) readByte() (byte, error) {
-	if s.ReadPosInFrame == len(s.StreamFrames[s.ReadPosFrameNo].Data) {
-		s.ReadPosFrameNo++
-		if s.ReadPosFrameNo == len(s.StreamFrames) {
-			return 0, io.EOF
-		}
-		s.ReadPosInFrame = 0
-	}
-	b := s.StreamFrames[s.ReadPosFrameNo].Data[s.ReadPosInFrame]
-	s.ReadPosInFrame++
-	return b, nil
 }
 
 // Read reads data
 func (s *Stream) Read(p []byte) (int, error) {
-	var err error
-	n := 0
-	if c := cap(p); c > 0 {
-		for n < c {
-			p[n], err = s.readByte()
-			n++
-			if err != nil {
-				break
+	bytesRead := 0
+	for bytesRead < len(p) {
+		if s.CurrentFrame == nil {
+			select {
+			case s.CurrentFrame = <-s.StreamFrames:
+			default:
+				if bytesRead == 0 {
+					s.CurrentFrame = <-s.StreamFrames
+				} else {
+					return bytesRead, nil
 				}
 			}
+			s.ReadPosInFrame = 0
 		}
-	return n, nil
+		m := utils.Min(len(p)-bytesRead, len(s.CurrentFrame.Data)-s.ReadPosInFrame)
+		copy(p[bytesRead:], s.CurrentFrame.Data[s.ReadPosInFrame:])
+		s.ReadPosInFrame += m
+		bytesRead += m
+		if s.ReadPosInFrame >= len(s.CurrentFrame.Data) {
+			s.CurrentFrame = nil
+		}
+	}
+
+	return bytesRead, nil
 }
 
 // AddStreamFrame adds a new stream frame
 func (s *Stream) AddStreamFrame(frame *frames.StreamFrame) error {
-	if frame.Offset != s.DataLen {
-		return errors.New("Stream: Wrong offset")
-	}
-	s.StreamFrames = append(s.StreamFrames, frame)
-	s.DataLen += uint64(len(frame.Data))
+	s.StreamFrames <- frame
 	return nil
 }

stream_test.go

@@ -1,6 +1,8 @@
 package quic
 
 import (
+	"time"
+
 	"github.com/lucas-clemente/quic-go/frames"
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
@@ -18,10 +20,28 @@ var _ = Describe("Stream", func() {
 		n, err := stream.Read(b)
 		Expect(err).ToNot(HaveOccurred())
 		Expect(n).To(Equal(4))
-		Expect(stream.DataLen).To(Equal(uint64(4)))
 		Expect(b).To(Equal([]byte{0xDE, 0xAD, 0xBE, 0xEF}))
 	})
 
+	It("reads all data available", func() {
+		frame1 := frames.StreamFrame{
+			Offset: 0,
+			Data:   []byte{0xDE, 0xAD},
+		}
+		frame2 := frames.StreamFrame{
+			Offset: 2,
+			Data:   []byte{0xBE, 0xEF},
+		}
+		stream := NewStream()
+		stream.AddStreamFrame(&frame1)
+		stream.AddStreamFrame(&frame2)
+		b := make([]byte, 6)
+		n, err := stream.Read(b)
+		Expect(err).ToNot(HaveOccurred())
+		Expect(n).To(Equal(4))
+		Expect(b).To(Equal([]byte{0xDE, 0xAD, 0xBE, 0xEF, 0x00, 0x00}))
+	})
+
 	It("assembles multiple StreamFrames", func() {
 		frame1 := frames.StreamFrame{
 			Offset: 0,
@@ -38,11 +58,26 @@ var _ = Describe("Stream", func() {
 		n, err := stream.Read(b)
 		Expect(err).ToNot(HaveOccurred())
 		Expect(n).To(Equal(4))
-		Expect(stream.DataLen).To(Equal(uint64(4)))
 		Expect(b).To(Equal([]byte{0xDE, 0xAD, 0xBE, 0xEF}))
 	})
 
-	It("rejects StreamFrames with wrong Offsets", func() {
+	It("waits until data is available", func() {
+		stream := NewStream()
+		go func() {
+			frame := frames.StreamFrame{
+				Offset: 0,
+				Data:   []byte{0xDE, 0xAD},
+			}
+			time.Sleep(time.Millisecond)
+			stream.AddStreamFrame(&frame)
+		}()
+		b := make([]byte, 2)
+		n, err := stream.Read(b)
+		Expect(err).ToNot(HaveOccurred())
+		Expect(n).To(Equal(2))
+	})
+
+	PIt("rejects StreamFrames with wrong Offsets", func() {
 		frame1 := frames.StreamFrame{
 			Offset: 0,
 			Data:   []byte{0xDE, 0xAD},

utils/utils.go

@@ -77,6 +77,22 @@ func WriteUint16(b *bytes.Buffer, i uint16) {
 	b.WriteByte(uint8((i >> 8) & 0xff))
 }
 
+// Max returns the maximum of two Ints
+func Max(a, b int) int {
+	if a < b {
+		return b
+	}
+	return a
+}
+
+// Min returns the minimum of two Ints
+func Min(a, b int) int {
+	if a < b {
+		return a
+	}
+	return b
+}
+
 // Uint32Slice attaches the methods of sort.Interface to []uint32, sorting in increasing order.
 type Uint32Slice []uint32
 

utils/utils_test.go

@@ -82,4 +82,16 @@ var _ = Describe("Utils", func() {
 			Expect(b.Bytes()).To(Equal([]byte{0x12, 0x35, 0xAC, 0xEF}))
 		})
 	})
+
+	Context("Max", func() {
+		It("returns the maximum", func() {
+			Expect(Max(5, 7)).To(Equal(7))
+		})
+	})
+
+	Context("Min", func() {
+		It("returns the minimum", func() {
+			Expect(Min(5, 7)).To(Equal(5))
+		})
+	})
 })