RTMP client: move writing of track info into WriteMetadata()

internal/clientrtmp/client.go

@@ -88,12 +88,8 @@ type Client struct {
 	pathMan             PathMan
 	parent              Parent
 
-	// read mode only
-	h264Decoder *rtph264.Decoder
-	videoTrack  *gortsplib.Track
-	aacDecoder  *rtpaac.Decoder
-	audioTrack  *gortsplib.Track
-	ringBuffer  *ringbuffer.RingBuffer
+	// read mode
+	ringBuffer *ringbuffer.RingBuffer
 
 	// in
 	terminate chan struct{}
@@ -225,10 +221,9 @@ func (c *Client) runRead() {
 	}
 
 	var videoTrack *gortsplib.Track
-	var h264SPS []byte
-	var h264PPS []byte
+	var h264Decoder *rtph264.Decoder
 	var audioTrack *gortsplib.Track
-	var aacConfig []byte
+	var aacDecoder *rtpaac.Decoder
 
 	err = func() error {
 		for i, t := range tracks {
@@ -236,25 +231,18 @@ func (c *Client) runRead() {
 				if videoTrack != nil {
 					return fmt.Errorf("can't read track %d with RTMP: too many tracks", i+1)
 				}
-				videoTrack = t
 
-				var err error
-				h264SPS, h264PPS, err = t.ExtractDataH264()
-				if err != nil {
-					return err
-				}
+				videoTrack = t
+				h264Decoder = rtph264.NewDecoder()
 
 			} else if t.IsAAC() {
 				if audioTrack != nil {
 					return fmt.Errorf("can't read track %d with RTMP: too many tracks", i+1)
 				}
-				audioTrack = t
 
-				var err error
-				aacConfig, err = t.ExtractDataAAC()
-				if err != nil {
-					return err
-				}
+				audioTrack = t
+				clockRate, _ := audioTrack.ClockRate()
+				aacDecoder = rtpaac.NewDecoder(clockRate)
 			}
 		}
 
@@ -265,21 +253,6 @@ func (c *Client) runRead() {
 		c.conn.NetConn().SetWriteDeadline(time.Now().Add(c.writeTimeout))
 		c.conn.WriteMetadata(videoTrack, audioTrack)
 
-		if videoTrack != nil {
-			c.conn.NetConn().SetWriteDeadline(time.Now().Add(c.writeTimeout))
-			c.conn.WriteH264Config(h264SPS, h264PPS)
-			c.h264Decoder = rtph264.NewDecoder()
-			c.videoTrack = videoTrack
-		}
-
-		if audioTrack != nil {
-			c.conn.NetConn().SetWriteDeadline(time.Now().Add(c.writeTimeout))
-			c.conn.WriteAACConfig(aacConfig)
-			clockRate, _ := audioTrack.ClockRate()
-			c.aacDecoder = rtpaac.NewDecoder(clockRate)
-			c.audioTrack = audioTrack
-		}
-
 		c.ringBuffer = ringbuffer.New(uint64(c.readBufferCount))
 
 		resc := make(chan client.PlayRes)
@@ -325,8 +298,8 @@ func (c *Client) runRead() {
 
 				now := time.Now()
 
-				if c.videoTrack != nil && pair.trackID == c.videoTrack.ID {
-					nts, err := c.h264Decoder.Decode(pair.buf)
+				if videoTrack != nil && pair.trackID == videoTrack.ID {
+					nts, err := h264Decoder.Decode(pair.buf)
 					if err != nil {
 						if err != rtph264.ErrMorePacketsNeeded {
 							c.log(logger.Warn, "unable to decode video track: %v", err)
@@ -356,8 +329,8 @@ func (c *Client) runRead() {
 						videoBuf = append(videoBuf, nt.NALU)
 					}
 
-				} else if c.audioTrack != nil && pair.trackID == c.audioTrack.ID {
-					ats, err := c.aacDecoder.Decode(pair.buf)
+				} else if audioTrack != nil && pair.trackID == audioTrack.ID {
+					ats, err := aacDecoder.Decode(pair.buf)
 					if err != nil {
 						c.log(logger.Warn, "unable to decode audio track: %v", err)
 						continue

internal/rtmp/aac.go

@@ -6,14 +6,6 @@ import (
 	"github.com/notedit/rtmp/av"
 )
 
-// WriteAACConfig writes an AAC config.
-func (c *Conn) WriteAACConfig(config []byte) error {
-	return c.WritePacket(av.Packet{
-		Type: av.AACDecoderConfig,
-		Data: config,
-	})
-}
-
 // WriteAAC writes an AAC AU.
 func (c *Conn) WriteAAC(au []byte, dts time.Duration) error {
 	return c.WritePacket(av.Packet{

internal/rtmp/h264.go

@@ -4,32 +4,10 @@ import (
 	"time"
 
 	"github.com/notedit/rtmp/av"
-	nh264 "github.com/notedit/rtmp/codec/h264"
 
 	"github.com/aler9/rtsp-simple-server/internal/h264"
 )
 
-// WriteH264Config writes a H264 config.
-func (c *Conn) WriteH264Config(sps []byte, pps []byte) error {
-	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]
-
-	return c.WritePacket(av.Packet{
-		Type: av.H264DecoderConfig,
-		Data: b,
-	})
-}
-
 // WriteH264 writes H264 NALUs.
 func (c *Conn) WriteH264(nalus [][]byte, dts time.Duration) error {
 	data, err := h264.EncodeAVCC(nalus)

internal/rtmp/metadata.go

@@ -15,12 +15,12 @@ const (
 )
 
 // ReadMetadata extracts track informations from a connection that is publishing.
-func (conn *Conn) ReadMetadata() (*gortsplib.Track, *gortsplib.Track, error) {
+func (c *Conn) ReadMetadata() (*gortsplib.Track, *gortsplib.Track, error) {
 	var videoTrack *gortsplib.Track
 	var audioTrack *gortsplib.Track
 
 	md, err := func() (flvio.AMFMap, error) {
-		pkt, err := conn.ReadPacket()
+		pkt, err := c.ReadPacket()
 		if err != nil {
 			return nil, err
 		}
@@ -111,7 +111,7 @@ func (conn *Conn) ReadMetadata() (*gortsplib.Track, *gortsplib.Track, error) {
 
 	for {
 		var pkt av.Packet
-		pkt, err = conn.ReadPacket()
+		pkt, err = c.ReadPacket()
 		if err != nil {
 			return nil, nil, err
 		}
@@ -157,8 +157,8 @@ func (conn *Conn) ReadMetadata() (*gortsplib.Track, *gortsplib.Track, error) {
 }
 
 // WriteMetadata writes track informations to a connection that is reading.
-func (conn *Conn) WriteMetadata(videoTrack *gortsplib.Track, audioTrack *gortsplib.Track) error {
-	return conn.WritePacket(av.Packet{
+func (c *Conn) WriteMetadata(videoTrack *gortsplib.Track, audioTrack *gortsplib.Track) error {
+	err := c.WritePacket(av.Packet{
 		Type: av.Metadata,
 		Data: flvio.FillAMF0ValMalloc(flvio.AMFMap{
 			{
@@ -189,4 +189,52 @@ func (conn *Conn) WriteMetadata(videoTrack *gortsplib.Track, audioTrack *gortspl
 			},
 		}),
 	})
+	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
 }