rtmp: improve video / audio messages

internal/rtmp/conn.go

@@ -10,6 +10,7 @@ import (
 	"github.com/aler9/gortsplib"
 	"github.com/aler9/gortsplib/pkg/aac"
 	"github.com/notedit/rtmp/av"
+	nhaac "github.com/notedit/rtmp/codec/aac"
 	nh264 "github.com/notedit/rtmp/codec/h264"
 	"github.com/notedit/rtmp/format/flv/flvio"
 	"github.com/notedit/rtmp/format/rtmp"
@@ -344,6 +345,11 @@ func (c *Conn) WriteTracks(videoTrack *gortsplib.TrackH264, audioTrack *gortspli
 
 		err = c.WritePacket(av.Packet{
 			Type: av.AACDecoderConfig,
+			AAC: &nhaac.Codec{
+				Config: nhaac.MPEG4AudioConfig{
+					ChannelLayout: nhaac.CH_STEREO,
+				},
+			},
 			Data: enc,
 		})
 		if err != nil {

internal/rtmp/conn_test.go

@@ -62,7 +62,7 @@ func TestReadTracks(t *testing.T) {
 	for _, ca := range []string{
 		"standard",
 		"metadata without codec id",
-		"no metadata",
+		"no metadata, video + audio",
 	} {
 		t.Run(ca, func(t *testing.T) {
 			ln, err := net.Listen("tcp", "127.0.0.1:9121")
@@ -100,7 +100,7 @@ func TestReadTracks(t *testing.T) {
 
 					require.Equal(t, (*gortsplib.TrackAAC)(nil), audioTrack)
 
-				case "no metadata":
+				case "no metadata, video + audio":
 					videoTrack2, err := gortsplib.NewTrackH264(96, sps, pps, nil)
 					require.NoError(t, err)
 					require.Equal(t, videoTrack2, videoTrack)
@@ -341,11 +341,12 @@ func TestReadTracks(t *testing.T) {
 				b := make([]byte, 128)
 				var n int
 				codec.ToConfig(b, &n)
-				body := append([]byte{flvio.FRAME_KEY<<4 | flvio.VIDEO_H264, 0, 0, 0, 0}, b[:n]...)
 				err = mw.Write(&message.MsgVideo{
 					ChunkStreamID:   6,
 					MessageStreamID: 1,
-					Body:            body,
+					IsKeyFrame:      true,
+					H264Type:        flvio.AVC_SEQHDR,
+					Payload:         b[:n],
 				})
 				require.NoError(t, err)
 
@@ -359,10 +360,11 @@ func TestReadTracks(t *testing.T) {
 				err = mw.Write(&message.MsgAudio{
 					ChunkStreamID:   4,
 					MessageStreamID: 1,
-					Body: append([]byte{
-						flvio.SOUND_AAC<<4 | flvio.SOUND_44Khz<<2 | flvio.SOUND_16BIT<<1 | flvio.SOUND_STEREO,
-						flvio.AAC_SEQHDR,
-					}, enc...),
+					Rate:            flvio.SOUND_44Khz,
+					Depth:           flvio.SOUND_16BIT,
+					Channels:        flvio.SOUND_STEREO,
+					AACType:         flvio.AAC_SEQHDR,
+					Payload:         enc,
 				})
 				require.NoError(t, err)
 
@@ -404,15 +406,16 @@ func TestReadTracks(t *testing.T) {
 				b := make([]byte, 128)
 				var n int
 				codec.ToConfig(b, &n)
-				body := append([]byte{flvio.FRAME_KEY<<4 | flvio.VIDEO_H264, 0, 0, 0, 0}, b[:n]...)
 				err = mw.Write(&message.MsgVideo{
 					ChunkStreamID:   6,
 					MessageStreamID: 1,
-					Body:            body,
+					IsKeyFrame:      true,
+					H264Type:        flvio.AVC_SEQHDR,
+					Payload:         b[:n],
 				})
 				require.NoError(t, err)
 
-			case "no metadata":
+			case "no metadata, video + audio":
 				// C->S H264 decoder config
 				codec := nh264.Codec{
 					SPS: map[int][]byte{
@@ -425,11 +428,12 @@ func TestReadTracks(t *testing.T) {
 				b := make([]byte, 128)
 				var n int
 				codec.ToConfig(b, &n)
-				body := append([]byte{flvio.FRAME_KEY<<4 | flvio.VIDEO_H264, 0, 0, 0, 0}, b[:n]...)
 				err = mw.Write(&message.MsgVideo{
 					ChunkStreamID:   6,
 					MessageStreamID: 1,
-					Body:            body,
+					IsKeyFrame:      true,
+					H264Type:        flvio.AVC_SEQHDR,
+					Payload:         b[:n],
 				})
 				require.NoError(t, err)
 			}
@@ -743,8 +747,10 @@ func TestWriteTracks(t *testing.T) {
 	require.Equal(t, &message.MsgVideo{
 		ChunkStreamID:   6,
 		MessageStreamID: 16777216,
-		Body: []byte{
-			0x17, 0x0, 0x0, 0x0, 0x0, 0x1, 0x64, 0x0,
+		IsKeyFrame:      true,
+		H264Type:        flvio.AVC_SEQHDR,
+		Payload: []byte{
+			0x1, 0x64, 0x0,
 			0xc, 0xff, 0xe1, 0x0, 0x15, 0x67, 0x64, 0x0,
 			0xc, 0xac, 0x3b, 0x50, 0xb0, 0x4b, 0x42, 0x0,
 			0x0, 0x3, 0x0, 0x2, 0x0, 0x0, 0x3, 0x0,
@@ -759,6 +765,10 @@ func TestWriteTracks(t *testing.T) {
 	require.Equal(t, &message.MsgAudio{
 		ChunkStreamID:   4,
 		MessageStreamID: 16777216,
-		Body:            []byte{0xae, 0x0, 0x12, 0x10},
+		Rate:            flvio.SOUND_44Khz,
+		Depth:           flvio.SOUND_16BIT,
+		Channels:        flvio.SOUND_STEREO,
+		AACType:         flvio.AAC_SEQHDR,
+		Payload:         []byte{0x12, 0x10},
 	}, msg)
 }

internal/rtmp/message/msg_audio.go

@@ -1,6 +1,10 @@
 package message
 
 import (
+	"fmt"
+
+	"github.com/notedit/rtmp/format/flv/flvio"
+
 	"github.com/aler9/rtsp-simple-server/internal/rtmp/chunk"
 	"github.com/aler9/rtsp-simple-server/internal/rtmp/rawmessage"
 )
@@ -9,23 +13,49 @@ import (
 type MsgAudio struct {
 	ChunkStreamID   byte
 	MessageStreamID uint32
-	Body            []byte
+	Rate            uint8
+	Depth           uint8
+	Channels        uint8
+	AACType         uint8
+	Payload         []byte
 }
 
 // Unmarshal implements Message.
 func (m *MsgAudio) Unmarshal(raw *rawmessage.Message) error {
 	m.ChunkStreamID = raw.ChunkStreamID
 	m.MessageStreamID = raw.MessageStreamID
-	m.Body = raw.Body
+
+	if len(raw.Body) < 2 {
+		return fmt.Errorf("invalid body size")
+	}
+
+	codec := raw.Body[0] >> 4
+	if codec != flvio.SOUND_AAC {
+		return fmt.Errorf("unsupported audio codec: %d", codec)
+	}
+
+	m.Rate = (raw.Body[0] >> 2) & 0x03
+	m.Depth = (raw.Body[0] >> 1) & 0x01
+	m.Channels = raw.Body[0] & 0x01
+	m.AACType = raw.Body[1]
+	m.Payload = raw.Body[2:]
+
 	return nil
 }
 
 // Marshal implements Message.
 func (m MsgAudio) Marshal() (*rawmessage.Message, error) {
+	body := make([]byte, 2+len(m.Payload))
+
+	body[0] = flvio.SOUND_AAC<<4 | m.Rate<<2 | m.Depth<<1 | m.Channels
+	body[1] = m.AACType
+
+	copy(body[2:], m.Payload)
+
 	return &rawmessage.Message{
 		ChunkStreamID:   m.ChunkStreamID,
 		Type:            chunk.MessageTypeAudio,
 		MessageStreamID: m.MessageStreamID,
-		Body:            m.Body,
+		Body:            body,
 	}, nil
 }

internal/rtmp/message/msg_video.go

@@ -1,6 +1,10 @@
 package message
 
 import (
+	"fmt"
+
+	"github.com/notedit/rtmp/format/flv/flvio"
+
 	"github.com/aler9/rtsp-simple-server/internal/rtmp/chunk"
 	"github.com/aler9/rtsp-simple-server/internal/rtmp/rawmessage"
 )
@@ -9,23 +13,55 @@ import (
 type MsgVideo struct {
 	ChunkStreamID   byte
 	MessageStreamID uint32
-	Body            []byte
+	IsKeyFrame      bool
+	H264Type        uint8
+	PTSDelta        uint32
+	Payload         []byte
 }
 
 // Unmarshal implements Message.
 func (m *MsgVideo) Unmarshal(raw *rawmessage.Message) error {
 	m.ChunkStreamID = raw.ChunkStreamID
 	m.MessageStreamID = raw.MessageStreamID
-	m.Body = raw.Body
+
+	if len(raw.Body) < 5 {
+		return fmt.Errorf("invalid body size")
+	}
+
+	m.IsKeyFrame = (raw.Body[0] >> 4) == flvio.FRAME_KEY
+
+	codec := raw.Body[0] & 0x0F
+	if codec != flvio.VIDEO_H264 {
+		return fmt.Errorf("unsupported video codec: %d", codec)
+	}
+
+	m.H264Type = raw.Body[1]
+	m.PTSDelta = uint32(raw.Body[2])<<16 | uint32(raw.Body[3])<<8 | uint32(raw.Body[4])
+	m.Payload = raw.Body[5:]
+
 	return nil
 }
 
 // Marshal implements Message.
 func (m MsgVideo) Marshal() (*rawmessage.Message, error) {
+	body := make([]byte, 5+len(m.Payload))
+
+	if m.IsKeyFrame {
+		body[0] = flvio.FRAME_KEY << 4
+	} else {
+		body[0] = flvio.FRAME_INTER << 4
+	}
+	body[0] |= flvio.VIDEO_H264
+	body[1] = m.H264Type
+	body[2] = uint8(m.PTSDelta >> 16)
+	body[3] = uint8(m.PTSDelta >> 8)
+	body[4] = uint8(m.PTSDelta)
+	copy(body[5:], m.Payload)
+
 	return &rawmessage.Message{
 		ChunkStreamID:   m.ChunkStreamID,
 		Type:            chunk.MessageTypeVideo,
 		MessageStreamID: m.MessageStreamID,
-		Body:            m.Body,
+		Body:            body,
 	}, nil
 }