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 18, 2024
1 parent 3ff9392 commit cd9025a
Show file tree
Hide file tree
Showing 13 changed files with 99 additions and 84 deletions.
4 changes: 3 additions & 1 deletion internal/allocation/allocation.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ type Allocation struct {
Protocol Protocol
TurnSocket net.PacketConn
RelaySocket net.PacketConn
Username string
fiveTuple *FiveTuple
permissionsLock sync.RWMutex
permissions map[string]*Permission
Expand All @@ -45,10 +46,11 @@ type Allocation struct {
}

// NewAllocation creates a new instance of NewAllocation.
func NewAllocation(turnSocket net.PacketConn, fiveTuple *FiveTuple, log logging.LeveledLogger) *Allocation {
func NewAllocation(turnSocket net.PacketConn, fiveTuple *FiveTuple, username string, log logging.LeveledLogger) *Allocation {
return &Allocation{
TurnSocket: turnSocket,
fiveTuple: fiveTuple,
Username: username,
permissions: make(map[string]*Permission, 64),
closed: make(chan interface{}),
log: log,
Expand Down
22 changes: 13 additions & 9 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 @@ -103,9 +103,9 @@ func (m *Manager) CreateAllocation(fiveTuple *FiveTuple, turnSocket net.PacketCo
if a := m.GetAllocation(fiveTuple); a != nil {
return nil, fmt.Errorf("%w: %v", errDupeFiveTuple, fiveTuple)
}
a := NewAllocation(turnSocket, fiveTuple, m.log)
a := NewAllocation(turnSocket, fiveTuple, username, 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
}

if err != nil {
return 0, err
}
Expand Down
65 changes: 35 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 @@ -22,7 +23,7 @@ import (
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 +43,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 +80,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 +114,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 +124,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 +142,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 +175,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, errors.New("unexpected user name")
}
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 +202,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
29 changes: 16 additions & 13 deletions internal/allocation/allocation_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ func TestAllocation(t *testing.T) {
}

func subTestGetPermission(t *testing.T) {
a := NewAllocation(nil, nil, nil)
a := NewAllocation(nil, nil, "", nil)

addr, err := net.ResolveUDPAddr("udp", "127.0.0.1:3478")
if err != nil {
Expand Down Expand Up @@ -88,7 +88,7 @@ func subTestGetPermission(t *testing.T) {
}

func subTestAddPermission(t *testing.T) {
a := NewAllocation(nil, nil, nil)
a := NewAllocation(nil, nil, "", nil)

addr, err := net.ResolveUDPAddr("udp", "127.0.0.1:3478")
if err != nil {
Expand All @@ -107,7 +107,7 @@ func subTestAddPermission(t *testing.T) {
}

func subTestRemovePermission(t *testing.T) {
a := NewAllocation(nil, nil, nil)
a := NewAllocation(nil, nil, "", nil)

addr, err := net.ResolveUDPAddr("udp", "127.0.0.1:3478")
if err != nil {
Expand All @@ -130,7 +130,7 @@ func subTestRemovePermission(t *testing.T) {
}

func subTestAddChannelBind(t *testing.T) {
a := NewAllocation(nil, nil, nil)
a := NewAllocation(nil, nil, "", nil)

addr, err := net.ResolveUDPAddr("udp", "127.0.0.1:3478")
if err != nil {
Expand All @@ -154,7 +154,7 @@ func subTestAddChannelBind(t *testing.T) {
}

func subTestGetChannelByNumber(t *testing.T) {
a := NewAllocation(nil, nil, nil)
a := NewAllocation(nil, nil, "", nil)

addr, err := net.ResolveUDPAddr("udp", "127.0.0.1:3478")
if err != nil {
Expand All @@ -173,7 +173,7 @@ func subTestGetChannelByNumber(t *testing.T) {
}

func subTestGetChannelByAddr(t *testing.T) {
a := NewAllocation(nil, nil, nil)
a := NewAllocation(nil, nil, "", nil)

addr, err := net.ResolveUDPAddr("udp", "127.0.0.1:3478")
if err != nil {
Expand All @@ -193,7 +193,7 @@ func subTestGetChannelByAddr(t *testing.T) {
}

func subTestRemoveChannelBind(t *testing.T) {
a := NewAllocation(nil, nil, nil)
a := NewAllocation(nil, nil, "", nil)

addr, err := net.ResolveUDPAddr("udp", "127.0.0.1:3478")
if err != nil {
Expand All @@ -214,7 +214,7 @@ func subTestRemoveChannelBind(t *testing.T) {
}

func subTestAllocationRefresh(t *testing.T) {
a := NewAllocation(nil, nil, nil)
a := NewAllocation(nil, nil, "", nil)

var wg sync.WaitGroup
wg.Add(1)
Expand All @@ -236,7 +236,7 @@ func subTestAllocationClose(t *testing.T) {
panic(err)
}

a := NewAllocation(nil, nil, nil)
a := NewAllocation(nil, nil, "", nil)
a.RelaySocket = l
// Add mock lifetimeTimer
a.lifetimeTimer = time.AfterFunc(proto.DefaultLifetime, func() {})
Expand All @@ -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 Expand Up @@ -357,7 +360,7 @@ func subTestPacketHandler(t *testing.T) {
}

func subTestResponseCache(t *testing.T) {
a := NewAllocation(nil, nil, nil)
a := NewAllocation(nil, nil, "", nil)
transactionID := [stun.TransactionIDSize]byte{1, 2, 3}
responseAttrs := []stun.Setter{
&proto.Lifetime{
Expand Down
2 changes: 1 addition & 1 deletion internal/allocation/channel_bind_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ func TestChannelBindReset(t *testing.T) {
}

func newChannelBind(lifetime time.Duration) *ChannelBind {
a := NewAllocation(nil, nil, nil)
a := NewAllocation(nil, nil, "", nil)

addr, _ := net.ResolveUDPAddr("udp", "0.0.0.0:0")
c := &ChannelBind{
Expand Down
Loading

0 comments on commit cd9025a

Please sign in to comment.