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 Nov 20, 2024
1 parent 3ff9392 commit 5cec847
Show file tree
Hide file tree
Showing 13 changed files with 462 additions and 81 deletions.
20 changes: 12 additions & 8 deletions internal/allocation/allocation_manager.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,8 @@ import (
// 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, username string) (net.PacketConn, net.Addr, error)
AllocateConn func(network string, requestedPort int, username string) (net.Conn, net.Addr, error)
PermissionHandler func(sourceAddr net.Addr, peerIP net.IP) bool
}

Expand All @@ -33,8 +33,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, username string) (net.PacketConn, net.Addr, error)
allocateConn func(network string, requestedPort int, username string) (net.Conn, net.Addr, error)
permissionHandler func(sourceAddr net.Addr, peerIP net.IP) bool
}

Expand Down Expand Up @@ -86,7 +86,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, username string) (*Allocation, error) {
switch {
case fiveTuple == nil:
return nil, errNilFiveTuple
Expand All @@ -105,7 +105,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, username)
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -180,9 +180,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(username string) (int, error) {
for i := 0; i < 128; i++ {
conn, addr, err := m.allocatePacketConn("udp4", 0)
conn, addr, err := m.allocatePacketConn("udp4", 0, username)
if err != nil {
return 0, err
}

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

View check run for this annotation

Codecov / codecov/patch

internal/allocation/allocation_manager.go#L187-L188 

Added lines #L187 - L188 were not covered by tests

if err != nil {
return 0, err
}
Expand Down
67 changes: 37 additions & 30 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,12 @@ import (
"github.com/stretchr/testify/assert"
)

var 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)
}{
{"CreateInvalidAllocation", subTestCreateInvalidAllocation},
{"CreateAllocation", subTestCreateAllocation},
Expand All @@ -42,34 +45,34 @@ 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_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, username string) {
m, err := newTestManager(username)
assert.NoError(t, err)

if a, err := m.CreateAllocation(nil, turnSocket, 0, proto.DefaultLifetime); a != nil || err == nil {
if a, err := m.CreateAllocation(nil, turnSocket, 0, proto.DefaultLifetime, username); 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, username); 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, username); 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, username string) {
m, err := newTestManager(username)
assert.NoError(t, err)

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

Expand All @@ -79,26 +82,26 @@ 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, username string) {
m, err := newTestManager(username)
assert.NoError(t, err)

fiveTuple := randomFiveTuple()
if a, err := m.CreateAllocation(fiveTuple, turnSocket, 0, proto.DefaultLifetime); a == nil || err != nil {
if a, err := m.CreateAllocation(fiveTuple, turnSocket, 0, proto.DefaultLifetime, username); 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, username); 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, username string) {
m, err := newTestManager(username)
assert.NoError(t, err)

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

Expand All @@ -113,8 +116,8 @@ 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, username string) {
m, err := newTestManager(username)
assert.NoError(t, err)

allocations := make([]*Allocation, 5)
Expand All @@ -123,7 +126,7 @@ func subTestAllocationTimeout(t *testing.T, turnSocket net.PacketConn) {
for index := range allocations {
fiveTuple := randomFiveTuple()

a, err := m.CreateAllocation(fiveTuple, turnSocket, 0, lifetime)
a, err := m.CreateAllocation(fiveTuple, turnSocket, 0, lifetime, username)
if err != nil {
t.Errorf("Failed to create allocation with %v", fiveTuple)
}
Expand All @@ -141,15 +144,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, username string) {
m, err := newTestManager(username)
assert.NoError(t, err)

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

a1, _ := m.CreateAllocation(randomFiveTuple(), turnSocket, 0, time.Second)
a1, _ := m.CreateAllocation(randomFiveTuple(), turnSocket, 0, time.Second, username)
allocations[0] = a1
a2, _ := m.CreateAllocation(randomFiveTuple(), turnSocket, 0, time.Minute)
a2, _ := m.CreateAllocation(randomFiveTuple(), turnSocket, 0, time.Minute, username)
allocations[1] = a2

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

func newTestManager() (*Manager, error) {
func newTestManager(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, username string) (net.PacketConn, net.Addr, error) {
if 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, string) (net.Conn, net.Addr, error) { return nil, nil, nil },
}

return NewManager(config)
}

Expand All @@ -197,11 +204,11 @@ 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, username string) {
m, err := newTestManager(username)
assert.NoError(t, err)

port, err := m.GetRandomEvenPort()
port, err := m.GetRandomEvenPort(username)
assert.NoError(t, err)
assert.True(t, port > 0)
assert.True(t, port%2 == 0)
Expand Down
9 changes: 6 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,12 @@ func subTestAllocationClose(t *testing.T) {
}

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

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

// TURN server initialization
turnSocket, err := net.ListenPacket(network, "127.0.0.1:0")
Expand Down Expand Up @@ -292,7 +295,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, testUsername)

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

Expand Down
31 changes: 16 additions & 15 deletions internal/server/turn.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,8 @@ func handleAllocateRequest(r Request, m *stun.Message) error {
// mechanism of [https://tools.ietf.org/html/rfc5389#section-10.2.2]
// unless the client and server agree to use another mechanism through
// some procedure outside the scope of this document.
messageIntegrity, hasAuth, err := authenticateRequest(r, m, stun.MethodAllocate)
if !hasAuth {
authResult, err := authenticateRequest(r, m, stun.MethodAllocate)
if !authResult.hasAuth {
return err
}

Expand All @@ -51,7 +51,7 @@ func handleAllocateRequest(r Request, m *stun.Message) error {
return buildAndSendErr(r.Conn, r.SrcAddr, errRelayAlreadyAllocatedForFiveTuple, msg...)
}
// A retry allocation
msg := buildMsg(m.TransactionID, stun.NewType(stun.MethodAllocate, stun.ClassSuccessResponse), append(attrs, messageIntegrity)...)
msg := buildMsg(m.TransactionID, stun.NewType(stun.MethodAllocate, stun.ClassSuccessResponse), append(attrs, authResult.messageIntegrity)...)

Check warning on line 54 in internal/server/turn.go

View check run for this annotation

Codecov / codecov/patch

internal/server/turn.go#L54 

Added line #L54 was not covered by tests
return buildAndSend(r.Conn, r.SrcAddr, msg...)
}

Expand Down Expand Up @@ -104,7 +104,7 @@ func handleAllocateRequest(r Request, m *stun.Message) error {
var evenPort proto.EvenPort
if err = evenPort.GetFrom(m); err == nil {
var randomPort int
randomPort, err = r.AllocationManager.GetRandomEvenPort()
randomPort, err = r.AllocationManager.GetRandomEvenPort(authResult.username)

Check warning on line 107 in internal/server/turn.go

View check run for this annotation

Codecov / codecov/patch

internal/server/turn.go#L107 

Added line #L107 was not covered by tests
if err != nil {
return buildAndSendErr(r.Conn, r.SrcAddr, err, insufficientCapacityMsg...)
}
Expand All @@ -131,7 +131,8 @@ func handleAllocateRequest(r Request, m *stun.Message) error {
fiveTuple,
r.Conn,
requestedPort,
lifetimeDuration)
lifetimeDuration,
authResult.username)
if err != nil {
return buildAndSendErr(r.Conn, r.SrcAddr, err, insufficientCapacityMsg...)
}
Expand Down Expand Up @@ -177,16 +178,16 @@ func handleAllocateRequest(r Request, m *stun.Message) error {
responseAttrs = append(responseAttrs, proto.ReservationToken([]byte(reservationToken)))
}

msg := buildMsg(m.TransactionID, stun.NewType(stun.MethodAllocate, stun.ClassSuccessResponse), append(responseAttrs, messageIntegrity)...)
msg := buildMsg(m.TransactionID, stun.NewType(stun.MethodAllocate, stun.ClassSuccessResponse), append(responseAttrs, authResult.messageIntegrity)...)
a.SetResponseCache(m.TransactionID, responseAttrs)
return buildAndSend(r.Conn, r.SrcAddr, msg...)
}

func handleRefreshRequest(r Request, m *stun.Message) error {
r.Log.Debugf("Received RefreshRequest from %s", r.SrcAddr)

messageIntegrity, hasAuth, err := authenticateRequest(r, m, stun.MethodRefresh)
if !hasAuth {
authResult, err := authenticateRequest(r, m, stun.MethodRefresh)
if !authResult.hasAuth {
return err
}

Expand All @@ -212,7 +213,7 @@ func handleRefreshRequest(r Request, m *stun.Message) error {
&proto.Lifetime{
Duration: lifetimeDuration,
},
messageIntegrity,
authResult.messageIntegrity,
}...)...)
}

Expand All @@ -228,8 +229,8 @@ func handleCreatePermissionRequest(r Request, m *stun.Message) error {
return fmt.Errorf("%w %v:%v", errNoAllocationFound, r.SrcAddr, r.Conn.LocalAddr())
}

messageIntegrity, hasAuth, err := authenticateRequest(r, m, stun.MethodCreatePermission)
if !hasAuth {
authResult, err := authenticateRequest(r, m, stun.MethodCreatePermission)
if !authResult.hasAuth {
return err
}

Expand Down Expand Up @@ -267,7 +268,7 @@ func handleCreatePermissionRequest(r Request, m *stun.Message) error {
respClass = stun.ClassErrorResponse
}

return buildAndSend(r.Conn, r.SrcAddr, buildMsg(m.TransactionID, stun.NewType(stun.MethodCreatePermission, respClass), []stun.Setter{messageIntegrity}...)...)
return buildAndSend(r.Conn, r.SrcAddr, buildMsg(m.TransactionID, stun.NewType(stun.MethodCreatePermission, respClass), []stun.Setter{authResult.messageIntegrity}...)...)
}

func handleSendIndication(r Request, m *stun.Message) error {
Expand Down Expand Up @@ -317,8 +318,8 @@ func handleChannelBindRequest(r Request, m *stun.Message) error {

badRequestMsg := buildMsg(m.TransactionID, stun.NewType(stun.MethodChannelBind, stun.ClassErrorResponse), &stun.ErrorCodeAttribute{Code: stun.CodeBadRequest})

messageIntegrity, hasAuth, err := authenticateRequest(r, m, stun.MethodChannelBind)
if !hasAuth {
authResult, err := authenticateRequest(r, m, stun.MethodChannelBind)
if !authResult.hasAuth {
return err
}

Expand Down Expand Up @@ -351,7 +352,7 @@ func handleChannelBindRequest(r Request, m *stun.Message) error {
return buildAndSendErr(r.Conn, r.SrcAddr, err, badRequestMsg...)
}

return buildAndSend(r.Conn, r.SrcAddr, buildMsg(m.TransactionID, stun.NewType(stun.MethodChannelBind, stun.ClassSuccessResponse), []stun.Setter{messageIntegrity}...)...)
return buildAndSend(r.Conn, r.SrcAddr, buildMsg(m.TransactionID, stun.NewType(stun.MethodChannelBind, stun.ClassSuccessResponse), []stun.Setter{authResult.messageIntegrity}...)...)
}

func handleChannelData(r Request, c *proto.ChannelData) error {
Expand Down
6 changes: 3 additions & 3 deletions internal/server/turn_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -64,15 +64,15 @@ func TestAllocationLifeTime(t *testing.T) {
logger := logging.NewDefaultLoggerFactory().NewLogger("turn")

allocationManager, err := allocation.NewManager(allocation.ManagerConfig{
AllocatePacketConn: func(network string, _ int) (net.PacketConn, net.Addr, error) {
AllocatePacketConn: func(network string, _ int, _ string) (net.PacketConn, net.Addr, error) {
conn, listenErr := net.ListenPacket(network, "0.0.0.0:0")
if err != nil {
return nil, nil, listenErr
}

return conn, conn.LocalAddr(), nil
},
AllocateConn: func(string, int) (net.Conn, net.Addr, error) {
AllocateConn: func(string, int, string) (net.Conn, net.Addr, error) {
return nil, nil, nil
},
LeveledLogger: logger,
Expand All @@ -97,7 +97,7 @@ func TestAllocationLifeTime(t *testing.T) {

fiveTuple := &allocation.FiveTuple{SrcAddr: r.SrcAddr, DstAddr: r.Conn.LocalAddr(), Protocol: allocation.UDP}

_, err = r.AllocationManager.CreateAllocation(fiveTuple, r.Conn, 0, time.Hour)
_, err = r.AllocationManager.CreateAllocation(fiveTuple, r.Conn, 0, time.Hour, "")
assert.NoError(t, err)

assert.NotNil(t, r.AllocationManager.GetAllocation(fiveTuple))
Expand Down
Loading

0 comments on commit 5cec847

Please sign in to comment.