add Track.IsH264, Track.ExtractDataH264, Track.IsAAC, Track.ExtractDataAAC

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 {

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)
+}