Move cmd module to internal

2023-05-01 12:55:14 +03:00
parent bc770f1a85
commit 75f61b38ac
58 changed files with 109 additions and 104 deletions
@@ -0,0 +1,41 @@
+package streams
+
+import (
+	"fmt"
+	"github.com/AlexxIT/go2rtc/pkg/core"
+	"strings"
+	"sync"
+)
+
+type Handler func(url string) (core.Producer, error)
+
+var handlers = map[string]Handler{}
+var handlersMu sync.Mutex
+
+func HandleFunc(scheme string, handler Handler) {
+	handlersMu.Lock()
+	handlers[scheme] = handler
+	handlersMu.Unlock()
+}
+
+func getHandler(url string) Handler {
+	i := strings.IndexByte(url, ':')
+	if i <= 0 { // TODO: i < 4 ?
+		return nil
+	}
+	handlersMu.Lock()
+	defer handlersMu.Unlock()
+	return handlers[url[:i]]
+}
+
+func HasProducer(url string) bool {
+	return getHandler(url) != nil
+}
+
+func GetProducer(url string) (core.Producer, error) {
+	handler := getHandler(url)
+	if handler == nil {
+		return nil, fmt.Errorf("unsupported scheme: %s", url)
+	}
+	return handler(url)
+}
@@ -0,0 +1,139 @@
+package streams
+
+import (
+	"encoding/json"
+	"github.com/AlexxIT/go2rtc/internal/api"
+	"github.com/AlexxIT/go2rtc/internal/app"
+	"github.com/AlexxIT/go2rtc/internal/app/store"
+	"github.com/rs/zerolog"
+	"net/http"
+	"net/url"
+)
+
+func Init() {
+	var cfg struct {
+		Mod map[string]any `yaml:"streams"`
+	}
+
+	app.LoadConfig(&cfg)
+
+	log = app.GetLogger("streams")
+
+	for name, item := range cfg.Mod {
+		streams[name] = NewStream(item)
+	}
+
+	for name, item := range store.GetDict("streams") {
+		streams[name] = NewStream(item)
+	}
+
+	api.HandleFunc("api/streams", streamsHandler)
+}
+
+func Get(name string) *Stream {
+	return streams[name]
+}
+
+func New(name string, source any) *Stream {
+	stream := NewStream(source)
+	streams[name] = stream
+	return stream
+}
+
+func NewTemplate(name string, source any) *Stream {
+	// check if source links to some stream name from go2rtc
+	if rawURL, ok := source.(string); ok {
+		if u, err := url.Parse(rawURL); err == nil && u.Scheme == "rtsp" {
+			if stream, ok := streams[u.Path[1:]]; ok {
+				streams[name] = stream
+				return stream
+			}
+		}
+	}
+
+	return New(name, "{input}")
+}
+
+func GetOrNew(src string) *Stream {
+	if stream, ok := streams[src]; ok {
+		return stream
+	}
+
+	if !HasProducer(src) {
+		return nil
+	}
+
+	log.Info().Str("url", src).Msg("[streams] create new stream")
+
+	return New(src, src)
+}
+
+func GetAll() (names []string) {
+	for name := range streams {
+		names = append(names, name)
+	}
+	return
+}
+
+func streamsHandler(w http.ResponseWriter, r *http.Request) {
+	query := r.URL.Query()
+	src := query.Get("src")
+
+	// without source - return all streams list
+	if src == "" && r.Method != "POST" {
+		_ = json.NewEncoder(w).Encode(streams)
+		return
+	}
+
+	// Not sure about all this API. Should be rewrited...
+	switch r.Method {
+	case "GET":
+		e := json.NewEncoder(w)
+		e.SetIndent("", "  ")
+		_ = e.Encode(streams[src])
+
+	case "PUT":
+		name := query.Get("name")
+		if name == "" {
+			name = src
+		}
+
+		New(name, src)
+
+	case "PATCH":
+		name := query.Get("name")
+		if name == "" {
+			http.Error(w, "", http.StatusBadRequest)
+			return
+		}
+
+		// support {input} templates: https://github.com/AlexxIT/go2rtc#module-hass
+		stream := Get(name)
+		if stream == nil {
+			stream = NewTemplate(name, src)
+		}
+		stream.SetSource(src)
+
+	case "POST":
+		// with dst - redirect source to dst
+		if dst := query.Get("dst"); dst != "" {
+			if stream := Get(dst); stream != nil {
+				if err := stream.Play(src); err != nil {
+					http.Error(w, err.Error(), http.StatusInternalServerError)
+				} else {
+					_ = json.NewEncoder(w).Encode(stream)
+				}
+			} else {
+				http.Error(w, "", http.StatusNotFound)
+			}
+		} else {
+			http.Error(w, "", http.StatusBadRequest)
+		}
+
+	case "DELETE":
+		delete(streams, src)
+	}
+}
+
+var log zerolog.Logger
+var streams = map[string]*Stream{}
@@ -0,0 +1,151 @@
+package streams
+
+import (
+	"errors"
+	"github.com/AlexxIT/go2rtc/pkg/core"
+)
+
+func (s *Stream) Play(source string) error {
+	s.mu.Lock()
+	for _, producer := range s.producers {
+		if producer.state == stateInternal && producer.conn != nil {
+			_ = producer.conn.Stop()
+		}
+	}
+	s.mu.Unlock()
+
+	if source == "" {
+		return nil
+	}
+
+	var src core.Producer
+
+	for _, producer := range s.producers {
+		if producer.conn == nil {
+			continue
+		}
+
+		cons, ok := producer.conn.(core.Consumer)
+		if !ok {
+			continue
+		}
+
+		if src == nil {
+			var err error
+			if src, err = GetProducer(source); err != nil {
+				return err
+			}
+		}
+
+		if !matchMedia(src, cons) {
+			continue
+		}
+
+		s.AddInternalProducer(src)
+
+		go func() {
+			_ = src.Start()
+			s.RemoveProducer(src)
+		}()
+
+		return nil
+	}
+
+	for _, producer := range s.producers {
+		// start new client
+		dst, err := GetProducer(producer.url)
+		if err != nil {
+			continue
+		}
+
+		// check if client support consumer interface
+		cons, ok := dst.(core.Consumer)
+		if !ok {
+			_ = dst.Stop()
+			continue
+		}
+
+		// start new producer
+		if src == nil {
+			if src, err = GetProducer(source); err != nil {
+				return err
+			}
+		}
+
+		if !matchMedia(src, cons) {
+			_ = dst.Stop()
+			continue
+		}
+
+		s.AddInternalProducer(src)
+		s.AddInternalConsumer(cons)
+
+		go func() {
+			_ = src.Start()
+			_ = dst.Stop()
+			s.RemoveProducer(src)
+		}()
+
+		go func() {
+			_ = dst.Start()
+			_ = src.Stop()
+			s.RemoveInternalConsumer(cons)
+		}()
+
+		return nil
+	}
+
+	return errors.New("can't find consumer")
+}
+
+func (s *Stream) AddInternalProducer(conn core.Producer) {
+	producer := &Producer{conn: conn, state: stateInternal}
+	s.mu.Lock()
+	s.producers = append(s.producers, producer)
+	s.mu.Unlock()
+}
+
+func (s *Stream) AddInternalConsumer(conn core.Consumer) {
+	s.mu.Lock()
+	s.consumers = append(s.consumers, conn)
+	s.mu.Unlock()
+}
+
+func (s *Stream) RemoveInternalConsumer(conn core.Consumer) {
+	s.mu.Lock()
+	for i, consumer := range s.consumers {
+		if consumer == conn {
+			s.consumers = append(s.consumers[:i], s.consumers[i+1:]...)
+			break
+		}
+	}
+	s.mu.Unlock()
+}
+
+func matchMedia(prod core.Producer, cons core.Consumer) bool {
+	for _, consMedia := range cons.GetMedias() {
+		for _, prodMedia := range prod.GetMedias() {
+			if prodMedia.Direction != core.DirectionRecvonly {
+				continue
+			}
+
+			prodCodec, consCodec := prodMedia.MatchMedia(consMedia)
+			if prodCodec == nil {
+				continue
+			}
+
+			track, err := prod.GetTrack(prodMedia, prodCodec)
+			if err != nil {
+				continue
+			}
+
+			if err = cons.AddTrack(consMedia, consCodec, track); err != nil {
+				continue
+			}
+
+			return true
+		}
+	}
+
+	return false
+}
@@ -0,0 +1,251 @@
+package streams
+
+import (
+	"encoding/json"
+	"errors"
+	"github.com/AlexxIT/go2rtc/pkg/core"
+	"strings"
+	"sync"
+	"time"
+)
+
+type state byte
+
+const (
+	stateNone state = iota
+	stateMedias
+	stateTracks
+	stateStart
+	stateExternal
+	stateInternal
+)
+
+type Producer struct {
+	core.Listener
+
+	url      string
+	template string
+
+	conn      core.Producer
+	receivers []*core.Receiver
+	senders   []*core.Receiver
+
+	state    state
+	mu       sync.Mutex
+	workerID int
+}
+
+func (p *Producer) Dial() error {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	if p.state == stateNone {
+		conn, err := GetProducer(p.url)
+		if err != nil {
+			return err
+		}
+
+		p.conn = conn
+		p.state = stateMedias
+	}
+
+	return nil
+}
+
+func (p *Producer) GetMedias() []*core.Media {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	return p.conn.GetMedias()
+}
+
+func (p *Producer) GetTrack(media *core.Media, codec *core.Codec) (*core.Receiver, error) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	if p.state == stateNone {
+		return nil, errors.New("get track from none state")
+	}
+
+	for _, track := range p.receivers {
+		if track.Codec == codec {
+			return track, nil
+		}
+	}
+
+	track, err := p.conn.GetTrack(media, codec)
+	if err != nil {
+		return nil, err
+	}
+
+	p.receivers = append(p.receivers, track)
+
+	if p.state == stateMedias {
+		p.state = stateTracks
+	}
+
+	return track, nil
+}
+
+func (p *Producer) AddTrack(media *core.Media, codec *core.Codec, track *core.Receiver) error {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	if p.state == stateNone {
+		return errors.New("add track from none state")
+	}
+
+	if err := p.conn.(core.Consumer).AddTrack(media, codec, track); err != nil {
+		return err
+	}
+
+	p.senders = append(p.senders, track)
+
+	if p.state == stateMedias {
+		p.state = stateTracks
+	}
+
+	return nil
+}
+
+func (p *Producer) SetSource(s string) {
+	if p.template == "" {
+		p.template = p.url
+	}
+	p.url = strings.Replace(p.template, "{input}", s, 1)
+}
+
+func (p *Producer) MarshalJSON() ([]byte, error) {
+	if p.conn != nil {
+		return json.Marshal(p.conn)
+	}
+
+	info := core.Info{URL: p.url}
+	return json.Marshal(info)
+}
+
+// internals
+
+func (p *Producer) start() {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	if p.state != stateTracks {
+		return
+	}
+
+	log.Debug().Msgf("[streams] start producer url=%s", p.url)
+
+	p.state = stateStart
+	p.workerID++
+
+	go p.worker(p.conn, p.workerID)
+}
+
+func (p *Producer) worker(conn core.Producer, workerID int) {
+	if err := conn.Start(); err != nil {
+		p.mu.Lock()
+		closed := p.workerID != workerID
+		p.mu.Unlock()
+
+		if closed {
+			return
+		}
+
+		log.Warn().Err(err).Str("url", p.url).Caller().Send()
+	}
+
+	p.reconnect(workerID, 0)
+}
+
+func (p *Producer) reconnect(workerID, retry int) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	if p.workerID != workerID {
+		log.Trace().Msgf("[streams] stop reconnect url=%s", p.url)
+		return
+	}
+
+	log.Debug().Msgf("[streams] retry=%d to url=%s", retry, p.url)
+
+	conn, err := GetProducer(p.url)
+	if err != nil {
+		log.Debug().Msgf("[streams] producer=%s", err)
+
+		timeout := time.Minute
+		if retry < 5 {
+			timeout = time.Second
+		} else if retry < 10 {
+			timeout = time.Second * 5
+		} else if retry < 20 {
+			timeout = time.Second * 10
+		}
+
+		time.AfterFunc(timeout, func() {
+			p.reconnect(workerID, retry+1)
+		})
+		return
+	}
+
+	for _, media := range conn.GetMedias() {
+		switch media.Direction {
+		case core.DirectionRecvonly:
+			for _, receiver := range p.receivers {
+				codec := media.MatchCodec(receiver.Codec)
+				if codec == nil {
+					continue
+				}
+
+				track, err := conn.GetTrack(media, codec)
+				if err != nil {
+					continue
+				}
+
+				receiver.Replace(track)
+				break
+			}
+
+		case core.DirectionSendonly:
+			for _, sender := range p.senders {
+				codec := media.MatchCodec(sender.Codec)
+				if codec == nil {
+					continue
+				}
+
+				_ = conn.(core.Consumer).AddTrack(media, codec, sender)
+			}
+		}
+	}
+
+	p.conn = conn
+
+	go p.worker(conn, workerID)
+}
+
+func (p *Producer) stop() {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	switch p.state {
+	case stateExternal:
+		log.Debug().Msgf("[streams] can't stop external producer")
+		return
+	case stateNone:
+		log.Debug().Msgf("[streams] can't stop none producer")
+		return
+	case stateStart:
+		p.workerID++
+	}
+
+	log.Debug().Msgf("[streams] stop producer url=%s", p.url)
+
+	if p.conn != nil {
+		_ = p.conn.Stop()
+		p.conn = nil
+	}
+
+	p.state = stateNone
+	p.receivers = nil
+	p.senders = nil
+}
@@ -0,0 +1,254 @@
+package streams
+
+import (
+	"encoding/json"
+	"errors"
+	"github.com/AlexxIT/go2rtc/pkg/core"
+	"strings"
+	"sync"
+	"sync/atomic"
+)
+
+type Stream struct {
+	producers []*Producer
+	consumers []core.Consumer
+	mu        sync.Mutex
+	requests  int32
+}
+
+func NewStream(source any) *Stream {
+	switch source := source.(type) {
+	case string:
+		s := new(Stream)
+		prod := &Producer{url: source}
+		s.producers = append(s.producers, prod)
+		return s
+	case []any:
+		s := new(Stream)
+		for _, source := range source {
+			prod := &Producer{url: source.(string)}
+			s.producers = append(s.producers, prod)
+		}
+		return s
+	case map[string]any:
+		return NewStream(source["url"])
+	case nil:
+		return new(Stream)
+	default:
+		panic(core.Caller())
+	}
+}
+
+func (s *Stream) SetSource(source string) {
+	for _, prod := range s.producers {
+		prod.SetSource(source)
+	}
+}
+
+func (s *Stream) AddConsumer(cons core.Consumer) (err error) {
+	// support for multiple simultaneous requests from different consumers
+	consN := atomic.AddInt32(&s.requests, 1) - 1
+
+	var statErrors []error
+	var statMedias []*core.Media
+	var statProds []*Producer // matched producers for consumer
+
+	// Step 1. Get consumer medias
+	for _, consMedia := range cons.GetMedias() {
+		log.Trace().Msgf("[streams] check cons=%d media=%s", consN, consMedia)
+
+	producers:
+		for prodN, prod := range s.producers {
+			if err = prod.Dial(); err != nil {
+				log.Trace().Err(err).Msgf("[streams] skip prod=%s", prod.url)
+				statErrors = append(statErrors, err)
+				continue
+			}
+
+			// Step 2. Get producer medias (not tracks yet)
+			for _, prodMedia := range prod.GetMedias() {
+				log.Trace().Msgf("[streams] check prod=%d media=%s", prodN, prodMedia)
+				statMedias = append(statMedias, prodMedia)
+
+				// Step 3. Match consumer/producer codecs list
+				prodCodec, consCodec := prodMedia.MatchMedia(consMedia)
+				if prodCodec == nil {
+					continue
+				}
+
+				var track *core.Receiver
+
+				switch prodMedia.Direction {
+				case core.DirectionRecvonly:
+					log.Trace().Msgf("[streams] match prod=%d => cons=%d", prodN, consN)
+
+					// Step 4. Get recvonly track from producer
+					if track, err = prod.GetTrack(prodMedia, prodCodec); err != nil {
+						log.Info().Err(err).Msg("[streams] can't get track")
+						continue
+					}
+					// Step 5. Add track to consumer
+					if err = cons.AddTrack(consMedia, consCodec, track); err != nil {
+						log.Info().Err(err).Msg("[streams] can't add track")
+						continue
+					}
+
+				case core.DirectionSendonly:
+					log.Trace().Msgf("[streams] match cons=%d => prod=%d", consN, prodN)
+
+					// Step 4. Get recvonly track from consumer (backchannel)
+					if track, err = cons.(core.Producer).GetTrack(consMedia, consCodec); err != nil {
+						log.Info().Err(err).Msg("[streams] can't get track")
+						continue
+					}
+					// Step 5. Add track to producer
+					if err = prod.AddTrack(prodMedia, prodCodec, track); err != nil {
+						log.Info().Err(err).Msg("[streams] can't add track")
+						continue
+					}
+				}
+
+				statProds = append(statProds, prod)
+
+				if !consMedia.MatchAll() {
+					break producers
+				}
+			}
+		}
+	}
+
+	// stop producers if they don't have readers
+	if atomic.AddInt32(&s.requests, -1) == 0 {
+		s.stopProducers()
+	}
+
+	if len(statProds) == 0 {
+		return formatError(statMedias, statErrors)
+	}
+
+	s.mu.Lock()
+	s.consumers = append(s.consumers, cons)
+	s.mu.Unlock()
+
+	// there may be duplicates, but that's not a problem
+	for _, prod := range statProds {
+		prod.start()
+	}
+
+	return nil
+}
+
+func (s *Stream) RemoveConsumer(cons core.Consumer) {
+	_ = cons.Stop()
+
+	s.mu.Lock()
+	for i, consumer := range s.consumers {
+		if consumer == cons {
+			s.consumers = append(s.consumers[:i], s.consumers[i+1:]...)
+			break
+		}
+	}
+	s.mu.Unlock()
+
+	s.stopProducers()
+}
+
+func (s *Stream) AddProducer(prod core.Producer) {
+	producer := &Producer{conn: prod, state: stateExternal}
+	s.mu.Lock()
+	s.producers = append(s.producers, producer)
+	s.mu.Unlock()
+}
+
+func (s *Stream) RemoveProducer(prod core.Producer) {
+	s.mu.Lock()
+	for i, producer := range s.producers {
+		if producer.conn == prod {
+			s.producers = append(s.producers[:i], s.producers[i+1:]...)
+			break
+		}
+	}
+	s.mu.Unlock()
+}
+
+func (s *Stream) stopProducers() {
+	s.mu.Lock()
+producers:
+	for _, producer := range s.producers {
+		for _, track := range producer.receivers {
+			if len(track.Senders()) > 0 {
+				continue producers
+			}
+		}
+		for _, track := range producer.senders {
+			if len(track.Senders()) > 0 {
+				continue producers
+			}
+		}
+		producer.stop()
+	}
+	s.mu.Unlock()
+}
+
+func (s *Stream) MarshalJSON() ([]byte, error) {
+	if !s.mu.TryLock() {
+		log.Warn().Msgf("[streams] json locked")
+		return json.Marshal(nil)
+	}
+
+	var info struct {
+		Producers []*Producer     `json:"producers"`
+		Consumers []core.Consumer `json:"consumers"`
+	}
+	info.Producers = s.producers
+	info.Consumers = s.consumers
+
+	s.mu.Unlock()
+
+	return json.Marshal(info)
+}
+
+func formatError(statMedias []*core.Media, statErrors []error) error {
+	var text string
+
+	for _, media := range statMedias {
+		if media.Direction == core.DirectionRecvonly {
+			continue
+		}
+
+		for _, codec := range media.Codecs {
+			name := codec.Name
+			if name == core.CodecAAC {
+				name = "AAC"
+			}
+			if strings.Contains(text, name) {
+				continue
+			}
+			if len(text) > 0 {
+				text += ","
+			}
+			text += name
+		}
+	}
+
+	if text != "" {
+		return errors.New(text)
+	}
+
+	for _, err := range statErrors {
+		s := err.Error()
+		if strings.Contains(text, s) {
+			continue
+		}
+		if len(text) > 0 {
+			text += ","
+		}
+		text += s
+	}
+
+	if text != "" {
+		return errors.New(text)
+	}
+
+	return errors.New("unknown error")
+}
@@ -0,0 +1,19 @@
+package streams
+
+import (
+	"github.com/stretchr/testify/require"
+	"testing"
+)
+
+func TestTemplate(t *testing.T) {
+	source1 := "does not matter"
+
+	stream1 := New("from_yaml", source1)
+	require.Len(t, streams, 1)
+
+	stream2 := NewTemplate("camera.from_hass", "rtsp://localhost:8554/from_yaml?video")
+
+	require.Equal(t, stream1, stream2)
+	require.Equal(t, stream2.producers[0].url, source1)
+	require.Len(t, streams, 2)
+}