diff --git a/pkg/config/config.go b/pkg/config/config.go index 79219dae8..a275dd33d 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -312,28 +312,28 @@ func (l LimitConfig) CheckAttributesSize(attributes map[string]string) bool { return uint32(total) <= l.MaxAttributesSize } -func (l LimitConfig) CheckAsyncAttributesSize(asyncAttributes map[string][]byte) bool { +func (l LimitConfig) CheckAsyncAttributesSize(asyncAttributes []*livekit.DataTrackSchemaDefinition) bool { if l.MaxAsyncAttributesSize == 0 { return true } total := 0 - for k, v := range asyncAttributes { - total += len(k) + len(v) + for _, asyncAttribute := range asyncAttributes { + total += len(asyncAttribute.Id.Name) + len(asyncAttribute.Definition) } return uint32(total) <= l.MaxAsyncAttributesSize } -func (l LimitConfig) CanAddAsyncAttribute(asyncAttributes map[string][]byte, toAddKey string, toAddValue []byte) bool { +func (l LimitConfig) CanAddAsyncAttribute(asyncAttributes []*livekit.DataTrackSchemaDefinition, toAdd *livekit.DataTrackSchemaDefinition) bool { if l.MaxAsyncAttributesSize == 0 { return true } total := 0 - for k, v := range asyncAttributes { - total += len(k) + len(v) + for _, asyncAttribute := range asyncAttributes { + total += len(asyncAttribute.Id.Name) + len(asyncAttribute.Definition) } - return uint32(total+len(toAddKey)+len(toAddValue)) <= l.MaxAsyncAttributesSize + return uint32(total+len(toAdd.Id.Name)+len(toAdd.Definition)) <= l.MaxAsyncAttributesSize } // --------------------------------- diff --git a/pkg/rtc/participant_async_attributes.go b/pkg/rtc/participant_async_attributes.go index aa546f46e..b17bba526 100644 --- a/pkg/rtc/participant_async_attributes.go +++ b/pkg/rtc/participant_async_attributes.go @@ -22,6 +22,7 @@ import ( "github.com/livekit/protocol/livekit" "github.com/livekit/protocol/logger" + "github.com/livekit/protocol/utils" ) type ParticipantAsyncAttributesParams struct { @@ -31,25 +32,25 @@ type ParticipantAsyncAttributesParams struct { type ParticipantAsyncAttributes struct { params ParticipantAsyncAttributesParams lock sync.Mutex - attributes map[string][]byte + attributes map[string]*livekit.DataTrackSchemaDefinition } func NewParticipantAsyncAttributes(params ParticipantAsyncAttributesParams) *ParticipantAsyncAttributes { return &ParticipantAsyncAttributes{ params: params, - attributes: make(map[string][]byte), + attributes: make(map[string]*livekit.DataTrackSchemaDefinition), } } -func (p *ParticipantAsyncAttributes) Add(id *livekit.DataTrackSchemaId, value []byte) { +func (p *ParticipantAsyncAttributes) Add(aa *livekit.DataTrackSchemaDefinition) { p.lock.Lock() defer p.lock.Unlock() - if id == nil { + if aa.Id == nil { return } - p.attributes[ToParticipantAsyncAttributeKey(id)] = value + p.attributes[ToParticipantAsyncAttributeKey(aa.Id)] = utils.CloneProto(aa) } func (p *ParticipantAsyncAttributes) Delete(id *livekit.DataTrackSchemaId) { @@ -71,22 +72,35 @@ func (p *ParticipantAsyncAttributes) Get(id *livekit.DataTrackSchemaId) *livekit return nil } - value, ok := p.attributes[ToParticipantAsyncAttributeKey(id)] + aa, ok := p.attributes[ToParticipantAsyncAttributeKey(id)] if !ok { return nil } - return &livekit.DataTrackSchemaDefinition{ - Id: id, - Definition: value, - } + return aa } -func (p *ParticipantAsyncAttributes) GetAll() map[string][]byte { +func (p *ParticipantAsyncAttributes) GetAll() []*livekit.DataTrackSchemaDefinition { p.lock.Lock() defer p.lock.Unlock() - return p.attributes + all := make([]*livekit.DataTrackSchemaDefinition, 0, len(p.attributes)) + for _, aa := range p.attributes { + all = append(all, utils.CloneProto(aa)) + } + return all +} + +func (p *ParticipantAsyncAttributes) GetAllIDs() []*livekit.DataTrackSchemaId { + p.lock.Lock() + defer p.lock.Unlock() + + ids := make([]*livekit.DataTrackSchemaId, 0, len(p.attributes)) + for _, aa := range p.attributes { + ids = append(ids, utils.CloneProto(aa.Id)) + } + + return ids } // ------------------------------- @@ -95,7 +109,7 @@ func ToParticipantAsyncAttributeKey(id *livekit.DataTrackSchemaId) string { return fmt.Sprintf("%s/%d", id.Name, id.Encoding) } -func FromParticipantAsyncAttributeKey(key string) *livekit.DataTrackSchemaId { +func fromParticipantAsyncAttributeKey(key string) *livekit.DataTrackSchemaId { parts := strings.Split(key, "/") encoding, _ := strconv.Atoi(parts[1]) return &livekit.DataTrackSchemaId{ diff --git a/pkg/rtc/participant_async_attributes_handler.go b/pkg/rtc/participant_async_attributes_handler.go index 993f81635..f5029f61c 100644 --- a/pkg/rtc/participant_async_attributes_handler.go +++ b/pkg/rtc/participant_async_attributes_handler.go @@ -57,11 +57,7 @@ func (p *ParticipantImpl) HandleDefineDataTrackSchemaRequest(req *livekit.Define return } - if !p.params.LimitConfig.CanAddAsyncAttribute( - p.asyncAttributes.GetAll(), - ToParticipantAsyncAttributeKey(req.SchemaDefinition.Id), - req.SchemaDefinition.Definition, - ) { + if !p.params.LimitConfig.CanAddAsyncAttribute(p.asyncAttributes.GetAll(), req.SchemaDefinition) { p.sendRequestResponse(&livekit.RequestResponse{ Reason: livekit.RequestResponse_LIMIT_EXCEEDED, Message: "async attribute definition exceeds limit", @@ -91,7 +87,7 @@ func (p *ParticipantImpl) HandleGetDataTrackSchemaRequest(req *livekit.GetDataTr } func (p *ParticipantImpl) AddDataTrackSchema(definition *livekit.DataTrackSchemaDefinition) { - p.asyncAttributes.Add(definition.Id, definition.Definition) + p.asyncAttributes.Add(definition) } func (p *ParticipantImpl) GetDataTrackSchema(id *livekit.DataTrackSchemaId) *livekit.DataTrackSchemaDefinition { @@ -124,3 +120,7 @@ func (p *ParticipantImpl) ProcessGetDataTrackSchemaRequest(req *livekit.GetDataT p.sendGetDataTrackSchemaResponse(asyncAttribute) } + +func (p *ParticipantImpl) GetAllAsyncAttributeIDs() []*livekit.DataTrackSchemaId { + return p.asyncAttributes.GetAllIDs() +} diff --git a/pkg/rtc/participant_async_attributes_handler_test.go b/pkg/rtc/participant_async_attributes_handler_test.go index 9d8201c12..97761e41f 100644 --- a/pkg/rtc/participant_async_attributes_handler_test.go +++ b/pkg/rtc/participant_async_attributes_handler_test.go @@ -242,7 +242,11 @@ func TestGetDataTrackSchema(t *testing.T) { } require.Nil(t, p.GetDataTrackSchema(id)) - p.asyncAttributes.Add(id, []byte("definition")) + definition := &livekit.DataTrackSchemaDefinition{ + Id: id, + Definition: []byte("definition"), + } + p.asyncAttributes.Add(definition) got := p.GetDataTrackSchema(id) require.NotNil(t, got) require.Equal(t, id.Name, got.Id.Name) @@ -317,3 +321,75 @@ func TestProcessGetDataTrackSchemaRequest(t *testing.T) { require.Equal(t, def, response.GetDataTrackSchemaResponse.SchemaDefinition) }) } + +func TestGetAllAsyncAttributeIDs(t *testing.T) { + t.Run("returns empty slice when no schemas defined", func(t *testing.T) { + p := newParticipantWithAsyncAttributes(t, true, 0) + require.Empty(t, p.GetAllAsyncAttributeIDs()) + }) + + t.Run("returns ids for all defined schemas", func(t *testing.T) { + p := newParticipantWithAsyncAttributes(t, true, 0) + + id1 := &livekit.DataTrackSchemaId{ + Name: "schema-1", + Encoding: livekit.DataTrackSchemaEncoding_DATA_TRACK_SCHEMA_ENCODING_PROTOBUF, + } + id2 := &livekit.DataTrackSchemaId{ + Name: "schema-2", + Encoding: livekit.DataTrackSchemaEncoding_DATA_TRACK_SCHEMA_ENCODING_FLATBUFFER, + } + id3 := &livekit.DataTrackSchemaId{ + Name: "schema-1", + Encoding: livekit.DataTrackSchemaEncoding_DATA_TRACK_SCHEMA_ENCODING_JSON_SCHEMA, + } + + p.AddDataTrackSchema(&livekit.DataTrackSchemaDefinition{Id: id1, Definition: []byte("def-1")}) + p.AddDataTrackSchema(&livekit.DataTrackSchemaDefinition{Id: id2, Definition: []byte("def-2")}) + p.AddDataTrackSchema(&livekit.DataTrackSchemaDefinition{Id: id3, Definition: []byte("def-3")}) + + ids := p.GetAllAsyncAttributeIDs() + require.Len(t, ids, 3) + + got := make(map[string]livekit.DataTrackSchemaEncoding, len(ids)) + for _, id := range ids { + got[id.Name+"|"+id.Encoding.String()] = id.Encoding + } + require.Contains(t, got, "schema-1|"+id1.Encoding.String()) + require.Contains(t, got, "schema-2|"+id2.Encoding.String()) + require.Contains(t, got, "schema-1|"+id3.Encoding.String()) + }) + + t.Run("reflects deletes", func(t *testing.T) { + p := newParticipantWithAsyncAttributes(t, true, 0) + + id := &livekit.DataTrackSchemaId{ + Name: "schema-1", + Encoding: livekit.DataTrackSchemaEncoding_DATA_TRACK_SCHEMA_ENCODING_PROTOBUF, + } + p.AddDataTrackSchema(&livekit.DataTrackSchemaDefinition{Id: id, Definition: []byte("def")}) + require.Len(t, p.GetAllAsyncAttributeIDs(), 1) + + p.asyncAttributes.Delete(id) + require.Empty(t, p.GetAllAsyncAttributeIDs()) + }) + + t.Run("returned ids are cloned and do not alias internal state", func(t *testing.T) { + p := newParticipantWithAsyncAttributes(t, true, 0) + + id := &livekit.DataTrackSchemaId{ + Name: "schema-1", + Encoding: livekit.DataTrackSchemaEncoding_DATA_TRACK_SCHEMA_ENCODING_PROTOBUF, + } + p.AddDataTrackSchema(&livekit.DataTrackSchemaDefinition{Id: id, Definition: []byte("def")}) + + ids := p.GetAllAsyncAttributeIDs() + require.Len(t, ids, 1) + ids[0].Name = "mutated" + + // mutating the returned slice must not affect what's stored + again := p.GetAllAsyncAttributeIDs() + require.Len(t, again, 1) + require.Equal(t, "schema-1", again[0].Name) + }) +} diff --git a/pkg/rtc/participant_async_attributes_test.go b/pkg/rtc/participant_async_attributes_test.go index 370fadb91..e64045108 100644 --- a/pkg/rtc/participant_async_attributes_test.go +++ b/pkg/rtc/participant_async_attributes_test.go @@ -39,7 +39,7 @@ func TestParticipantAsyncAttributes_AddAndGet(t *testing.T) { } value := []byte("definition-bytes") - a.Add(id, value) + a.Add(&livekit.DataTrackSchemaDefinition{Id: id, Definition: value}) got := a.Get(id) require.NotNil(t, got) @@ -55,14 +55,18 @@ func TestParticipantAsyncAttributes_AddOverwrites(t *testing.T) { Name: "schema-1", Encoding: livekit.DataTrackSchemaEncoding_DATA_TRACK_SCHEMA_ENCODING_PROTOBUF, } - a.Add(id, []byte("v1")) - a.Add(id, []byte("v2")) + a.Add(&livekit.DataTrackSchemaDefinition{Id: id, Definition: []byte("v1")}) + a.Add(&livekit.DataTrackSchemaDefinition{Id: id, Definition: []byte("v2")}) got := a.Get(id) require.NotNil(t, got) require.Equal(t, []byte("v2"), got.Definition) require.Len(t, a.GetAll(), 1) + allIDs := a.GetAllIDs() + require.Len(t, allIDs, 1) + require.Equal(t, id.Name, allIDs[0].Name) + require.Equal(t, id.Encoding, allIDs[0].Encoding) } func TestParticipantAsyncAttributes_DifferentEncodingsAreDistinct(t *testing.T) { @@ -77,8 +81,8 @@ func TestParticipantAsyncAttributes_DifferentEncodingsAreDistinct(t *testing.T) Encoding: livekit.DataTrackSchemaEncoding_DATA_TRACK_SCHEMA_ENCODING_JSON_SCHEMA, } - a.Add(idProto, []byte("proto-def")) - a.Add(idJSON, []byte("json-def")) + a.Add(&livekit.DataTrackSchemaDefinition{Id: idProto, Definition: []byte("proto-def")}) + a.Add(&livekit.DataTrackSchemaDefinition{Id: idJSON, Definition: []byte("json-def")}) gotProto := a.Get(idProto) require.NotNil(t, gotProto) @@ -89,6 +93,7 @@ func TestParticipantAsyncAttributes_DifferentEncodingsAreDistinct(t *testing.T) require.Equal(t, []byte("json-def"), gotJSON.Definition) require.Len(t, a.GetAll(), 2) + require.Len(t, a.GetAllIDs(), 2) } func TestParticipantAsyncAttributes_Delete(t *testing.T) { @@ -98,7 +103,10 @@ func TestParticipantAsyncAttributes_Delete(t *testing.T) { Name: "schema-1", Encoding: livekit.DataTrackSchemaEncoding_DATA_TRACK_SCHEMA_ENCODING_PROTOBUF, } - a.Add(id, []byte("definition")) + a.Add(&livekit.DataTrackSchemaDefinition{ + Id: id, + Definition: []byte("definition"), + }) a.Delete(id) require.Nil(t, a.Get(id)) @@ -113,7 +121,7 @@ func TestParticipantAsyncAttributes_NilId(t *testing.T) { a := newTestAsyncAttributes() // nil id should be silently ignored, not panic - a.Add(nil, []byte("definition")) + a.Add(&livekit.DataTrackSchemaDefinition{Definition: []byte("definition")}) require.Empty(t, a.GetAll()) require.Nil(t, a.Get(nil)) @@ -144,13 +152,21 @@ func TestParticipantAsyncAttributes_GetAllContents(t *testing.T) { Encoding: livekit.DataTrackSchemaEncoding_DATA_TRACK_SCHEMA_ENCODING_FLATBUFFER, } - a.Add(id1, []byte("def-1")) - a.Add(id2, []byte("def-2")) + a.Add(&livekit.DataTrackSchemaDefinition{Id: id1, Definition: []byte("def-1")}) + a.Add(&livekit.DataTrackSchemaDefinition{Id: id2, Definition: []byte("def-2")}) all := a.GetAll() require.Len(t, all, 2) - require.Equal(t, []byte("def-1"), all[ToParticipantAsyncAttributeKey(id1)]) - require.Equal(t, []byte("def-2"), all[ToParticipantAsyncAttributeKey(id2)]) + for _, aa := range all { + switch aa.Id.Name { + case "schema-1": + require.Equal(t, []byte("def-1"), aa.Definition) + case "schema-2": + require.Equal(t, []byte("def-2"), aa.Definition) + default: + require.Fail(t, "unexpected name", aa.Id.Name) + } + } } func TestParticipantAsyncAttributes_ConcurrentAccess(t *testing.T) { @@ -169,7 +185,7 @@ func TestParticipantAsyncAttributes_ConcurrentAccess(t *testing.T) { Name: "schema", Encoding: livekit.DataTrackSchemaEncoding(g % 8), } - a.Add(id, []byte("v")) + a.Add(&livekit.DataTrackSchemaDefinition{Id: id, Definition: []byte("v")}) _ = a.Get(id) _ = a.GetAll() if i%3 == 0 { @@ -199,7 +215,7 @@ func TestToKeyFromKey(t *testing.T) { for _, id := range ids { t.Run(id.Name, func(t *testing.T) { key := ToParticipantAsyncAttributeKey(id) - roundTripped := FromParticipantAsyncAttributeKey(key) + roundTripped := fromParticipantAsyncAttributeKey(key) require.Equal(t, id.Name, roundTripped.Name) require.Equal(t, id.Encoding, roundTripped.Encoding) }) diff --git a/pkg/rtc/types/interfaces.go b/pkg/rtc/types/interfaces.go index 9a07cc823..d6baaed88 100644 --- a/pkg/rtc/types/interfaces.go +++ b/pkg/rtc/types/interfaces.go @@ -576,6 +576,8 @@ type LocalParticipant interface { ClearParticipantListener() GetNextSubscribedDataTrackHandle() uint16 + + GetAllAsyncAttributeIDs() []*livekit.DataTrackSchemaId } // --------------------------------------------- diff --git a/pkg/rtc/types/typesfakes/fake_local_participant.go b/pkg/rtc/types/typesfakes/fake_local_participant.go index b05c9983b..d9c7e0a5b 100644 --- a/pkg/rtc/types/typesfakes/fake_local_participant.go +++ b/pkg/rtc/types/typesfakes/fake_local_participant.go @@ -221,6 +221,16 @@ type FakeLocalParticipant struct { getAdaptiveStreamReturnsOnCall map[int]struct { result1 bool } + GetAllAsyncAttributeIDsStub func() []*livekit.DataTrackSchemaId + getAllAsyncAttributeIDsMutex sync.RWMutex + getAllAsyncAttributeIDsArgsForCall []struct { + } + getAllAsyncAttributeIDsReturns struct { + result1 []*livekit.DataTrackSchemaId + } + getAllAsyncAttributeIDsReturnsOnCall map[int]struct { + result1 []*livekit.DataTrackSchemaId + } GetAnswerStub func() (webrtc.SessionDescription, uint32, error) getAnswerMutex sync.RWMutex getAnswerArgsForCall []struct { @@ -2583,6 +2593,59 @@ func (fake *FakeLocalParticipant) GetAdaptiveStreamReturnsOnCall(i int, result1 }{result1} } +func (fake *FakeLocalParticipant) GetAllAsyncAttributeIDs() []*livekit.DataTrackSchemaId { + fake.getAllAsyncAttributeIDsMutex.Lock() + ret, specificReturn := fake.getAllAsyncAttributeIDsReturnsOnCall[len(fake.getAllAsyncAttributeIDsArgsForCall)] + fake.getAllAsyncAttributeIDsArgsForCall = append(fake.getAllAsyncAttributeIDsArgsForCall, struct { + }{}) + stub := fake.GetAllAsyncAttributeIDsStub + fakeReturns := fake.getAllAsyncAttributeIDsReturns + fake.recordInvocation("GetAllAsyncAttributeIDs", []interface{}{}) + fake.getAllAsyncAttributeIDsMutex.Unlock() + if stub != nil { + return stub() + } + if specificReturn { + return ret.result1 + } + return fakeReturns.result1 +} + +func (fake *FakeLocalParticipant) GetAllAsyncAttributeIDsCallCount() int { + fake.getAllAsyncAttributeIDsMutex.RLock() + defer fake.getAllAsyncAttributeIDsMutex.RUnlock() + return len(fake.getAllAsyncAttributeIDsArgsForCall) +} + +func (fake *FakeLocalParticipant) GetAllAsyncAttributeIDsCalls(stub func() []*livekit.DataTrackSchemaId) { + fake.getAllAsyncAttributeIDsMutex.Lock() + defer fake.getAllAsyncAttributeIDsMutex.Unlock() + fake.GetAllAsyncAttributeIDsStub = stub +} + +func (fake *FakeLocalParticipant) GetAllAsyncAttributeIDsReturns(result1 []*livekit.DataTrackSchemaId) { + fake.getAllAsyncAttributeIDsMutex.Lock() + defer fake.getAllAsyncAttributeIDsMutex.Unlock() + fake.GetAllAsyncAttributeIDsStub = nil + fake.getAllAsyncAttributeIDsReturns = struct { + result1 []*livekit.DataTrackSchemaId + }{result1} +} + +func (fake *FakeLocalParticipant) GetAllAsyncAttributeIDsReturnsOnCall(i int, result1 []*livekit.DataTrackSchemaId) { + fake.getAllAsyncAttributeIDsMutex.Lock() + defer fake.getAllAsyncAttributeIDsMutex.Unlock() + fake.GetAllAsyncAttributeIDsStub = nil + if fake.getAllAsyncAttributeIDsReturnsOnCall == nil { + fake.getAllAsyncAttributeIDsReturnsOnCall = make(map[int]struct { + result1 []*livekit.DataTrackSchemaId + }) + } + fake.getAllAsyncAttributeIDsReturnsOnCall[i] = struct { + result1 []*livekit.DataTrackSchemaId + }{result1} +} + func (fake *FakeLocalParticipant) GetAnswer() (webrtc.SessionDescription, uint32, error) { fake.getAnswerMutex.Lock() ret, specificReturn := fake.getAnswerReturnsOnCall[len(fake.getAnswerArgsForCall)]