mediamtx/internal/rtmp/metadata.go

241 lines
4.2 KiB
Go
Raw Normal View History

2021-04-03 09:39:19 +00:00
package rtmp
2021-01-31 15:44:32 +00:00
import (
"fmt"
"github.com/aler9/gortsplib"
"github.com/notedit/rtmp/av"
2021-04-03 09:39:19 +00:00
nh264 "github.com/notedit/rtmp/codec/h264"
2021-01-31 15:44:32 +00:00
"github.com/notedit/rtmp/format/flv/flvio"
)
const (
codecH264 = 7
codecAAC = 10
)
2021-04-03 08:52:59 +00:00
// ReadMetadata extracts track informations from a connection that is publishing.
func (c *Conn) ReadMetadata() (*gortsplib.Track, *gortsplib.Track, error) {
2021-04-03 08:52:59 +00:00
var videoTrack *gortsplib.Track
var audioTrack *gortsplib.Track
2021-01-31 15:44:32 +00:00
2021-04-03 08:52:59 +00:00
md, err := func() (flvio.AMFMap, error) {
pkt, err := c.ReadPacket()
2021-04-03 08:52:59 +00:00
if err != nil {
return nil, err
}
2021-01-31 15:44:32 +00:00
2021-04-03 08:52:59 +00:00
if pkt.Type != av.Metadata {
return nil, fmt.Errorf("first packet must be metadata")
}
2021-01-31 15:44:32 +00:00
2021-04-03 08:52:59 +00:00
arr, err := flvio.ParseAMFVals(pkt.Data, false)
if err != nil {
return nil, err
}
2021-01-31 15:44:32 +00:00
2021-04-03 08:52:59 +00:00
if len(arr) != 1 {
return nil, fmt.Errorf("invalid metadata")
}
2021-01-31 15:44:32 +00:00
2021-04-03 08:52:59 +00:00
ma, ok := arr[0].(flvio.AMFMap)
if !ok {
return nil, fmt.Errorf("invalid metadata")
}
2021-01-31 15:44:32 +00:00
2021-04-03 08:52:59 +00:00
return ma, nil
}()
2021-01-31 15:44:32 +00:00
if err != nil {
return nil, nil, err
}
2021-03-05 10:15:27 +00:00
hasVideo, err := func() (bool, error) {
v, ok := md.GetV("videocodecid")
if !ok {
return false, nil
2021-01-31 15:44:32 +00:00
}
2021-03-05 10:15:27 +00:00
switch vt := v.(type) {
case float64:
switch vt {
case 0:
return false, nil
case codecH264:
return true, nil
}
case string:
2021-03-20 13:14:41 +00:00
if vt == "avc1" {
2021-03-05 10:15:27 +00:00
return true, nil
}
}
return false, fmt.Errorf("unsupported video codec %v", v)
}()
if err != nil {
return nil, nil, err
2021-01-31 15:44:32 +00:00
}
2021-03-05 10:15:27 +00:00
hasAudio, err := func() (bool, error) {
v, ok := md.GetV("audiocodecid")
if !ok {
return false, nil
2021-01-31 15:44:32 +00:00
}
2021-03-05 10:15:27 +00:00
switch vt := v.(type) {
case float64:
switch vt {
case 0:
return false, nil
case codecAAC:
return true, nil
}
case string:
2021-03-20 13:14:41 +00:00
if vt == "mp4a" {
2021-03-05 10:15:27 +00:00
return true, nil
}
}
return false, fmt.Errorf("unsupported audio codec %v", v)
}()
if err != nil {
return nil, nil, err
2021-01-31 15:44:32 +00:00
}
if !hasVideo && !hasAudio {
return nil, nil, fmt.Errorf("stream has no tracks")
}
for {
var pkt av.Packet
pkt, err = c.ReadPacket()
2021-01-31 15:44:32 +00:00
if err != nil {
return nil, nil, err
}
switch pkt.Type {
case av.H264DecoderConfig:
if !hasVideo {
return nil, nil, fmt.Errorf("unexpected video packet")
}
if videoTrack != nil {
return nil, nil, fmt.Errorf("video track setupped twice")
}
2021-04-03 09:39:19 +00:00
codec, err := nh264.FromDecoderConfig(pkt.Data)
2021-01-31 15:44:32 +00:00
if err != nil {
return nil, nil, err
}
videoTrack, err = gortsplib.NewTrackH264(96, codec.SPS[0], codec.PPS[0])
if err != nil {
return nil, nil, err
}
case av.AACDecoderConfig:
if !hasAudio {
return nil, nil, fmt.Errorf("unexpected audio packet")
}
if audioTrack != nil {
return nil, nil, fmt.Errorf("audio track setupped twice")
}
audioTrack, err = gortsplib.NewTrackAAC(96, pkt.Data)
if err != nil {
return nil, nil, err
}
}
if (!hasVideo || videoTrack != nil) &&
(!hasAudio || audioTrack != nil) {
return videoTrack, audioTrack, nil
}
}
}
2021-03-10 14:06:45 +00:00
2021-04-03 08:52:59 +00:00
// WriteMetadata writes track informations to a connection that is reading.
func (c *Conn) WriteMetadata(videoTrack *gortsplib.Track, audioTrack *gortsplib.Track) error {
err := c.WritePacket(av.Packet{
2021-03-10 14:06:45 +00:00
Type: av.Metadata,
Data: flvio.FillAMF0ValMalloc(flvio.AMFMap{
{
K: "videodatarate",
V: float64(0),
},
{
K: "videocodecid",
V: func() float64 {
if videoTrack != nil {
return codecH264
}
return 0
}(),
},
{
K: "audiodatarate",
V: float64(0),
},
{
K: "audiocodecid",
V: func() float64 {
if audioTrack != nil {
return codecAAC
}
return 0
}(),
},
}),
})
if err != nil {
return err
}
if videoTrack != nil {
sps, pps, err := videoTrack.ExtractDataH264()
if err != nil {
return err
}
codec := nh264.Codec{
SPS: map[int][]byte{
0: sps,
},
PPS: map[int][]byte{
0: pps,
},
}
b := make([]byte, 128)
var n int
codec.ToConfig(b, &n)
b = b[:n]
err = c.WritePacket(av.Packet{
Type: av.H264DecoderConfig,
Data: b,
})
if err != nil {
return err
}
}
if audioTrack != nil {
config, err := audioTrack.ExtractDataAAC()
if err != nil {
return err
}
err = c.WritePacket(av.Packet{
Type: av.AACDecoderConfig,
Data: config,
})
if err != nil {
return err
}
}
return nil
2021-03-10 14:06:45 +00:00
}