support publishing, reading, proxying with SRT (#2068)

internal/core/srt_source.go

@@ -0,0 +1,221 @@
+package core
+
+import (
+	"context"
+	"time"
+
+	"github.com/bluenviron/gortsplib/v3/pkg/formats"
+	"github.com/bluenviron/gortsplib/v3/pkg/media"
+	"github.com/bluenviron/mediacommon/pkg/formats/mpegts"
+	"github.com/datarhei/gosrt"
+
+	"github.com/bluenviron/mediamtx/internal/conf"
+	"github.com/bluenviron/mediamtx/internal/formatprocessor"
+	"github.com/bluenviron/mediamtx/internal/logger"
+	"github.com/bluenviron/mediamtx/internal/stream"
+)
+
+type srtSourceParent interface {
+	logger.Writer
+	sourceStaticImplSetReady(req pathSourceStaticSetReadyReq) pathSourceStaticSetReadyRes
+	sourceStaticImplSetNotReady(req pathSourceStaticSetNotReadyReq)
+}
+
+type srtSource struct {
+	readTimeout conf.StringDuration
+	parent      srtSourceParent
+}
+
+func newSRTSource(
+	readTimeout conf.StringDuration,
+	parent srtSourceParent,
+) *srtSource {
+	s := &srtSource{
+		readTimeout: readTimeout,
+		parent:      parent,
+	}
+
+	return s
+}
+
+func (s *srtSource) Log(level logger.Level, format string, args ...interface{}) {
+	s.parent.Log(level, "[srt source] "+format, args...)
+}
+
+// run implements sourceStaticImpl.
+func (s *srtSource) run(ctx context.Context, cnf *conf.PathConf, reloadConf chan *conf.PathConf) error {
+	s.Log(logger.Debug, "connecting")
+
+	conf := srt.DefaultConfig()
+	address, err := conf.UnmarshalURL(cnf.Source)
+	if err != nil {
+		return err
+	}
+
+	err = conf.Validate()
+	if err != nil {
+		return err
+	}
+
+	sconn, err := srt.Dial("srt", address, conf)
+	if err != nil {
+		return err
+	}
+
+	readDone := make(chan error)
+	go func() {
+		readDone <- s.runReader(sconn)
+	}()
+
+	for {
+		select {
+		case err := <-readDone:
+			sconn.Close()
+			return err
+
+		case <-reloadConf:
+
+		case <-ctx.Done():
+			sconn.Close()
+			<-readDone
+			return nil
+		}
+	}
+}
+
+func (s *srtSource) runReader(sconn srt.Conn) error {
+	sconn.SetReadDeadline(time.Now().Add(time.Duration(s.readTimeout)))
+	r, err := mpegts.NewReader(mpegts.NewBufferedReader(sconn))
+	if err != nil {
+		return err
+	}
+
+	var medias media.Medias
+	var stream *stream.Stream
+
+	var td *mpegts.TimeDecoder
+	decodeTime := func(t int64) time.Duration {
+		if td == nil {
+			td = mpegts.NewTimeDecoder(t)
+		}
+		return td.Decode(t)
+	}
+
+	for _, track := range r.Tracks() { //nolint:dupl
+		var medi *media.Media
+
+		switch tcodec := track.Codec.(type) {
+		case *mpegts.CodecH264:
+			medi = &media.Media{
+				Type: media.TypeVideo,
+				Formats: []formats.Format{&formats.H264{
+					PayloadTyp:        96,
+					PacketizationMode: 1,
+				}},
+			}
+
+			r.OnDataH26x(track, func(pts int64, _ int64, au [][]byte) error {
+				stream.WriteUnit(medi, medi.Formats[0], &formatprocessor.UnitH264{
+					BaseUnit: formatprocessor.BaseUnit{
+						NTP: time.Now(),
+					},
+					PTS: decodeTime(pts),
+					AU:  au,
+				})
+				return nil
+			})
+
+		case *mpegts.CodecH265:
+			medi = &media.Media{
+				Type: media.TypeVideo,
+				Formats: []formats.Format{&formats.H265{
+					PayloadTyp: 96,
+				}},
+			}
+
+			r.OnDataH26x(track, func(pts int64, _ int64, au [][]byte) error {
+				stream.WriteUnit(medi, medi.Formats[0], &formatprocessor.UnitH265{
+					BaseUnit: formatprocessor.BaseUnit{
+						NTP: time.Now(),
+					},
+					PTS: decodeTime(pts),
+					AU:  au,
+				})
+				return nil
+			})
+
+		case *mpegts.CodecMPEG4Audio:
+			medi = &media.Media{
+				Type: media.TypeAudio,
+				Formats: []formats.Format{&formats.MPEG4Audio{
+					PayloadTyp:       96,
+					SizeLength:       13,
+					IndexLength:      3,
+					IndexDeltaLength: 3,
+					Config:           &tcodec.Config,
+				}},
+			}
+
+			r.OnDataMPEG4Audio(track, func(pts int64, _ int64, aus [][]byte) error {
+				stream.WriteUnit(medi, medi.Formats[0], &formatprocessor.UnitMPEG4AudioGeneric{
+					BaseUnit: formatprocessor.BaseUnit{
+						NTP: time.Now(),
+					},
+					PTS: decodeTime(pts),
+					AUs: aus,
+				})
+				return nil
+			})
+
+		case *mpegts.CodecOpus:
+			medi = &media.Media{
+				Type: media.TypeAudio,
+				Formats: []formats.Format{&formats.Opus{
+					PayloadTyp: 96,
+					IsStereo:   (tcodec.ChannelCount == 2),
+				}},
+			}
+
+			r.OnDataOpus(track, func(pts int64, _ int64, packets [][]byte) error {
+				stream.WriteUnit(medi, medi.Formats[0], &formatprocessor.UnitOpus{
+					BaseUnit: formatprocessor.BaseUnit{
+						NTP: time.Now(),
+					},
+					PTS:     decodeTime(pts),
+					Packets: packets,
+				})
+				return nil
+			})
+		}
+
+		medias = append(medias, medi)
+	}
+
+	res := s.parent.sourceStaticImplSetReady(pathSourceStaticSetReadyReq{
+		medias:             medias,
+		generateRTPPackets: true,
+	})
+	if res.err != nil {
+		return res.err
+	}
+
+	s.Log(logger.Info, "ready: %s", sourceMediaInfo(medias))
+
+	stream = res.stream
+
+	for {
+		sconn.SetReadDeadline(time.Now().Add(time.Duration(s.readTimeout)))
+		err := r.Read()
+		if err != nil {
+			return err
+		}
+	}
+}
+
+// apiSourceDescribe implements sourceStaticImpl.
+func (*srtSource) apiSourceDescribe() pathAPISourceOrReader {
+	return pathAPISourceOrReader{
+		Type: "srtSource",
+		ID:   "",
+	}
+}