diff --git a/pkg/rtc/participant_data_track.go b/pkg/rtc/participant_data_track.go index f025c6452..dba7c61af 100644 --- a/pkg/rtc/participant_data_track.go +++ b/pkg/rtc/participant_data_track.go @@ -25,6 +25,10 @@ import ( "github.com/livekit/protocol/utils/guid" ) +func isValidDataTrackHandle(h uint32) bool { + return h > 0 && h <= 65535 +} + func (p *ParticipantImpl) HandlePublishDataTrackRequest(req *livekit.PublishDataTrackRequest) { if !p.CanPublishData() || !p.params.EnableDataTracks { p.pubLogger.Warnw("no permission to publish data track", nil, "req", logger.Proto(req)) @@ -38,7 +42,7 @@ func (p *ParticipantImpl) HandlePublishDataTrackRequest(req *livekit.PublishData return } - if req.PubHandle == 0 || req.PubHandle > 65535 { + if !isValidDataTrackHandle(req.PubHandle) { p.pubLogger.Warnw("invalid data track handle", nil, "req", logger.Proto(req)) p.sendRequestResponse(&livekit.RequestResponse{ Reason: livekit.RequestResponse_INVALID_HANDLE, @@ -139,6 +143,18 @@ func (p *ParticipantImpl) HandlePublishDataTrackRequest(req *livekit.PublishData } func (p *ParticipantImpl) HandleUnpublishDataTrackRequest(req *livekit.UnpublishDataTrackRequest) { + if !isValidDataTrackHandle(req.PubHandle) { + p.pubLogger.Warnw("invalid data track handle", nil, "req", logger.Proto(req)) + p.sendRequestResponse(&livekit.RequestResponse{ + Reason: livekit.RequestResponse_INVALID_HANDLE, + Message: "handle should be > 0 AND < 65536", + Request: &livekit.RequestResponse_UnpublishDataTrack{ + UnpublishDataTrack: proto.CloneOf(req), + }, + }) + return + } + dt := p.UpDataTrackManager.GetPublishedDataTrack(uint16(req.PubHandle)) if dt == nil { p.pubLogger.Warnw("unpublish data track not found", nil, "req", logger.Proto(req)) diff --git a/pkg/rtc/participant_data_track_test.go b/pkg/rtc/participant_data_track_test.go new file mode 100644 index 000000000..7a419f6a4 --- /dev/null +++ b/pkg/rtc/participant_data_track_test.go @@ -0,0 +1,96 @@ +// Copyright 2026 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 ( + "fmt" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/livekit/livekit-server/pkg/routing/routingfakes" + "github.com/livekit/protocol/livekit" +) + +func TestUnpublishDataTrackRejectsInvalidHandle(t *testing.T) { + for _, test := range []struct { + name string + handle uint32 + }{ + {name: "handle wraps to published track", handle: 65537}, + {name: "zero handle", handle: 0}, + {name: "first out-of-range handle", handle: 65536}, + {name: "maximum uint32 handle", handle: ^uint32(0)}, + } { + t.Run(test.name, func(t *testing.T) { + p := newParticipantForTest("test") + p.params.EnableDataTracks = true + sink := p.params.Sink.(*routingfakes.FakeMessageSink) + + p.HandlePublishDataTrackRequest(&livekit.PublishDataTrackRequest{PubHandle: 1, Name: "test-track"}) + require.NotNil(t, p.UpDataTrackManager.GetPublishedDataTrack(1)) + require.Equal(t, 1, sink.WriteMessageCallCount()) + + p.HandleUnpublishDataTrackRequest(&livekit.UnpublishDataTrackRequest{PubHandle: test.handle}) + + require.Equal(t, 2, sink.WriteMessageCallCount()) + response := sink.WriteMessageArgsForCall(1).(*livekit.SignalResponse).GetRequestResponse() + if assert.NotNil(t, response, "expected INVALID_HANDLE request response for handle %d", test.handle) { + assert.Equal(t, livekit.RequestResponse_INVALID_HANDLE, response.Reason) + } + assert.NotNil(t, p.UpDataTrackManager.GetPublishedDataTrack(1), "published track must remain after handle %d", test.handle) + }) + } +} + +func TestPublishDataTrackRejectsInvalidHandle(t *testing.T) { + for _, handle := range []uint32{0, 65536, 65537, ^uint32(0)} { + t.Run(fmt.Sprint(handle), func(t *testing.T) { + p := newParticipantForTest("test") + p.params.EnableDataTracks = true + sink := p.params.Sink.(*routingfakes.FakeMessageSink) + + p.HandlePublishDataTrackRequest(&livekit.PublishDataTrackRequest{PubHandle: handle, Name: "test-track"}) + + require.Equal(t, 1, sink.WriteMessageCallCount()) + response := lastRequestResponse(t, sink, 0) + require.Equal(t, livekit.RequestResponse_INVALID_HANDLE, response.Reason) + require.Empty(t, p.UpDataTrackManager.GetPublishedDataTracks()) + }) + } +} + +func TestDataTrackValidHandles(t *testing.T) { + for _, handle := range []uint32{1, 32768, 65535} { + t.Run(fmt.Sprint(handle), func(t *testing.T) { + p := newParticipantForTest("test") + p.params.EnableDataTracks = true + sink := p.params.Sink.(*routingfakes.FakeMessageSink) + + p.HandlePublishDataTrackRequest(&livekit.PublishDataTrackRequest{PubHandle: handle, Name: "test-track"}) + + require.NotNil(t, p.UpDataTrackManager.GetPublishedDataTrack(uint16(handle))) + require.Equal(t, 1, sink.WriteMessageCallCount()) + require.NotNil(t, sink.WriteMessageArgsForCall(0).(*livekit.SignalResponse).GetPublishDataTrackResponse()) + + p.HandleUnpublishDataTrackRequest(&livekit.UnpublishDataTrackRequest{PubHandle: handle}) + + require.Empty(t, p.UpDataTrackManager.GetPublishedDataTracks()) + require.Equal(t, 2, sink.WriteMessageCallCount()) + require.NotNil(t, sink.WriteMessageArgsForCall(1).(*livekit.SignalResponse).GetUnpublishDataTrackResponse()) + }) + } +}