From d3797d3139f166226c95eb73fc722b7b09f8d2ca Mon Sep 17 00:00:00 2001 From: aler9 <46489434+aler9@users.noreply.github.com> Date: Tue, 7 Jun 2022 22:48:10 +0200 Subject: [PATCH] rtmp: improve video / audio messages --- internal/rtmp/conn.go | 6 +++++ internal/rtmp/conn_test.go | 42 ++++++++++++++++++------------ internal/rtmp/message/msg_audio.go | 36 ++++++++++++++++++++++--- internal/rtmp/message/msg_video.go | 42 +++++++++++++++++++++++++++--- 4 files changed, 104 insertions(+), 22 deletions(-) diff --git a/internal/rtmp/conn.go b/internal/rtmp/conn.go index 681b175d..fa1fd903 100644 --- a/internal/rtmp/conn.go +++ b/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 { diff --git a/internal/rtmp/conn_test.go b/internal/rtmp/conn_test.go index dbaae924..97864750 100644 --- a/internal/rtmp/conn_test.go +++ b/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) } diff --git a/internal/rtmp/message/msg_audio.go b/internal/rtmp/message/msg_audio.go index 10a74e24..82c98350 100644 --- a/internal/rtmp/message/msg_audio.go +++ b/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 } diff --git a/internal/rtmp/message/msg_video.go b/internal/rtmp/message/msg_video.go index 81f2e579..975ab5d7 100644 --- a/internal/rtmp/message/msg_video.go +++ b/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 }