Skip to content

Commit

Permalink
User-aware relay address generator
Browse files Browse the repository at this point in the history 
Resolves #420
  • Loading branch information
@enobufs
enobufs committed Dec 16, 2024
1 parent 3ff9392 commit 93ca8f6
Show file tree
Hide file tree
Showing 13 changed files with 464 additions and 83 deletions.
28 changes: 20 additions & 8 deletions internal/allocation/allocation_manager.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,11 +12,19 @@ import (
"github.com/pion/logging"
)

// Context contains contextual information for TURN server allocation tasks.
// It provides necessary context for the allocation of relay addresses and related
// operations during the TURN server's handling of allocation requests.
type Context struct {
Realm string
Username string
}

// ManagerConfig a bag of config params for Manager.
type ManagerConfig struct {
LeveledLogger logging.LeveledLogger
AllocatePacketConn func(network string, requestedPort int) (net.PacketConn, net.Addr, error)
AllocateConn func(network string, requestedPort int) (net.Conn, net.Addr, error)
AllocatePacketConn func(network string, requestedPort int, allocCtx Context) (net.PacketConn, net.Addr, error)
AllocateConn func(network string, requestedPort int, allocCtx Context) (net.Conn, net.Addr, error)
PermissionHandler func(sourceAddr net.Addr, peerIP net.IP) bool
}

Expand All @@ -33,8 +41,8 @@ type Manager struct {
allocations map[FiveTupleFingerprint]*Allocation
reservations []*reservation

allocatePacketConn func(network string, requestedPort int) (net.PacketConn, net.Addr, error)
allocateConn func(network string, requestedPort int) (net.Conn, net.Addr, error)
allocatePacketConn func(network string, requestedPort int, allocCtx Context) (net.PacketConn, net.Addr, error)
allocateConn func(network string, requestedPort int, allocCtx Context) (net.Conn, net.Addr, error)
permissionHandler func(sourceAddr net.Addr, peerIP net.IP) bool
}

Expand Down Expand Up @@ -86,7 +94,7 @@ func (m *Manager) Close() error {
}

// CreateAllocation creates a new allocation and starts relaying
func (m *Manager) CreateAllocation(fiveTuple *FiveTuple, turnSocket net.PacketConn, requestedPort int, lifetime time.Duration) (*Allocation, error) {
func (m *Manager) CreateAllocation(fiveTuple *FiveTuple, turnSocket net.PacketConn, requestedPort int, lifetime time.Duration, allocCtx Context) (*Allocation, error) {
switch {
case fiveTuple == nil:
return nil, errNilFiveTuple
Expand All @@ -105,7 +113,7 @@ func (m *Manager) CreateAllocation(fiveTuple *FiveTuple, turnSocket net.PacketCo
}
a := NewAllocation(turnSocket, fiveTuple, m.log)

conn, relayAddr, err := m.allocatePacketConn("udp4", requestedPort)
conn, relayAddr, err := m.allocatePacketConn("udp4", requestedPort, allocCtx)
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -180,9 +188,13 @@ func (m *Manager) GetReservation(reservationToken string) (int, bool) {
}

// GetRandomEvenPort returns a random un-allocated udp4 port
func (m *Manager) GetRandomEvenPort() (int, error) {
func (m *Manager) GetRandomEvenPort(allocCtx Context) (int, error) {
for i := 0; i < 128; i++ {
conn, addr, err := m.allocatePacketConn("udp4", 0)
conn, addr, err := m.allocatePacketConn("udp4", 0, allocCtx)
if err != nil {
return 0, err
}

Check warning on line 196 in internal/allocation/allocation_manager.go

View check run for this annotation

Codecov / codecov/patch

internal/allocation/allocation_manager.go#L195-L196 

Added lines #L195 - L196 were not covered by tests

if err != nil {
return 0, err
}
Expand Down
82 changes: 50 additions & 32 deletions internal/allocation/allocation_manager_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
package allocation

import (
"errors"
"io"
"math/rand"
"net"
Expand All @@ -19,10 +20,15 @@ import (
"github.com/stretchr/testify/assert"
)

var (
errUnexpectedTestRealm = errors.New("unexpected test realm")
errUnexpectedTestUsername = errors.New("unexpected user name")
)

func TestManager(t *testing.T) {
tt := []struct {
name string
f func(*testing.T, net.PacketConn)
f func(*testing.T, net.PacketConn, string, string)
}{
{"CreateInvalidAllocation", subTestCreateInvalidAllocation},
{"CreateAllocation", subTestCreateAllocation},
Expand All @@ -42,34 +48,36 @@ func TestManager(t *testing.T) {
for _, tc := range tt {
f := tc.f
t.Run(tc.name, func(t *testing.T) {
f(t, turnSocket)
f(t, turnSocket, "test_realm_1", "test_user_1")
})
}
}

// Test invalid Allocation creations
func subTestCreateInvalidAllocation(t *testing.T, turnSocket net.PacketConn) {
m, err := newTestManager()
func subTestCreateInvalidAllocation(t *testing.T, turnSocket net.PacketConn, realm, username string) {
m, err := newTestManager(realm, username)
assert.NoError(t, err)

if a, err := m.CreateAllocation(nil, turnSocket, 0, proto.DefaultLifetime); a != nil || err == nil {
allocCtx := Context{Realm: realm, Username: username}
if a, err := m.CreateAllocation(nil, turnSocket, 0, proto.DefaultLifetime, allocCtx); a != nil || err == nil {
t.Errorf("Illegally created allocation with nil FiveTuple")
}
if a, err := m.CreateAllocation(randomFiveTuple(), nil, 0, proto.DefaultLifetime); a != nil || err == nil {
if a, err := m.CreateAllocation(randomFiveTuple(), nil, 0, proto.DefaultLifetime, allocCtx); a != nil || err == nil {
t.Errorf("Illegally created allocation with nil turnSocket")
}
if a, err := m.CreateAllocation(randomFiveTuple(), turnSocket, 0, 0); a != nil || err == nil {
if a, err := m.CreateAllocation(randomFiveTuple(), turnSocket, 0, 0, allocCtx); a != nil || err == nil {
t.Errorf("Illegally created allocation with 0 lifetime")
}
}

// Test valid Allocation creations
func subTestCreateAllocation(t *testing.T, turnSocket net.PacketConn) {
m, err := newTestManager()
func subTestCreateAllocation(t *testing.T, turnSocket net.PacketConn, realm, username string) {
m, err := newTestManager(realm, username)
assert.NoError(t, err)

fiveTuple := randomFiveTuple()
if a, err := m.CreateAllocation(fiveTuple, turnSocket, 0, proto.DefaultLifetime); a == nil || err != nil {
allocCtx := Context{Realm: realm, Username: username}
if a, err := m.CreateAllocation(fiveTuple, turnSocket, 0, proto.DefaultLifetime, allocCtx); a == nil || err != nil {
t.Errorf("Failed to create allocation %v %v", a, err)
}

Expand All @@ -79,26 +87,28 @@ func subTestCreateAllocation(t *testing.T, turnSocket net.PacketConn) {
}

// Test that two allocations can't be created with the same FiveTuple
func subTestCreateAllocationDuplicateFiveTuple(t *testing.T, turnSocket net.PacketConn) {
m, err := newTestManager()
func subTestCreateAllocationDuplicateFiveTuple(t *testing.T, turnSocket net.PacketConn, realm, username string) {
m, err := newTestManager(realm, username)
assert.NoError(t, err)

fiveTuple := randomFiveTuple()
if a, err := m.CreateAllocation(fiveTuple, turnSocket, 0, proto.DefaultLifetime); a == nil || err != nil {
allocCtx := Context{Realm: realm, Username: username}
if a, err := m.CreateAllocation(fiveTuple, turnSocket, 0, proto.DefaultLifetime, allocCtx); a == nil || err != nil {
t.Errorf("Failed to create allocation %v %v", a, err)
}

if a, err := m.CreateAllocation(fiveTuple, turnSocket, 0, proto.DefaultLifetime); a != nil || err == nil {
if a, err := m.CreateAllocation(fiveTuple, turnSocket, 0, proto.DefaultLifetime, allocCtx); a != nil || err == nil {
t.Errorf("Was able to create allocation with same FiveTuple twice")
}
}

func subTestDeleteAllocation(t *testing.T, turnSocket net.PacketConn) {
m, err := newTestManager()
func subTestDeleteAllocation(t *testing.T, turnSocket net.PacketConn, realm, username string) {
m, err := newTestManager(realm, username)
assert.NoError(t, err)

fiveTuple := randomFiveTuple()
if a, err := m.CreateAllocation(fiveTuple, turnSocket, 0, proto.DefaultLifetime); a == nil || err != nil {
allocCtx := Context{Realm: realm, Username: username}
if a, err := m.CreateAllocation(fiveTuple, turnSocket, 0, proto.DefaultLifetime, allocCtx); a == nil || err != nil {
t.Errorf("Failed to create allocation %v %v", a, err)
}

Expand All @@ -113,17 +123,17 @@ func subTestDeleteAllocation(t *testing.T, turnSocket net.PacketConn) {
}

// Test that allocation should be closed if timeout
func subTestAllocationTimeout(t *testing.T, turnSocket net.PacketConn) {
m, err := newTestManager()
func subTestAllocationTimeout(t *testing.T, turnSocket net.PacketConn, realm, username string) {
m, err := newTestManager(realm, username)
assert.NoError(t, err)

allocations := make([]*Allocation, 5)
lifetime := time.Second

allocCtx := Context{Realm: realm, Username: username}
for index := range allocations {
fiveTuple := randomFiveTuple()

a, err := m.CreateAllocation(fiveTuple, turnSocket, 0, lifetime)
a, err := m.CreateAllocation(fiveTuple, turnSocket, 0, lifetime, allocCtx)
if err != nil {
t.Errorf("Failed to create allocation with %v", fiveTuple)
}
Expand All @@ -141,15 +151,15 @@ func subTestAllocationTimeout(t *testing.T, turnSocket net.PacketConn) {
}

// Test for manager close
func subTestManagerClose(t *testing.T, turnSocket net.PacketConn) {
m, err := newTestManager()
func subTestManagerClose(t *testing.T, turnSocket net.PacketConn, realm, username string) {
m, err := newTestManager(realm, username)
assert.NoError(t, err)

allocations := make([]*Allocation, 2)

a1, _ := m.CreateAllocation(randomFiveTuple(), turnSocket, 0, time.Second)
allocCtx := Context{Realm: realm, Username: username}
a1, _ := m.CreateAllocation(randomFiveTuple(), turnSocket, 0, time.Second, allocCtx)
allocations[0] = a1
a2, _ := m.CreateAllocation(randomFiveTuple(), turnSocket, 0, time.Minute)
a2, _ := m.CreateAllocation(randomFiveTuple(), turnSocket, 0, time.Minute, allocCtx)
allocations[1] = a2

// Make a1 timeout
Expand All @@ -174,21 +184,28 @@ func randomFiveTuple() *FiveTuple {
}
}

func newTestManager() (*Manager, error) {
func newTestManager(expectedRealm, expectedUsername string) (*Manager, error) {
loggerFactory := logging.NewDefaultLoggerFactory()

config := ManagerConfig{
LeveledLogger: loggerFactory.NewLogger("test"),
AllocatePacketConn: func(string, int) (net.PacketConn, net.Addr, error) {
AllocatePacketConn: func(_ string, _ int, allocCtx Context) (net.PacketConn, net.Addr, error) {
if allocCtx.Realm != expectedRealm {
return nil, nil, errUnexpectedTestRealm
}
if allocCtx.Username != expectedUsername {
return nil, nil, errUnexpectedTestUsername
}
conn, err := net.ListenPacket("udp4", "0.0.0.0:0")
if err != nil {
return nil, nil, err
}

return conn, conn.LocalAddr(), nil
},
AllocateConn: func(string, int) (net.Conn, net.Addr, error) { return nil, nil, nil },
AllocateConn: func(string, int, Context) (net.Conn, net.Addr, error) { return nil, nil, nil },
}

return NewManager(config)
}

Expand All @@ -197,11 +214,12 @@ func isClose(conn io.Closer) bool {
return closeErr != nil && strings.Contains(closeErr.Error(), "use of closed network connection")
}

func subTestGetRandomEvenPort(t *testing.T, _ net.PacketConn) {
m, err := newTestManager()
func subTestGetRandomEvenPort(t *testing.T, _ net.PacketConn, realm, username string) {
m, err := newTestManager(realm, username)
assert.NoError(t, err)

port, err := m.GetRandomEvenPort()
allocCtx := Context{Realm: realm, Username: username}
port, err := m.GetRandomEvenPort(allocCtx)
assert.NoError(t, err)
assert.True(t, port > 0)
assert.True(t, port%2 == 0)
Expand Down
10 changes: 7 additions & 3 deletions internal/allocation/allocation_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -259,9 +259,13 @@ func subTestAllocationClose(t *testing.T) {
}

func subTestPacketHandler(t *testing.T) {
network := "udp"
const (
network = "udp"
testRealm = "test_realm_2"
testUsername = "test_user_2"
)

m, _ := newTestManager()
m, _ := newTestManager(testRealm, testUsername)

// TURN server initialization
turnSocket, err := net.ListenPacket(network, "127.0.0.1:0")
Expand Down Expand Up @@ -292,7 +296,7 @@ func subTestPacketHandler(t *testing.T) {
a, err := m.CreateAllocation(&FiveTuple{
SrcAddr: clientListener.LocalAddr(),
DstAddr: turnSocket.LocalAddr(),
}, turnSocket, 0, proto.DefaultLifetime)
}, turnSocket, 0, proto.DefaultLifetime, Context{Realm: testRealm, Username: testUsername})

assert.Nil(t, err, "should succeed")

Expand Down
Loading

0 comments on commit 93ca8f6

Please sign in to comment.