Skip to content

Commit 93ca8f6

Browse files
enobufsenobufs
committed
User-aware relay address generator
Resolves #420
1 parent 3ff9392 commit 93ca8f6

13 files changed

Lines changed: 464 additions & 83 deletions

‎internal/allocation/allocation_manager.go‎

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,19 @@ import (
1212
"github.com/pion/logging"
1313
)
1414

15+
// Context contains contextual information for TURN server allocation tasks.
16+
// It provides necessary context for the allocation of relay addresses and related
17+
// operations during the TURN server's handling of allocation requests.
18+
type Context struct {
19+
Realm string
20+
Username string
21+
}
22+
1523
// ManagerConfig a bag of config params for Manager.
1624
type ManagerConfig struct {
1725
LeveledLogger logging.LeveledLogger
18-
AllocatePacketConn func(network string, requestedPort int) (net.PacketConn, net.Addr, error)
19-
AllocateConn func(network string, requestedPort int) (net.Conn, net.Addr, error)
26+
AllocatePacketConn func(network string, requestedPort int, allocCtx Context) (net.PacketConn, net.Addr, error)
27+
AllocateConn func(network string, requestedPort int, allocCtx Context) (net.Conn, net.Addr, error)
2028
PermissionHandler func(sourceAddr net.Addr, peerIP net.IP) bool
2129
}
2230

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

36-
allocatePacketConn func(network string, requestedPort int) (net.PacketConn, net.Addr, error)
37-
allocateConn func(network string, requestedPort int) (net.Conn, net.Addr, error)
44+
allocatePacketConn func(network string, requestedPort int, allocCtx Context) (net.PacketConn, net.Addr, error)
45+
allocateConn func(network string, requestedPort int, allocCtx Context) (net.Conn, net.Addr, error)
3846
permissionHandler func(sourceAddr net.Addr, peerIP net.IP) bool
3947
}
4048

@@ -86,7 +94,7 @@ func (m *Manager) Close() error {
8694
}
8795

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

108-
conn, relayAddr, err := m.allocatePacketConn("udp4", requestedPort)
116+
conn, relayAddr, err := m.allocatePacketConn("udp4", requestedPort, allocCtx)
109117
if err != nil {
110118
return nil, err
111119
}
@@ -180,9 +188,13 @@ func (m *Manager) GetReservation(reservationToken string) (int, bool) {
180188
}
181189

182190
// GetRandomEvenPort returns a random un-allocated udp4 port
183-
func (m *Manager) GetRandomEvenPort() (int, error) {
191+
func (m *Manager) GetRandomEvenPort(allocCtx Context) (int, error) {
184192
for i := 0; i < 128; i++ {
185-
conn, addr, err := m.allocatePacketConn("udp4", 0)
193+
conn, addr, err := m.allocatePacketConn("udp4", 0, allocCtx)
194+
if err != nil {
195+
return 0, err
196+
}
197+
186198
if err != nil {
187199
return 0, err
188200
}

‎internal/allocation/allocation_manager_test.go‎

Lines changed: 50 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
package allocation
88

99
import (
10+
"errors"
1011
"io"
1112
"math/rand"
1213
"net"
@@ -19,10 +20,15 @@ import (
1920
"github.com/stretchr/testify/assert"
2021
)
2122

23+
var (
24+
errUnexpectedTestRealm = errors.New("unexpected test realm")
25+
errUnexpectedTestUsername = errors.New("unexpected user name")
26+
)
27+
2228
func TestManager(t *testing.T) {
2329
tt := []struct {
2430
name string
25-
f func(*testing.T, net.PacketConn)
31+
f func(*testing.T, net.PacketConn, string, string)
2632
}{
2733
{"CreateInvalidAllocation", subTestCreateInvalidAllocation},
2834
{"CreateAllocation", subTestCreateAllocation},
@@ -42,34 +48,36 @@ func TestManager(t *testing.T) {
4248
for _, tc := range tt {
4349
f := tc.f
4450
t.Run(tc.name, func(t *testing.T) {
45-
f(t, turnSocket)
51+
f(t, turnSocket, "test_realm_1", "test_user_1")
4652
})
4753
}
4854
}
4955

5056
// Test invalid Allocation creations
51-
func subTestCreateInvalidAllocation(t *testing.T, turnSocket net.PacketConn) {
52-
m, err := newTestManager()
57+
func subTestCreateInvalidAllocation(t *testing.T, turnSocket net.PacketConn, realm, username string) {
58+
m, err := newTestManager(realm, username)
5359
assert.NoError(t, err)
5460

55-
if a, err := m.CreateAllocation(nil, turnSocket, 0, proto.DefaultLifetime); a != nil || err == nil {
61+
allocCtx := Context{Realm: realm, Username: username}
62+
if a, err := m.CreateAllocation(nil, turnSocket, 0, proto.DefaultLifetime, allocCtx); a != nil || err == nil {
5663
t.Errorf("Illegally created allocation with nil FiveTuple")
5764
}
58-
if a, err := m.CreateAllocation(randomFiveTuple(), nil, 0, proto.DefaultLifetime); a != nil || err == nil {
65+
if a, err := m.CreateAllocation(randomFiveTuple(), nil, 0, proto.DefaultLifetime, allocCtx); a != nil || err == nil {
5966
t.Errorf("Illegally created allocation with nil turnSocket")
6067
}
61-
if a, err := m.CreateAllocation(randomFiveTuple(), turnSocket, 0, 0); a != nil || err == nil {
68+
if a, err := m.CreateAllocation(randomFiveTuple(), turnSocket, 0, 0, allocCtx); a != nil || err == nil {
6269
t.Errorf("Illegally created allocation with 0 lifetime")
6370
}
6471
}
6572

6673
// Test valid Allocation creations
67-
func subTestCreateAllocation(t *testing.T, turnSocket net.PacketConn) {
68-
m, err := newTestManager()
74+
func subTestCreateAllocation(t *testing.T, turnSocket net.PacketConn, realm, username string) {
75+
m, err := newTestManager(realm, username)
6976
assert.NoError(t, err)
7077

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

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

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

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

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

96-
func subTestDeleteAllocation(t *testing.T, turnSocket net.PacketConn) {
97-
m, err := newTestManager()
105+
func subTestDeleteAllocation(t *testing.T, turnSocket net.PacketConn, realm, username string) {
106+
m, err := newTestManager(realm, username)
98107
assert.NoError(t, err)
99108

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

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

115125
// Test that allocation should be closed if timeout
116-
func subTestAllocationTimeout(t *testing.T, turnSocket net.PacketConn) {
117-
m, err := newTestManager()
126+
func subTestAllocationTimeout(t *testing.T, turnSocket net.PacketConn, realm, username string) {
127+
m, err := newTestManager(realm, username)
118128
assert.NoError(t, err)
119129

120130
allocations := make([]*Allocation, 5)
121131
lifetime := time.Second
122-
132+
allocCtx := Context{Realm: realm, Username: username}
123133
for index := range allocations {
124134
fiveTuple := randomFiveTuple()
125135

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

143153
// Test for manager close
144-
func subTestManagerClose(t *testing.T, turnSocket net.PacketConn) {
145-
m, err := newTestManager()
154+
func subTestManagerClose(t *testing.T, turnSocket net.PacketConn, realm, username string) {
155+
m, err := newTestManager(realm, username)
146156
assert.NoError(t, err)
147157

148158
allocations := make([]*Allocation, 2)
149-
150-
a1, _ := m.CreateAllocation(randomFiveTuple(), turnSocket, 0, time.Second)
159+
allocCtx := Context{Realm: realm, Username: username}
160+
a1, _ := m.CreateAllocation(randomFiveTuple(), turnSocket, 0, time.Second, allocCtx)
151161
allocations[0] = a1
152-
a2, _ := m.CreateAllocation(randomFiveTuple(), turnSocket, 0, time.Minute)
162+
a2, _ := m.CreateAllocation(randomFiveTuple(), turnSocket, 0, time.Minute, allocCtx)
153163
allocations[1] = a2
154164

155165
// Make a1 timeout
@@ -174,21 +184,28 @@ func randomFiveTuple() *FiveTuple {
174184
}
175185
}
176186

177-
func newTestManager() (*Manager, error) {
187+
func newTestManager(expectedRealm, expectedUsername string) (*Manager, error) {
178188
loggerFactory := logging.NewDefaultLoggerFactory()
179189

180190
config := ManagerConfig{
181191
LeveledLogger: loggerFactory.NewLogger("test"),
182-
AllocatePacketConn: func(string, int) (net.PacketConn, net.Addr, error) {
192+
AllocatePacketConn: func(_ string, _ int, allocCtx Context) (net.PacketConn, net.Addr, error) {
193+
if allocCtx.Realm != expectedRealm {
194+
return nil, nil, errUnexpectedTestRealm
195+
}
196+
if allocCtx.Username != expectedUsername {
197+
return nil, nil, errUnexpectedTestUsername
198+
}
183199
conn, err := net.ListenPacket("udp4", "0.0.0.0:0")
184200
if err != nil {
185201
return nil, nil, err
186202
}
187203

188204
return conn, conn.LocalAddr(), nil
189205
},
190-
AllocateConn: func(string, int) (net.Conn, net.Addr, error) { return nil, nil, nil },
206+
AllocateConn: func(string, int, Context) (net.Conn, net.Addr, error) { return nil, nil, nil },
191207
}
208+
192209
return NewManager(config)
193210
}
194211

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

200-
func subTestGetRandomEvenPort(t *testing.T, _ net.PacketConn) {
201-
m, err := newTestManager()
217+
func subTestGetRandomEvenPort(t *testing.T, _ net.PacketConn, realm, username string) {
218+
m, err := newTestManager(realm, username)
202219
assert.NoError(t, err)
203220

204-
port, err := m.GetRandomEvenPort()
221+
allocCtx := Context{Realm: realm, Username: username}
222+
port, err := m.GetRandomEvenPort(allocCtx)
205223
assert.NoError(t, err)
206224
assert.True(t, port > 0)
207225
assert.True(t, port%2 == 0)

‎internal/allocation/allocation_test.go‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -259,9 +259,13 @@ func subTestAllocationClose(t *testing.T) {
259259
}
260260

261261
func subTestPacketHandler(t *testing.T) {
262-
network := "udp"
262+
const (
263+
network = "udp"
264+
testRealm = "test_realm_2"
265+
testUsername = "test_user_2"
266+
)
263267

264-
m, _ := newTestManager()
268+
m, _ := newTestManager(testRealm, testUsername)
265269

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

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

0 commit comments

Comments
 (0)