mirror of
https://github.com/netbirdio/netbird.git
synced 2026-05-07 01:10:03 +00:00
Add test for port allocator
This commit is contained in:
53
client/internal/lazyconn/listener/allocator.go
Normal file
53
client/internal/lazyconn/listener/allocator.go
Normal 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)
|
||||||
|
}
|
||||||
34
client/internal/lazyconn/listener/allocator_test.go
Normal file
34
client/internal/lazyconn/listener/allocator_test.go
Normal 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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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)) {
|
||||||
|
|||||||
@@ -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
|
|
||||||
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 {
|
if err != nil {
|
||||||
log.Debugf("failed to allocate port: %d: %v", addr.Port, err)
|
return fmt.Errorf("failed to bind lazy connection: %v", 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)
|
||||||
|
|||||||
Reference in New Issue
Block a user