diff --git a/track.go b/track.go index 627acf1a..a73ce772 100644 --- a/track.go +++ b/track.go @@ -26,7 +26,7 @@ type Track struct { Media *psdp.MediaDescription } -// NewTrackH264 initializes an H264 track. +// NewTrackH264 initializes an H264 track from a SPS and PPS. func NewTrackH264(payloadType uint8, sps []byte, pps []byte) (*Track, error) { spropParameterSets := base64.StdEncoding.EncodeToString(sps) + "," + base64.StdEncoding.EncodeToString(pps) @@ -57,7 +57,71 @@ func NewTrackH264(payloadType uint8, sps []byte, pps []byte) (*Track, error) { }, nil } -// NewTrackAAC initializes an AAC track. +// IsH264 checks whether the track is a H264 track. +func (t *Track) IsH264() bool { + v, ok := t.Media.Attribute("rtpmap") + if !ok { + return false + } + + vals := strings.Split(v, " ") + if len(vals) != 2 { + return false + } + + return vals[1] == "H264/90000" +} + +// ExtractDataH264 extracts the SPS and PPS from an H264 track. +func (t *Track) ExtractDataH264() ([]byte, []byte, error) { + v, ok := t.Media.Attribute("fmtp") + if !ok { + return nil, nil, fmt.Errorf("unable to find fmtp") + } + + tmp := strings.SplitN(v, " ", 2) + if len(tmp) != 2 { + return nil, nil, fmt.Errorf("unable to parse fmtp (%v)", v) + } + + var sps []byte + var pps []byte + + for _, kv := range strings.Split(tmp[1], ";") { + kv = strings.Trim(kv, " ") + + tmp := strings.SplitN(kv, "=", 2) + if len(tmp) != 2 { + return nil, nil, fmt.Errorf("unable to parse fmtp (%v)", v) + } + + if tmp[0] == "sprop-parameter-sets" { + tmp := strings.SplitN(tmp[1], ",", 2) + if len(tmp) != 2 { + return nil, nil, fmt.Errorf("unable to parse sprop-parameter-sets (%v)", v) + } + + var err error + sps, err = base64.StdEncoding.DecodeString(tmp[0]) + if err != nil { + return nil, nil, fmt.Errorf("unable to parse sprop-parameter-sets (%v)", v) + } + + pps, err = base64.StdEncoding.DecodeString(tmp[1]) + if err != nil { + return nil, nil, fmt.Errorf("unable to parse sprop-parameter-sets (%v)", v) + } + } + } + + if sps == nil || pps == nil { + return nil, nil, fmt.Errorf("unable to find SPS or PPS (%v)", v) + } + + return sps, pps, nil +} + +// NewTrackAAC initializes an AAC track from a configuration. func NewTrackAAC(payloadType uint8, config []byte) (*Track, error) { codec, err := aac.FromMPEG4AudioConfigBytes(config) if err != nil { @@ -109,6 +173,60 @@ func NewTrackAAC(payloadType uint8, config []byte) (*Track, error) { }, nil } +// IsAAC checks whether the track is a AAC track. +func (t *Track) IsAAC() bool { + v, ok := t.Media.Attribute("rtpmap") + if !ok { + return false + } + + vals := strings.Split(v, " ") + if len(vals) != 2 { + return false + } + + return vals[1] == "MPEG4-GENERIC/48000/2" +} + +// ExtractDataAAC extracts the config from an AAC track. +func (t *Track) ExtractDataAAC() ([]byte, error) { + v, ok := t.Media.Attribute("fmtp") + if !ok { + return nil, fmt.Errorf("unable to find fmtp") + } + + tmp := strings.SplitN(v, " ", 2) + if len(tmp) != 2 { + return nil, fmt.Errorf("unable to parse fmtp (%v)", v) + } + + var config []byte + + for _, kv := range strings.Split(tmp[1], ";") { + kv = strings.Trim(kv, " ") + + tmp := strings.SplitN(kv, "=", 2) + if len(tmp) != 2 { + return nil, fmt.Errorf("unable to parse fmtp (%v)", v) + } + + if tmp[0] == "config" { + var err error + config, err = hex.DecodeString(tmp[1]) + if err != nil { + return nil, fmt.Errorf("unable to parse config (%v)", v) + } + break + } + } + + if config == nil { + return nil, fmt.Errorf("unable to find config (%v)", v) + } + + return config, nil +} + // ClockRate returns the clock rate of the track. func (t *Track) ClockRate() (int, error) { if len(t.Media.MediaName.Formats) != 1 { diff --git a/track_test.go b/track_test.go index c8927b37..1be9bc9c 100644 --- a/track_test.go +++ b/track_test.go @@ -3,6 +3,7 @@ package gortsplib import ( "testing" + psdp "github.com/pion/sdp/v3" "github.com/stretchr/testify/require" ) @@ -82,3 +83,74 @@ func TestTrackClockRate(t *testing.T) { }) } } + +var testH264SPS = []byte("\x67\x64\x00\x0c\xac\x3b\x50\xb0\x4b\x42\x00\x00\x03\x00\x02\x00\x00\x03\x00\x3d\x08") + +var testH264PPS = []byte("\x68\xee\x3c\x80") + +var testH264Track = &Track{ + Media: &psdp.MediaDescription{ + MediaName: psdp.MediaName{ + Media: "video", + Protos: []string{"RTP", "AVP"}, + Formats: []string{"96"}, + }, + Attributes: []psdp.Attribute{ + { + Key: "rtpmap", + Value: "96 H264/90000", + }, + { + Key: "fmtp", + Value: "96 packetization-mode=1; sprop-parameter-sets=Z2QADKw7ULBLQgAAAwACAAADAD0I,aO48gA==; profile-level-id=64000C", + }, + }, + }, +} + +func TestTrackH264New(t *testing.T) { + tr, err := NewTrackH264(96, testH264SPS, testH264PPS) + require.NoError(t, err) + require.Equal(t, testH264Track, tr) +} + +func TestTrackH264Extract(t *testing.T) { + sps, pps, err := testH264Track.ExtractDataH264() + require.NoError(t, err) + require.Equal(t, testH264SPS, sps) + require.Equal(t, testH264PPS, pps) +} + +var testAACConfig = []byte{17, 144} + +var testAACTrack = &Track{ + Media: &psdp.MediaDescription{ + MediaName: psdp.MediaName{ + Media: "audio", + Protos: []string{"RTP", "AVP"}, + Formats: []string{"96"}, + }, + Attributes: []psdp.Attribute{ + { + Key: "rtpmap", + Value: "96 MPEG4-GENERIC/48000/2", + }, + { + Key: "fmtp", + Value: "96 profile-level-id=1; mode=AAC-hbr; sizelength=13; indexlength=3; indexdeltalength=3; config=1190", + }, + }, + }, +} + +func TestTrackAACNew(t *testing.T) { + tr, err := NewTrackAAC(96, testAACConfig) + require.NoError(t, err) + require.Equal(t, testAACTrack, tr) +} + +func TestTrackAACExtract(t *testing.T) { + config, err := testAACTrack.ExtractDataAAC() + require.NoError(t, err) + require.Equal(t, testAACConfig, config) +}