Add test for port allocator

client/internal/lazyconn/listener/allocator.go

@@ -0,0 +1,53 @@
+package listener
+
+import (
+	"fmt"
+	"net"
+)
+
+const (
+	retryLimit = 100
+)
+
+var (
+	listenIP      = net.ParseIP("127.0.0.254")
+	ErrNoFreePort = fmt.Errorf("no free port")
+)
+
+// portAllocator lookup for free port and allocate it
+type portAllocator struct {
+	nextFreePort uint16
+}
+
+func newPortAllocator() *portAllocator {
+	return &portAllocator{
+		nextFreePort: 65535,
+	}
+}
+
+func (p *portAllocator) newConn() (*net.UDPConn, *net.UDPAddr, error) {
+	for i := 0; i < retryLimit; i++ {
+		addr := &net.UDPAddr{
+			Port: p.nextPort(),
+			IP:   listenIP,
+		}
+
+		conn, err := net.ListenUDP("udp", addr)
+		if err != nil {
+			// port could be allocated by another process
+			continue
+		}
+
+		return conn, addr, nil
+	}
+	return nil, nil, ErrNoFreePort
+}
+
+func (p *portAllocator) nextPort() int {
+	port := p.nextFreePort
+	p.nextFreePort--
+	if p.nextFreePort == 0 {
+		p.nextFreePort = 65535
+	}
+	return int(port)
+}

client/internal/lazyconn/listener/allocator_test.go

@@ -0,0 +1,34 @@
+package listener
+
+import (
+	"testing"
+)
+
+func Test_portAllocator_newConn(t *testing.T) {
+	pa := newPortAllocator()
+	for i := 65535; i > 65535-10; i-- {
+		conn, addr, err := pa.newConn()
+		if err != nil {
+			t.Errorf("newConn() error = %v, want nil", err)
+		}
+		if addr.Port != i {
+			t.Errorf("newConn() addr.Port = %v, want %d", addr.Port, i)
+		}
+		_ = conn.Close()
+	}
+}
+
+func Test_portAllocator_port_zero(t *testing.T) {
+	pa := newPortAllocator()
+	pa.nextFreePort = 1
+
+	port := pa.nextPort()
+	if port != 1 {
+		t.Errorf("nextPort() = %v, want %d", port, 1)
+	}
+
+	port = pa.nextPort()
+	if port != 65535 {
+		t.Errorf("nextPort() = %v, want %d", port, 65535)
+	}
+}

client/internal/lazyconn/listener/listener.go

@@ -15,17 +15,12 @@ type Listener struct {
 	wg sync.WaitGroup
 }
 
-func NewListener(peerID string, addr *net.UDPAddr) (*Listener, error) {
-	conn, err := net.ListenUDP("udp", addr)
-	if err != nil {
-		return nil, err
-	}
-
+func NewListener(peerID string, conn *net.UDPConn) *Listener {
 	d := &Listener{
 		conn:   conn,
 		peerID: peerID,
 	}
-	return d, nil
+	return d
 }
 
 func (d *Listener) ReadPackets(trigger func(peerID string)) {

client/internal/lazyconn/listener/manager.go

@@ -9,31 +9,12 @@ import (
 	"github.com/netbirdio/netbird/client/internal/lazyconn"
 )
 
-type portGenerator struct {
-	nextFreePort uint16
-}
-
-func newPortGenerator() *portGenerator {
-	return &portGenerator{
-		nextFreePort: 65535,
-	}
-}
-
-func (p *portGenerator) nextPort() int {
-	port := p.nextFreePort
-	p.nextFreePort--
-	if p.nextFreePort == 0 {
-		p.nextFreePort = 65535
-	}
-	return int(port)
-}
-
 type Manager struct {
 	TrafficStartChan chan string
 
 	wgIface lazyconn.WGIface
 
-	portGenerator *portGenerator
+	portGenerator *portAllocator
 	// todo peers add/remove is not thread safe because of the callback function
 	peers map[string]*Listener
 	done  chan struct{}
@@ -43,7 +24,7 @@ func NewManager(wgIface lazyconn.WGIface) *Manager {
 	m := &Manager{
 		TrafficStartChan: make(chan string, 1),
 		wgIface:          wgIface,
-		portGenerator:    newPortGenerator(),
+		portGenerator:    newPortAllocator(),
 		peers:            make(map[string]*Listener),
 		done:             make(chan struct{}),
 	}
@@ -86,26 +67,12 @@ func (m *Manager) Close() {
 }
 
 func (m *Manager) createFakePeer(peerCfg lazyconn.PeerConfig) error {
-	var (
-		listener *Listener
-		err      error
-		addr     *net.UDPAddr
-	)
-	for i := 0; i < 100; i++ {
-		addr = &net.UDPAddr{
-			Port: m.portGenerator.nextPort(),
-			IP:   net.ParseIP("127.0.0.254"),
-		}
-		listener, err = NewListener(peerCfg.PublicKey, addr)
-		if err != nil {
-			log.Debugf("failed to allocate port: %d: %v", addr.Port, err)
-			continue
-		}
+	conn, addr, err := m.portGenerator.newConn()
+	if err != nil {
+		return fmt.Errorf("failed to bind lazy connection: %v", err)
 	}
 
-	if listener == nil {
-		return fmt.Errorf("failed to allocate lazy connection port for: %s", peerCfg.PublicKey)
-	}
+	listener := NewListener(peerCfg.PublicKey, conn)
 
 	if err := m.createEndpoint(peerCfg, addr); err != nil {
 		log.Errorf("failed to create endpoint for %s: %v", peerCfg.PublicKey, err)