mpegts: add buffered reader (#25)

pkg/formats/mpegts/buffered_reader.go

@@ -0,0 +1,44 @@
+package mpegts
+
+import (
+	"fmt"
+	"io"
+)
+
+// BufferedReader is a buffered reader optimized for MPEG-TS.
+type BufferedReader struct {
+	r         io.Reader
+	midbuf    []byte
+	midbufpos int
+}
+
+// NewBufferedReader allocates a BufferedReader.
+func NewBufferedReader(r io.Reader) *BufferedReader {
+	return &BufferedReader{
+		r:      r,
+		midbuf: make([]byte, 0, 1500),
+	}
+}
+
+// Read implements io.Reader.
+func (r *BufferedReader) Read(p []byte) (int, error) {
+	if r.midbufpos < len(r.midbuf) {
+		n := copy(p, r.midbuf[r.midbufpos:])
+		r.midbufpos += n
+		return n, nil
+	}
+
+	mn, err := r.r.Read(r.midbuf[:cap(r.midbuf)])
+	if err != nil {
+		return 0, err
+	}
+
+	if (mn % 188) != 0 {
+		return 0, fmt.Errorf("received packet with size %d not multiple of 188", mn)
+	}
+
+	r.midbuf = r.midbuf[:mn]
+	n := copy(p, r.midbuf)
+	r.midbufpos = n
+	return n, nil
+}

pkg/formats/mpegts/buffered_reader_test.go

@@ -0,0 +1,43 @@
+package mpegts
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestBufferedReader(t *testing.T) {
+	var buf bytes.Buffer
+
+	buf.Write(bytes.Repeat([]byte{1}, 188))
+	buf.Write(bytes.Repeat([]byte{2}, 188))
+	buf.Write(bytes.Repeat([]byte{3}, 188))
+
+	r := NewBufferedReader(&buf)
+
+	byts := make([]byte, 188)
+	n, err := r.Read(byts)
+	require.NoError(t, err)
+	require.Equal(t, 188, n)
+	require.Equal(t, bytes.Repeat([]byte{1}, 188), byts[:n])
+
+	require.Equal(t, 0, len(buf.Bytes()))
+
+	byts = make([]byte, 188)
+	n, err = r.Read(byts)
+	require.NoError(t, err)
+	require.Equal(t, 188, n)
+	require.Equal(t, bytes.Repeat([]byte{2}, 188), byts[:n])
+}
+
+func TestBufferedReaderError(t *testing.T) {
+	var buf bytes.Buffer
+
+	buf.Write(bytes.Repeat([]byte{1}, 1000))
+
+	r := NewBufferedReader(&buf)
+	byts := make([]byte, 188)
+	_, err := r.Read(byts)
+	require.EqualError(t, err, "received packet with size 1000 not multiple of 188")
+}