mpegts: add pes payload encoding, fix h264 nalbuffer append

mpegts/h264.go

@@ -62,7 +62,7 @@ func (h *H264Parser) InitPacket() ([]byte, error) {
 	copy(pesData[len(h.sps):], h.pps)
 
 	// encode PES payload
-	pesPayload, err := encodePES(pesData)
+	pesPayload, err := encodeVideoPES(pesData)
 	if err != nil {
 		return nil, err
 	}
@@ -249,7 +249,7 @@ func (h *H264Parser) Parse(rd *io.PipeReader) {
 	// read remaining bytes
 	buffered := brd.Buffered()
 	if buffered > 0 {
-		tmp := make([]byte, 0, buffered)
+		tmp := make([]byte, buffered)
 		n, _ := brd.Read(tmp)
 		nalBuffer = append(nalBuffer, tmp[:n]...)
 	}

mpegts/parser_test.go

@@ -26,7 +26,7 @@ func TestParser_ParseH264(t *testing.T) {
 		t.Error("Init should not be nil")
 	}
 	numPackets := len(pkts)
-	if numPackets != 20 {
-		t.Errorf("Expected 3 init packets, got %d", numPackets)
+	if numPackets != 32 {
+		t.Errorf("Expected 32 init packets, got %d", numPackets)
 	}
 }

mpegts/pes.go

@@ -1,6 +1,7 @@
 package mpegts
 
 import (
+	"encoding/binary"
 	"io"
 	"log"
 )
@@ -9,6 +10,7 @@ import (
 const (
 	PESStartCode   = 0x000001
 	MaxPayloadSize = PacketLen - 4
+	PESHeaderSize  = 6
 	PESMaxLength   = 200 * 1024
 
 	PESStreamIDAudio = 0xc0
@@ -26,13 +28,27 @@ type CodecParser interface {
 	Parse(*io.PipeReader)
 }
 
-func encodePES(data []byte) ([]byte, error) {
+// Encode video codec PES with custom payload
+func encodeVideoPES(data []byte) ([]byte, error) {
 	// Early check, but MPEG-TS packet will check again
-	if len(data)+9+5 > PacketLen-HeaderLen {
+	if len(data)+PESHeaderSize > PacketLen-HeaderLen {
 		return nil, ErrDataTooLong
 	}
-	// len = len + 9 + 5;//pes len
-	return nil, nil
+
+	pes := make([]byte, PESHeaderSize, MaxPayloadSize)
+
+	// write pes header (Video only for now)
+	offset := 0
+	tmp := uint32(PESStartCode<<8) | PESStreamIDVideo
+	binary.BigEndian.PutUint32(pes[offset:offset+4], tmp)
+	offset += 4
+
+	// write pes length (can be 0 only for Video packets)
+	binary.BigEndian.PutUint16(pes[offset:offset+2], 0)
+
+	pes = append(pes, data...)
+
+	return pes, nil
 }
 
 // Parse PES packet

mpegts/pes_test.go

@@ -0,0 +1,72 @@
+package mpegts
+
+import (
+	"bytes"
+	"io"
+	"testing"
+)
+
+// Mock codec storing all parsed data
+type mockCodec struct {
+	data []byte
+	done chan struct{}
+	err  error
+}
+
+func (m *mockCodec) HasInit() bool {
+	return false
+}
+
+func (m *mockCodec) InitPacket() ([]byte, error) {
+	return []byte{}, nil
+}
+
+func (m *mockCodec) Parse(rd *io.PipeReader) {
+	tmp := make([]byte, MaxPayloadSize)
+	var n int
+	n, m.err = rd.Read(tmp)
+	m.data = append(m.data, tmp[:n]...)
+	rd.Close()
+	m.done <- struct{}{}
+}
+
+func (m *mockCodec) Data() ([]byte, error) {
+	<-m.done
+	return m.data, m.err
+}
+
+// Test PES encoding against our PES parser
+func TestPES_encodeVideoPES(t *testing.T) {
+	parser := &mockCodec{
+		done: make(chan struct{}),
+	}
+	es := ElementaryStream{
+		parser: parser,
+	}
+
+	expected := []byte{1, 2, 3, 4}
+	payload, err := encodeVideoPES(expected)
+
+	if err != nil {
+		t.Fatal("Encode failed", err)
+	}
+
+	pkt := Packet{
+		PID:     256,
+		PUSI:    true,
+		Payload: payload,
+	}
+
+	err = es.ParsePES(&pkt)
+	if err != nil {
+		t.Fatal("Parse failed", err)
+	}
+
+	res, err := parser.Data()
+	if err != nil {
+		t.Fatal("Mock Parser failed", err)
+	}
+	if bytes.Compare(res, expected) != 0 {
+		t.Errorf("Got payload %v, expected %v", res, expected)
+	}
+}