Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

User-aware relay address generator #421

Closed
wants to merge 1 commit into from
Closed

User-aware relay address generator #421

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion internal/allocation/allocation.go
View file
Open in desktop
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
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,9 +15,12 @@
// 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 @@
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 @@
allocatePacketConn: config.AllocatePacketConn,
allocateConn: config.AllocateConn,
permissionHandler: config.PermissionHandler,

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

Expand Down Expand Up @@ -86,7 +95,7 @@
}

// 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 @@
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 @@
}

// 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
Open in desktop
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 @@
"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 @@
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 @@
}

// 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 @@
}

// 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 @@
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 @@
}

// 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 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 @@
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
Loading