neko/internal/webrtc/track.go

package webrtc

import (
	"errors"
	"fmt"
	"io"
	"sync"

	"github.com/pion/rtcp"
	"github.com/pion/webrtc/v3"
	"github.com/pion/webrtc/v3/pkg/media"
	"github.com/rs/zerolog"

	"github.com/demodesk/neko/pkg/types"
	"github.com/demodesk/neko/pkg/types/codec"
)

type Track struct {
	logger      zerolog.Logger
	track       *webrtc.TrackLocalStaticSample
	paused      bool
	videoAuto   bool
	videoAutoMu sync.RWMutex
	listener    func(sample types.Sample)

	stream   types.StreamSinkManager
	streamMu sync.Mutex

	onRtcp   func(rtcp.Packet)
	onRtcpMu sync.RWMutex

	bitrateChange func(int) (bool, error)
	videoChange   func(string) (bool, error)
}

type option func(*Track)

func WithVideoAuto(auto bool) option {
	return func(t *Track) {
		t.videoAuto = auto
	}
}

func NewTrack(logger zerolog.Logger, codec codec.RTPCodec, connection *webrtc.PeerConnection, opts ...option) (*Track, error) {
	id := codec.Type.String()
	track, err := webrtc.NewTrackLocalStaticSample(codec.Capability, id, "stream")
	if err != nil {
		return nil, err
	}

	logger = logger.With().Str("id", id).Logger()

	t := &Track{
		logger: logger,
		track:  track,
	}

	for _, opt := range opts {
		opt(t)
	}

	t.listener = func(sample types.Sample) {
		if t.paused {
			return
		}

		err := track.WriteSample(media.Sample{
			Data:     sample.Data,
			Duration: sample.Duration,
		})
		if err != nil && !errors.Is(err, io.ErrClosedPipe) {
			logger.Warn().Err(err).Msg("failed to write sample to track")
		}
	}

	sender, err := connection.AddTrack(t.track)
	if err != nil {
		return nil, err
	}

	go t.rtcpReader(sender)

	return t, nil
}

func (t *Track) rtcpReader(sender *webrtc.RTPSender) {
	rtcpBuf := make([]byte, 1500)
	for {
		n, _, err := sender.Read(rtcpBuf)
		if err != nil {
			if err == io.EOF || err == io.ErrClosedPipe {
				return
			}

			t.logger.Err(err).Msg("RTCP read error")
			continue
		}

		packets, err := rtcp.Unmarshal(rtcpBuf[:n])
		if err != nil {
			t.logger.Err(err).Msg("RTCP unmarshal error")
			continue
		}

		t.onRtcpMu.RLock()
		handler := t.onRtcp
		t.onRtcpMu.RUnlock()

		for _, packet := range packets {
			if handler != nil {
				go handler(packet)
			}
		}
	}
}

func (t *Track) SetStream(stream types.StreamSinkManager) (bool, error) {
	t.streamMu.Lock()
	defer t.streamMu.Unlock()

	// if we already listen to the stream, do nothing
	if t.stream == stream {
		return false, nil
	}

	var err error
	if t.stream != nil {
		err = t.stream.MoveListenerTo(&t.listener, stream)
	} else {
		err = stream.AddListener(&t.listener)
	}
	if err != nil {
		return false, err
	}

	t.stream = stream

	return true, nil
}

func (t *Track) RemoveStream() {
	t.streamMu.Lock()
	defer t.streamMu.Unlock()

	if t.stream != nil {
		_ = t.stream.RemoveListener(&t.listener)
		t.stream = nil
	}
}

func (t *Track) SetPaused(paused bool) {
	t.paused = paused
}

func (t *Track) OnRTCP(f func(rtcp.Packet)) {
	t.onRtcpMu.Lock()
	defer t.onRtcpMu.Unlock()

	t.onRtcp = f
}

func (t *Track) SetBitrate(bitrate int) (bool, error) {
	if t.bitrateChange == nil {
		return false, fmt.Errorf("bitrate change not supported")
	}

	return t.bitrateChange(bitrate)
}

func (t *Track) SetVideoID(videoID string) (bool, error) {
	if t.videoChange == nil {
		return false, fmt.Errorf("video change not supported")
	}

	return t.videoChange(videoID)
}

func (t *Track) OnBitrateChange(f func(bitrate int) (bool, error)) {
	t.bitrateChange = f
}

func (t *Track) OnVideoChange(f func(string) (bool, error)) {
	t.videoChange = f
}

func (t *Track) SetVideoAuto(auto bool) {
	t.videoAutoMu.Lock()
	defer t.videoAutoMu.Unlock()
	t.videoAuto = auto
}

func (t *Track) VideoAuto() bool {
	t.videoAutoMu.RLock()
	defer t.videoAutoMu.RUnlock()
	return t.videoAuto
}