mpegts: add Reader, rewrite Writer

pkg/codecs/h264/annexb.go

@@ -4,7 +4,7 @@
 	"fmt"
 )
 
-// AnnexBUnmarshal decodes NALUs from the Annex-B stream format.
+// AnnexBUnmarshal decodes an access unit from the Annex-B stream format.
 // Specification: ITU-T Rec. H.264, Annex B
 func AnnexBUnmarshal(byts []byte) ([][]byte, error) {
 	bl := len(byts)
@@ -60,9 +60,9 @@
 		}
 	}
 
-	if (n + 1) > MaxNALUsPerGroup {
+	if (n + 1) > MaxNALUsPerAccessUnit {
 		return nil, fmt.Errorf("NALU count (%d) exceeds maximum allowed (%d)",
-			n+1, MaxNALUsPerGroup)
+			n+1, MaxNALUsPerAccessUnit)
 	}
 
 	ret := make([][]byte, n+1)
@@ -113,21 +113,21 @@
 	return ret, nil
 }
 
-func annexBMarshalSize(nalus [][]byte) int {
+func annexBMarshalSize(au [][]byte) int {
 	n := 0
-	for _, nalu := range nalus {
+	for _, nalu := range au {
 		n += 4 + len(nalu)
 	}
 	return n
 }
 
-// AnnexBMarshal encodes NALUs into the Annex-B stream format.
+// AnnexBMarshal encodes an access unit into the Annex-B stream format.
 // Specification: ITU-T Rec. H.264, Annex B
-func AnnexBMarshal(nalus [][]byte) ([]byte, error) {
-	buf := make([]byte, annexBMarshalSize(nalus))
+func AnnexBMarshal(au [][]byte) ([]byte, error) {
+	buf := make([]byte, annexBMarshalSize(au))
 	pos := 0
 
-	for _, nalu := range nalus {
+	for _, nalu := range au {
 		pos += copy(buf[pos:], []byte{0x00, 0x00, 0x00, 0x01})
 		pos += copy(buf[pos:], nalu)
 	}

pkg/codecs/h264/avcc.go

@@ -4,7 +4,7 @@
 	"fmt"
 )
 
-// AVCCUnmarshal decodes NALUs from the AVCC stream format.
+// AVCCUnmarshal decodes an access unit from the AVCC stream format.
 // Specification: ?
 func AVCCUnmarshal(buf []byte) ([][]byte, error) {
 	bl := len(buf)
@@ -27,9 +27,9 @@
 			return nil, fmt.Errorf("NALU size (%d) is too big, maximum is %d", l, MaxNALUSize)
 		}
 
-		if (len(ret) + 1) > MaxNALUsPerGroup {
+		if (len(ret) + 1) > MaxNALUsPerAccessUnit {
 			return nil, fmt.Errorf("NALU count (%d) exceeds maximum allowed (%d)",
-				len(ret)+1, MaxNALUsPerGroup)
+				len(ret)+1, MaxNALUsPerAccessUnit)
 		}
 
 		if (bl - pos) < l {
@@ -47,21 +47,21 @@
 	return ret, nil
 }
 
-func avccMarshalSize(nalus [][]byte) int {
+func avccMarshalSize(au [][]byte) int {
 	n := 0
-	for _, nalu := range nalus {
+	for _, nalu := range au {
 		n += 4 + len(nalu)
 	}
 	return n
 }
 
-// AVCCMarshal encodes NALUs into the AVCC stream format.
+// AVCCMarshal encodes an access unit into the AVCC stream format.
 // Specification: ?
-func AVCCMarshal(nalus [][]byte) ([]byte, error) {
-	buf := make([]byte, avccMarshalSize(nalus))
+func AVCCMarshal(au [][]byte) ([]byte, error) {
+	buf := make([]byte, avccMarshalSize(au))
 	pos := 0
 
-	for _, nalu := range nalus {
+	for _, nalu := range au {
 		naluLen := len(nalu)
 		buf[pos] = byte(naluLen >> 24)
 		buf[pos+1] = byte(naluLen >> 16)

pkg/codecs/h264/dts_extractor_test.go

@@ -9,9 +9,9 @@ import (
 
 func TestDTSExtractor(t *testing.T) {
 	type sequenceSample struct {
-		nalus [][]byte
-		dts   time.Duration
-		pts   time.Duration
+		au  [][]byte
+		dts time.Duration
+		pts time.Duration
 	}
 
 	for _, ca := range []struct {
@@ -201,7 +201,7 @@ func TestDTSExtractor(t *testing.T) {
 		t.Run(ca.name, func(t *testing.T) {
 			ex := NewDTSExtractor()
 			for _, sample := range ca.sequence {
-				dts, err := ex.Extract(sample.nalus, sample.pts)
+				dts, err := ex.Extract(sample.au, sample.pts)
 				require.NoError(t, err)
 				require.Equal(t, sample.dts, dts)
 			}

pkg/codecs/h264/h264.go

@@ -6,6 +6,6 @@ const (
 	// with a 250 Mbps H264 video, the maximum NALU size is 2.2MB
 	MaxNALUSize = 3 * 1024 * 1024
 
-	// MaxNALUsPerGroup is the maximum number of NALUs per group.
-	MaxNALUsPerGroup = 20
+	// MaxNALUsPerAccessUnit is the maximum number of NALUs per access unit.
+	MaxNALUsPerAccessUnit = 20
 )

pkg/codecs/h264/idrpresent.go

@@ -1,8 +1,8 @@
 package h264
 
-// IDRPresent check whether there's an IDR inside provided NALUs.
-func IDRPresent(nalus [][]byte) bool {
-	for _, nalu := range nalus {
+// IDRPresent check whether there's an IDR inside the access unit.
+func IDRPresent(au [][]byte) bool {
+	for _, nalu := range au {
 		typ := NALUType(nalu[0] & 0x1F)
 		if typ == NALUTypeIDR {
 			return true

pkg/codecs/h265/dts_extractor_test.go

@@ -9,9 +9,9 @@ import (
 
 func TestDTSExtractor(t *testing.T) {
 	type sequenceSample struct {
-		nalus [][]byte
-		pts   time.Duration
-		dts   time.Duration
+		au  [][]byte
+		pts time.Duration
+		dts time.Duration
 	}
 
 	for _, ca := range []struct {
@@ -49,7 +49,7 @@ func TestDTSExtractor(t *testing.T) {
 		t.Run(ca.name, func(t *testing.T) {
 			ex := NewDTSExtractor()
 			for _, sample := range ca.sequence {
-				dts, err := ex.Extract(sample.nalus, sample.pts)
+				dts, err := ex.Extract(sample.au, sample.pts)
 				require.NoError(t, err)
 				require.Equal(t, sample.dts, dts)
 			}

pkg/codecs/h265/h265.go

@@ -6,6 +6,6 @@ const (
 	// with a 250 Mbps H265 video, the maximum NALU size is 2.2MB
 	MaxNALUSize = 3 * 1024 * 1024
 
-	// MaxNALUsPerGroup is the maximum number of NALUs per group.
-	MaxNALUsPerGroup = 20
+	// MaxNALUsPerAccessUnit is the maximum number of NALUs per access unit.
+	MaxNALUsPerAccessUnit = 20
 )

pkg/formats/mpegts/codec.go

@@ -1,38 +1,11 @@
 package mpegts
 
 import (
-	"github.com/bluenviron/mediacommon/pkg/codecs/mpeg4audio"
+	"github.com/asticode/go-astits"
 )
 
 // Codec is a MPEG-TS codec.
 type Codec interface {
-	isCodec()
-}
-
-// CodecH264 is a H264 codec.
-type CodecH264 struct{}
-
-func (*CodecH264) isCodec() {
-}
-
-// CodecH265 is a H265 codec.
-type CodecH265 struct{}
-
-func (*CodecH265) isCodec() {
-}
-
-// CodecMPEG4Audio is a MPEG-4 Audio codec.
-type CodecMPEG4Audio struct {
-	mpeg4audio.Config
-}
-
-func (*CodecMPEG4Audio) isCodec() {
-}
-
-// CodecOpus is a Opus codec.
-type CodecOpus struct {
-	Channels int
-}
-
-func (*CodecOpus) isCodec() {
+	marshal(pid uint16) (*astits.PMTElementaryStream, error)
+	isVideo() bool
 }

pkg/formats/mpegts/codec_h264.go

@@ -0,0 +1,20 @@
+package mpegts
+
+import (
+	"github.com/asticode/go-astits"
+)
+
+// CodecH264 is a H264 codec.
+type CodecH264 struct{}
+
+func (c CodecH264) marshal(pid uint16) (*astits.PMTElementaryStream, error) {
+	return &astits.PMTElementaryStream{
+		ElementaryPID:               pid,
+		ElementaryStreamDescriptors: nil,
+		StreamType:                  astits.StreamTypeH264Video,
+	}, nil
+}
+
+func (CodecH264) isVideo() bool {
+	return true
+}

pkg/formats/mpegts/codec_h265.go

@@ -0,0 +1,20 @@
+package mpegts
+
+import (
+	"github.com/asticode/go-astits"
+)
+
+// CodecH265 is a H265 codec.
+type CodecH265 struct{}
+
+func (c CodecH265) marshal(pid uint16) (*astits.PMTElementaryStream, error) {
+	return &astits.PMTElementaryStream{
+		ElementaryPID:               pid,
+		ElementaryStreamDescriptors: nil,
+		StreamType:                  astits.StreamTypeH265Video,
+	}, nil
+}
+
+func (CodecH265) isVideo() bool {
+	return true
+}

pkg/formats/mpegts/codec_mpeg4audio.go

@@ -0,0 +1,24 @@
+package mpegts
+
+import (
+	"github.com/asticode/go-astits"
+
+	"github.com/bluenviron/mediacommon/pkg/codecs/mpeg4audio"
+)
+
+// CodecMPEG4Audio is a MPEG-4 Audio codec.
+type CodecMPEG4Audio struct {
+	mpeg4audio.Config
+}
+
+func (c CodecMPEG4Audio) marshal(pid uint16) (*astits.PMTElementaryStream, error) {
+	return &astits.PMTElementaryStream{
+		ElementaryPID:               pid,
+		ElementaryStreamDescriptors: nil,
+		StreamType:                  astits.StreamTypeAACAudio,
+	}, nil
+}
+
+func (CodecMPEG4Audio) isVideo() bool {
+	return false
+}

pkg/formats/mpegts/codec_opus.go

@@ -0,0 +1,38 @@
+package mpegts
+
+import (
+	"github.com/asticode/go-astits"
+)
+
+// CodecOpus is a Opus codec.
+type CodecOpus struct {
+	ChannelCount int
+}
+
+func (c CodecOpus) marshal(pid uint16) (*astits.PMTElementaryStream, error) {
+	return &astits.PMTElementaryStream{
+		ElementaryPID: pid,
+		ElementaryStreamDescriptors: []*astits.Descriptor{
+			{
+				Length: 4,
+				Tag:    astits.DescriptorTagRegistration,
+				Registration: &astits.DescriptorRegistration{
+					FormatIdentifier: opusIdentifier,
+				},
+			},
+			{
+				Length: 2,
+				Tag:    astits.DescriptorTagExtension,
+				Extension: &astits.DescriptorExtension{
+					Tag:     0x80,
+					Unknown: &[]uint8{uint8(c.ChannelCount)},
+				},
+			},
+		},
+		StreamType: astits.StreamTypePrivateData,
+	}, nil
+}
+
+func (CodecOpus) isVideo() bool {
+	return false
+}

pkg/formats/mpegts/opus_access_unit.go

@@ -4,15 +4,13 @@
 	"fmt"
 )
 
-// OpusAccessUnit is a MPEG-TS Opus access unit.
-type OpusAccessUnit struct {
-	ControlHeader OpusControlHeader
-	Frame         []byte
+type opusAccessUnit struct {
+	ControlHeader opusControlHeader
+	Packet        []byte
 }
 
-// Unmarshal decodes an access unit.
-func (au *OpusAccessUnit) Unmarshal(buf []byte) (int, error) {
-	n, err := au.ControlHeader.Unmarshal(buf)
+func (au *opusAccessUnit) unmarshal(buf []byte) (int, error) {
+	n, err := au.ControlHeader.unmarshal(buf)
 	if err != nil {
 		return 0, fmt.Errorf("could not decode Opus control header: %v", err)
 	}
@@ -22,7 +20,7 @@
 		return 0, fmt.Errorf("buffer is too small")
 	}
 
-	au.Frame = buf[:au.ControlHeader.PayloadSize]
+	au.Packet = buf[:au.ControlHeader.PayloadSize]
 
 	return n + int(au.ControlHeader.PayloadSize), nil
 }

pkg/formats/mpegts/opus_control_header.go

@@ -6,15 +6,13 @@
 	"github.com/bluenviron/mediacommon/pkg/bits"
 )
 
-// OpusControlHeader is a MPEG-TS Opus control header.
-type OpusControlHeader struct {
+type opusControlHeader struct {
 	PayloadSize uint64
 	StartTrim   uint16
 	EndTrim     uint16
 }
 
-// Unmarshal decodes a control header.
-func (h *OpusControlHeader) Unmarshal(buf []byte) (int, error) {
+func (h *opusControlHeader) unmarshal(buf []byte) (int, error) {
 	pos := 0
 
 	err := bits.HasSpace(buf, pos, 16)

pkg/formats/mpegts/playback_reader.go

@@ -0,0 +1,25 @@
+package mpegts
+
+import (
+	"io"
+)
+
+type playbackReader struct {
+	r   io.Reader
+	buf []byte
+}
+
+func (r *playbackReader) Read(p []byte) (int, error) {
+	if len(r.buf) > 0 {
+		n := copy(p, r.buf)
+		r.buf = r.buf[n:]
+
+		if len(r.buf) == 0 { // release buf from memory
+			r.buf = nil
+		}
+
+		return n, nil
+	}
+
+	return r.r.Read(p)
+}