RTSP source refactoring

pkg/rtsp/consumer.go

@@ -0,0 +1,102 @@
+package rtsp
+
+import (
+	"encoding/binary"
+	"github.com/AlexxIT/go2rtc/pkg/aac"
+	"github.com/AlexxIT/go2rtc/pkg/h264"
+	"github.com/AlexxIT/go2rtc/pkg/h265"
+	"github.com/AlexxIT/go2rtc/pkg/mjpeg"
+	"github.com/AlexxIT/go2rtc/pkg/streamer"
+	"github.com/pion/rtp"
+)
+
+func (c *Conn) AddTrack(media *streamer.Media, track *streamer.Track) *streamer.Track {
+	switch track.Direction {
+	// send our track to RTSP consumer (ex. FFmpeg)
+	case streamer.DirectionSendonly:
+		i := len(c.tracks)
+		channelID := byte(i << 1)
+
+		codec := track.Codec.Clone()
+		codec.PayloadType = uint8(96 + i)
+
+		if media.MatchAll() {
+			// fill consumer medias list
+			c.Medias = append(c.Medias, &streamer.Media{
+				Kind: media.Kind, Direction: media.Direction,
+				Codecs: []*streamer.Codec{codec},
+			})
+		} else {
+			// find consumer media and replace codec with right one
+			for i, m := range c.Medias {
+				if m == media {
+					media.Codecs = []*streamer.Codec{codec}
+					c.Medias[i] = media
+					break
+				}
+			}
+		}
+
+		track = c.bindTrack(track, channelID, codec.PayloadType)
+		track.Codec = codec
+		c.tracks = append(c.tracks, track)
+
+		return track
+
+	case streamer.DirectionRecvonly:
+		panic("not implemented")
+	}
+
+	panic("wrong direction")
+}
+
+func (c *Conn) bindTrack(
+	track *streamer.Track, channel uint8, payloadType uint8,
+) *streamer.Track {
+	push := func(packet *rtp.Packet) error {
+		if c.state == StateNone {
+			return nil
+		}
+		packet.Header.PayloadType = payloadType
+
+		size := packet.MarshalSize()
+
+		//log.Printf("[RTP] codec: %s, size: %6d, ts: %10d, pt: %2d, ssrc: %d, seq: %d, mark: %v", track.Codec.Name, len(packet.Payload), packet.Timestamp, packet.PayloadType, packet.SSRC, packet.SequenceNumber, packet.Marker)
+
+		data := make([]byte, 4+size)
+		data[0] = '$'
+		data[1] = channel
+		binary.BigEndian.PutUint16(data[2:], uint16(size))
+
+		if _, err := packet.MarshalTo(data[4:]); err != nil {
+			return nil
+		}
+
+		if _, err := c.conn.Write(data); err != nil {
+			return err
+		}
+
+		c.send += size
+
+		return nil
+	}
+
+	if !track.Codec.IsRTP() {
+		switch track.Codec.Name {
+		case streamer.CodecH264:
+			wrapper := h264.RTPPay(1500)
+			push = wrapper(push)
+		case streamer.CodecH265:
+			wrapper := h265.RTPPay(1500)
+			push = wrapper(push)
+		case streamer.CodecAAC:
+			wrapper := aac.RTPPay(1500)
+			push = wrapper(push)
+		case streamer.CodecJPEG:
+			wrapper := mjpeg.RTPPay()
+			push = wrapper(push)
+		}
+	}
+
+	return track.Bind(push)
+}