204 lines
6.8 KiB
Go
204 lines
6.8 KiB
Go
// 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)])
|
|
}
|
|
}
|