Add test for port allocator

This commit is contained in:
Zoltán Papp
2025-02-26 11:19:46 +01:00
parent 8542674a83
commit 309e825d58
4 changed files with 95 additions and 46 deletions

View File

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

View File

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

View File

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

View File

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