snapshot of https://github.com/livekit/livekit.git @ 08c5433449c106943231037d0b892895f6e92d3b (2026-09-15)

pkg/service/turnquota_test.go

@@ -0,0 +1,203 @@
+// 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 service
+
+import (
+	"fmt"
+	"net"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func udpAddr(t *testing.T, port int) net.Addr {
+	t.Helper()
+	return &net.UDPAddr{IP: net.IPv4(10, 0, 0, 1), Port: port}
+}
+
+func TestTURNAllocationQuota_LimitsPerUser(t *testing.T) {
+	q := newTURNAllocationQuota(4)
+
+	// four distinct source 5-tuples for the same participant are allowed
+	for i := 0; i < 4; i++ {
+		require.True(t, q.Allow("participantA", LivekitRealm, udpAddr(t, 5000+i)), "allocation %d", i)
+	}
+	// the fifth is rejected (Pion returns 486 Allocation Quota Reached)
+	require.False(t, q.Allow("participantA", LivekitRealm, udpAddr(t, 6000)))
+}
+
+func TestTURNAllocationQuota_IsolatedPerUser(t *testing.T) {
+	q := newTURNAllocationQuota(2)
+
+	require.True(t, q.Allow("A", LivekitRealm, udpAddr(t, 5000)))
+	require.True(t, q.Allow("A", LivekitRealm, udpAddr(t, 5001)))
+	require.False(t, q.Allow("A", LivekitRealm, udpAddr(t, 5002)))
+
+	// a different participant has its own budget
+	require.True(t, q.Allow("B", LivekitRealm, udpAddr(t, 5000)))
+	require.True(t, q.Allow("B", LivekitRealm, udpAddr(t, 5001)))
+	require.False(t, q.Allow("B", LivekitRealm, udpAddr(t, 5002)))
+}
+
+func TestTURNAllocationQuota_ReleaseFreesSlot(t *testing.T) {
+	q := newTURNAllocationQuota(1)
+
+	addr := udpAddr(t, 5000)
+	require.True(t, q.Allow("A", LivekitRealm, addr))
+	require.False(t, q.Allow("A", LivekitRealm, udpAddr(t, 5001)))
+
+	// once the allocation is deleted the slot is reclaimed
+	q.OnDeleted(addr, nil, "udp", "A", LivekitRealm)
+	require.True(t, q.Allow("A", LivekitRealm, udpAddr(t, 5001)))
+}
+
+func TestTURNAllocationQuota_RetransmitIsIdempotent(t *testing.T) {
+	q := newTURNAllocationQuota(1)
+
+	addr := udpAddr(t, 5000)
+	// same 5-tuple retried multiple times must consume only one slot
+	require.True(t, q.Allow("A", LivekitRealm, addr))
+	require.True(t, q.Allow("A", LivekitRealm, addr))
+	require.True(t, q.Allow("A", LivekitRealm, addr))
+
+	// a different 5-tuple is still over quota
+	require.False(t, q.Allow("A", LivekitRealm, udpAddr(t, 5001)))
+}
+
+func TestTURNAllocationQuota_ForgetsUserWhenEmpty(t *testing.T) {
+	q := newTURNAllocationQuota(1)
+
+	addr := udpAddr(t, 5000)
+	require.True(t, q.Allow("A", LivekitRealm, addr))
+	q.OnDeleted(addr, nil, "udp", "A", LivekitRealm)
+
+	q.mu.Lock()
+	_, tracked := q.users["A"]
+	q.mu.Unlock()
+	require.False(t, tracked, "user with no active allocations should not be retained")
+}
+
+func TestTURNAllocationQuota_UnconfirmedReservationReclaimed(t *testing.T) {
+	// an Allocate that passes the quota but never produces an allocation (e.g.
+	// relay ports exhausted) must not hold its slot forever
+	q := newTURNAllocationQuota(1)
+	q.reservationTTL = 10 * time.Millisecond
+
+	require.True(t, q.Allow("A", LivekitRealm, udpAddr(t, 5000)))
+	require.False(t, q.Allow("A", LivekitRealm, udpAddr(t, 5001)), "at quota while reservation pending")
+
+	require.Eventually(t, func() bool {
+		return q.Allow("A", LivekitRealm, udpAddr(t, 5002))
+	}, time.Second, 5*time.Millisecond, "slot should be reclaimed after TTL")
+
+	// the reclaimed user should not linger once fully drained
+	q.OnDeleted(udpAddr(t, 5002), nil, "udp", "A", LivekitRealm)
+	q.mu.Lock()
+	_, tracked := q.users["A"]
+	q.mu.Unlock()
+	require.False(t, tracked)
+}
+
+func TestTURNAllocationQuota_ConfirmedReservationSurvivesTTL(t *testing.T) {
+	// a confirmed (live) allocation must never be reclaimed by the pending timer
+	q := newTURNAllocationQuota(1)
+	q.reservationTTL = 10 * time.Millisecond
+
+	addr := udpAddr(t, 5000)
+	require.True(t, q.Allow("A", LivekitRealm, addr))
+	q.OnCreated(addr, nil, "udp", "A", LivekitRealm, nil, 0)
+
+	time.Sleep(30 * time.Millisecond) // past the TTL
+
+	require.False(t, q.Allow("A", LivekitRealm, udpAddr(t, 5001)), "confirmed allocation must still hold its slot")
+
+	q.OnDeleted(addr, nil, "udp", "A", LivekitRealm)
+	require.True(t, q.Allow("A", LivekitRealm, udpAddr(t, 5001)))
+}
+
+func TestTURNAllocationQuota_StaleReclaimDoesNotEvictReplacement(t *testing.T) {
+	// a reclaim timer that fires after its slot was released must not delete a
+	// fresh reservation created for the same participant + client address
+	q := newTURNAllocationQuota(1)
+	q.reservationTTL = time.Hour // keep real timers from firing during the test
+
+	addr := udpAddr(t, 5000)
+	key := srcAddrKey(addr)
+
+	require.True(t, q.Allow("A", LivekitRealm, addr))
+	q.mu.Lock()
+	stale := q.users["A"][key]
+	q.mu.Unlock()
+
+	// the original allocation is torn down, then the same 5-tuple is reused and
+	// its replacement allocation is confirmed
+	q.OnDeleted(addr, nil, "udp", "A", LivekitRealm)
+	require.True(t, q.Allow("A", LivekitRealm, addr))
+	q.OnCreated(addr, nil, "udp", "A", LivekitRealm, nil, 0)
+
+	// the original slot's timer fires late: it must be a no-op
+	q.reclaimPending("A", key, stale)
+
+	q.mu.Lock()
+	_, live := q.users["A"][key]
+	q.mu.Unlock()
+	require.True(t, live, "replacement slot must survive a stale reclaim")
+	require.False(t, q.Allow("A", LivekitRealm, udpAddr(t, 5001)), "replacement must still count against quota")
+}
+
+// TestTURNAllocationQuota_ConcurrentAllocatesRespectLimit is the core security
+// property: a burst of concurrent Allocate requests reusing one credential must
+// not be able to race past the limit.
+func TestTURNAllocationQuota_ConcurrentAllocatesRespectLimit(t *testing.T) {
+	const limit = 4
+	q := newTURNAllocationQuota(limit)
+
+	var granted atomic.Int32
+	var wg sync.WaitGroup
+	for i := 0; i < 200; i++ {
+		wg.Add(1)
+		go func(port int) {
+			defer wg.Done()
+			if q.Allow("attacker", LivekitRealm, udpAddr(t, port)) {
+				granted.Add(1)
+			}
+		}(7000 + i)
+	}
+	wg.Wait()
+
+	require.Equal(t, int32(limit), granted.Load(), "concurrent allocations must not exceed the quota")
+}
+
+func TestTURNAllocationQuota_SrcAddrKeyDistinguishesNetwork(t *testing.T) {
+	udp := &net.UDPAddr{IP: net.IPv4(10, 0, 0, 1), Port: 5000}
+	tcp := &net.TCPAddr{IP: net.IPv4(10, 0, 0, 1), Port: 5000}
+	require.NotEqual(t, srcAddrKey(udp), srcAddrKey(tcp))
+	require.Equal(t, "", srcAddrKey(nil))
+}
+
+func BenchmarkTURNAllocationQuota_Allow(b *testing.B) {
+	q := newTURNAllocationQuota(1000000)
+	addrs := make([]net.Addr, 0, 64)
+	for i := 0; i < 64; i++ {
+		addrs = append(addrs, &net.UDPAddr{IP: net.IPv4(10, 0, 0, 1), Port: 5000 + i})
+	}
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		q.Allow(fmt.Sprintf("u%d", i%128), LivekitRealm, addrs[i%len(addrs)])
+	}
+}