livekit/pkg/rtc/datatrack.go

// Copyright 2023 LiveKit, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package rtc

import (
	"errors"
	"slices"
	"sync"

	"github.com/frostbyte73/core"
	"google.golang.org/protobuf/proto"

	"github.com/livekit/livekit-server/pkg/rtc/types"
	sfuutils "github.com/livekit/livekit-server/pkg/sfu/utils"
	"github.com/livekit/protocol/datatrack"
	"github.com/livekit/protocol/livekit"
	"github.com/livekit/protocol/logger"
)

var (
	errReceiverClosed = errors.New("datatrack is closed")
)

var _ types.DataTrack = (*DataTrack)(nil)

type DataTrackParams struct {
	Logger              logger.Logger
	ParticipantID       func() livekit.ParticipantID
	ParticipantIdentity livekit.ParticipantIdentity
	BytesTrackStats     *BytesTrackStats
}

type subscribedDataTrack struct {
	subscriber    types.LocalParticipant
	dataDownTrack types.DataDownTrack
}

type DataTrack struct {
	params DataTrackParams

	logger logger.Logger

	lock             sync.Mutex
	dti              *livekit.DataTrackInfo
	subscribedTracks map[livekit.ParticipantID]subscribedDataTrack

	downTrackSpreader *sfuutils.DownTrackSpreader[types.DataTrackSender]

	stats *dataTrackStats

	closed core.Fuse
}

func NewDataTrack(params DataTrackParams, dti *livekit.DataTrackInfo) *DataTrack {
	d := &DataTrack{
		params:           params,
		dti:              dti,
		subscribedTracks: make(map[livekit.ParticipantID]subscribedDataTrack),
	}
	d.logger = params.Logger.WithValues("name", d.Name(), "handle", dti.PubHandle)
	d.downTrackSpreader = sfuutils.NewDownTrackSpreader[types.DataTrackSender](sfuutils.DownTrackSpreaderParams{
		Threshold: 20,
		Logger:    d.logger,
	})
	d.stats = newDataTrackStats(dataTrackStatsParams{Logger: d.logger})
	d.logger.Infow("created data track", "dataTrackInfo", logger.Proto(d.dti))
	return d
}

func (d *DataTrack) Close() {
	d.logger.Infow("closing data track")
	d.closed.Break()

	d.stats.Close()
	if d.params.BytesTrackStats != nil {
		d.params.BytesTrackStats.Stop()
	}
}

func (d *DataTrack) PublisherID() livekit.ParticipantID {
	return d.params.ParticipantID()
}

func (d *DataTrack) PublisherIdentity() livekit.ParticipantIdentity {
	return d.params.ParticipantIdentity
}

func (d *DataTrack) ToProto() *livekit.DataTrackInfo {
	return proto.CloneOf(d.dti)
}

func (d *DataTrack) PubHandle() uint16 {
	return uint16(d.dti.PubHandle)
}

func (d *DataTrack) ID() livekit.TrackID {
	return livekit.TrackID(d.dti.Sid)
}

func (d *DataTrack) Name() string {
	return d.dti.Name
}

func (d *DataTrack) AddSubscriber(sub types.LocalParticipant) (types.DataDownTrack, error) {
	d.lock.Lock()
	defer d.lock.Unlock()

	if _, ok := d.subscribedTracks[sub.ID()]; ok {
		return nil, errAlreadySubscribed
	}

	bytesStats := NewBytesTrackStats(
		sub.GetCountry(),
		d.ID(),
		sub.ID(),
		sub.Kind(),
		sub.KindDetails(),
		sub.GetTelemetryListener,
		sub.GetReporter(),
	)
	dataDownTrack, err := NewDataDownTrack(DataDownTrackParams{
		Logger:           sub.GetLogger().WithValues("trackID", d.ID()),
		SubscriberID:     sub.ID(),
		PublishDataTrack: d,
		Handle:           sub.GetNextSubscribedDataTrackHandle(),
		Transport:        sub.GetDataTrackTransport(),
		BytesTrackStats:  bytesStats,
	})
	if err != nil {
		bytesStats.Stop()
		return nil, err
	}

	d.subscribedTracks[sub.ID()] = subscribedDataTrack{
		subscriber:    sub,
		dataDownTrack: dataDownTrack,
	}
	return dataDownTrack, nil
}

func (d *DataTrack) RemoveSubscriber(subID livekit.ParticipantID) {
	d.lock.Lock()
	subscribedTrack, ok := d.subscribedTracks[subID]
	delete(d.subscribedTracks, subID)
	d.lock.Unlock()

	if ok {
		subscribedTrack.dataDownTrack.Close()
	}
}

func (d *DataTrack) IsSubscriber(subID livekit.ParticipantID) bool {
	d.lock.Lock()
	defer d.lock.Unlock()

	_, ok := d.subscribedTracks[subID]
	return ok
}

func (d *DataTrack) RevokeDisallowedSubscribers(allowedSubscriberIdentities []livekit.ParticipantIdentity) []livekit.ParticipantIdentity {
	var revokedSubscriberIdentities []livekit.ParticipantIdentity

	d.lock.Lock()
	disallowed := make(map[livekit.ParticipantID]livekit.ParticipantIdentity)
	for subID, subscribedTrack := range d.subscribedTracks {
		if IsParticipantExemptFromTrackPermissionsRestrictions(subscribedTrack.subscriber) {
			continue
		}

		if !slices.Contains(allowedSubscriberIdentities, subscribedTrack.subscriber.Identity()) {
			disallowed[subID] = subscribedTrack.subscriber.Identity()
		}
	}
	d.lock.Unlock()

	for subID, subIdentity := range disallowed {
		d.logger.Infow("revoking data track subscription",
			"subscriber", subIdentity,
			"subscriberID", subID,
		)
		d.RemoveSubscriber(subID)
		revokedSubscriberIdentities = append(revokedSubscriberIdentities, subIdentity)
	}

	return revokedSubscriberIdentities
}

func (d *DataTrack) AddDataDownTrack(dts types.DataTrackSender) error {
	if d.closed.IsBroken() {
		return errReceiverClosed
	}

	if d.downTrackSpreader.HasDownTrack(dts.SubscriberID()) {
		d.logger.Infow("subscriberID already exists, replacing data downtrack", "subscriberID", dts.SubscriberID())
	}

	d.downTrackSpreader.Store(dts)
	d.logger.Infow("data downtrack added", "subscriberID", dts.SubscriberID())
	return nil
}

func (d *DataTrack) DeleteDataDownTrack(subscriberID livekit.ParticipantID) {
	d.downTrackSpreader.Free(subscriberID)
	d.logger.Infow("data downtrack deleted", "subscriberID", subscriberID)
}

func (d *DataTrack) HandlePacket(data []byte, packet *datatrack.Packet, arrivalTime int64) {
	d.stats.Update(packet, arrivalTime, len(data))
	if d.params.BytesTrackStats != nil {
		d.params.BytesTrackStats.AddBytes(uint64(len(data)), false)
	}

	d.downTrackSpreader.Broadcast(func(dts types.DataTrackSender) {
		dts.WritePacket(data, packet, arrivalTime)
	})
}