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 6e85c3f
Show file tree
Hide file tree
Showing 13 changed files with 239 additions and 76 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
47 changes: 38 additions & 9 deletions internal/allocation/allocation_manager.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,9 +15,12 @@ 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) (net.PacketConn, net.Addr, error) // Deprecated: Use AllocatePacketConnForUser instead
AllocateConn func(network string, requestedPort int) (net.Conn, net.Addr, error) // Deprecated: Use AllocateConnForUser instead
PermissionHandler func(sourceAddr net.Addr, peerIP net.IP) bool

AllocatePacketConnForUser func(network string, requestedPort int, username string) (net.PacketConn, net.Addr, error)
AllocateConnForUser func(network string, requestedPort int, username string) (net.Conn, net.Addr, error)
}

type reservation struct {
Expand All @@ -36,14 +39,17 @@ type Manager struct {
allocatePacketConn func(network string, requestedPort int) (net.PacketConn, net.Addr, error)
allocateConn func(network string, requestedPort int) (net.Conn, net.Addr, error)
permissionHandler func(sourceAddr net.Addr, peerIP net.IP) bool

allocatePacketConnForUser func(network string, requestedPort int, username string) (net.PacketConn, net.Addr, error)
allocateConnForUser func(network string, requestedPort int, username string) (net.Conn, net.Addr, error)
}

// NewManager creates a new instance of Manager.
func NewManager(config ManagerConfig) (*Manager, error) {
switch {
case config.AllocatePacketConn == nil:
case config.AllocatePacketConn == nil && config.AllocatePacketConnForUser == nil:

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

View check run for this annotation

Codecov / codecov/patch

internal/allocation/allocation_manager.go#L50 

Added line #L50 was not covered by tests
return nil, errAllocatePacketConnMustBeSet
case config.AllocateConn == nil:
case config.AllocateConn == nil && config.AllocateConnForUser == nil:

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

View check run for this annotation

Codecov / codecov/patch

internal/allocation/allocation_manager.go#L52 

Added line #L52 was not covered by tests
return nil, errAllocateConnMustBeSet
case config.LeveledLogger == nil:
return nil, errLeveledLoggerMustBeSet
Expand All @@ -55,6 +61,9 @@ func NewManager(config ManagerConfig) (*Manager, error) {
allocatePacketConn: config.AllocatePacketConn,
allocateConn: config.AllocateConn,
permissionHandler: config.PermissionHandler,

allocatePacketConnForUser: config.AllocatePacketConnForUser,
allocateConnForUser: config.AllocateConnForUser,
}, nil
}

Expand Down Expand Up @@ -86,7 +95,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 +112,17 @@ 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)
var conn net.PacketConn
var relayAddr net.Addr
var err error

if m.allocatePacketConnForUser != nil {
conn, relayAddr, err = m.allocatePacketConnForUser("udp4", requestedPort, username)
} else {
conn, relayAddr, err = m.allocatePacketConn("udp4", requestedPort) // fallback
}
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -180,9 +197,21 @@ 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)
var conn net.PacketConn
var addr net.Addr
var err error

if m.allocatePacketConnForUser != nil {
conn, addr, err = m.allocatePacketConnForUser("udp4", 0, username)
} else {
conn, addr, err = m.allocatePacketConn("udp4", 0)
}
if err != nil {
return 0, err
}

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

View check run for this annotation

Codecov / codecov/patch

internal/allocation/allocation_manager.go#L212-L213 

Added lines #L212 - L213 were not covered by tests

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

type testManagerMode int

const (
testManagerModeLegacy testManagerMode = iota
testManagerModeUseAllocatePacketConnForUser
)

func (m testManagerMode) String() string {
switch m {
case testManagerModeLegacy:
return "Legacy"
case testManagerModeUseAllocatePacketConnForUser:
return "UseAllocatePacketConnForUser"
default:
return "Unknown"
}
}

const testUsername = "test_user_1"

func TestManager(t *testing.T) {
tt := []struct {
name string
f func(*testing.T, net.PacketConn)
f func(*testing.T, net.PacketConn, testManagerMode)
}{
{"CreateInvalidAllocation", subTestCreateInvalidAllocation},
{"CreateAllocation", subTestCreateAllocation},
Expand All @@ -39,37 +60,39 @@ func TestManager(t *testing.T) {
panic(err)
}

for _, tc := range tt {
f := tc.f
t.Run(tc.name, func(t *testing.T) {
f(t, turnSocket)
})
for _, mode := range []testManagerMode{testManagerModeLegacy, testManagerModeUseAllocatePacketConnForUser} {
for _, tc := range tt {
f := tc.f
t.Run(tc.name+" with mode "+mode.String(), func(t *testing.T) {
f(t, turnSocket, mode)
})
}
}
}

// Test invalid Allocation creations
func subTestCreateInvalidAllocation(t *testing.T, turnSocket net.PacketConn) {
m, err := newTestManager()
func subTestCreateInvalidAllocation(t *testing.T, turnSocket net.PacketConn, mode testManagerMode) {
m, err := newTestManager(mode)
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, testUsername); 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, testUsername); 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, testUsername); 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, mode testManagerMode) {
m, err := newTestManager(mode)
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, testUsername); a == nil || err != nil {
t.Errorf("Failed to create allocation %v %v", a, err)
}

Expand All @@ -79,26 +102,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, mode testManagerMode) {
m, err := newTestManager(mode)
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, testUsername); 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, testUsername); 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, mode testManagerMode) {
m, err := newTestManager(mode)
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, testUsername); a == nil || err != nil {
t.Errorf("Failed to create allocation %v %v", a, err)
}

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

allocations := make([]*Allocation, 5)
Expand All @@ -123,7 +146,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, testUsername)
if err != nil {
t.Errorf("Failed to create allocation with %v", fiveTuple)
}
Expand All @@ -141,15 +164,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, mode testManagerMode) {
m, err := newTestManager(mode)
assert.NoError(t, err)

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

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

// Make a1 timeout
Expand All @@ -174,20 +197,36 @@ func randomFiveTuple() *FiveTuple {
}
}

func newTestManager() (*Manager, error) {
func newTestManager(mode testManagerMode) (*Manager, error) {
loggerFactory := logging.NewDefaultLoggerFactory()

config := ManagerConfig{
LeveledLogger: loggerFactory.NewLogger("test"),
AllocatePacketConn: func(string, int) (net.PacketConn, net.Addr, error) {
}
switch mode {
case testManagerModeLegacy:
config.AllocatePacketConn = func(string, int) (net.PacketConn, net.Addr, error) {
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 },
}
config.AllocateConn = func(string, int) (net.Conn, net.Addr, error) { return nil, nil, nil }
case testManagerModeUseAllocatePacketConnForUser:
config.AllocatePacketConnForUser = func(_ string, _ int, username string) (net.PacketConn, net.Addr, error) {
if username != testUsername {
return nil, nil, errors.New("unexpected user")

Check failure on line 220 in internal/allocation/allocation_manager_test.go

View workflow job for this annotation

GitHub Actions / lint / Go 

do not define dynamic errors, use wrapped static errors instead: "errors.New(\"unexpected user\")" (err113)
}
conn, err := net.ListenPacket("udp4", "0.0.0.0:0")
if err != nil {
return nil, nil, err
}

return conn, conn.LocalAddr(), nil
}
config.AllocateConnForUser = func(string, int, string) (net.Conn, net.Addr, error) { return nil, nil, nil }
}
return NewManager(config)
}
Expand All @@ -197,11 +236,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, mode testManagerMode) {
m, err := newTestManager(mode)
assert.NoError(t, err)

port, err := m.GetRandomEvenPort()
port, err := m.GetRandomEvenPort(testUsername)
assert.NoError(t, err)
assert.True(t, port > 0)
assert.True(t, port%2 == 0)
Expand Down
Loading

0 comments on commit 6e85c3f

Please sign in to comment.