diff --git a/client/internal/lazyconn/listener/allocator.go b/client/internal/lazyconn/listener/allocator.go new file mode 100644 index 000000000..752f17aa9 --- /dev/null +++ b/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) +} diff --git a/client/internal/lazyconn/listener/allocator_test.go b/client/internal/lazyconn/listener/allocator_test.go new file mode 100644 index 000000000..7b4fb95f1 --- /dev/null +++ b/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) + } +} diff --git a/client/internal/lazyconn/listener/listener.go b/client/internal/lazyconn/listener/listener.go index 94d9f9667..9db3d50be 100644 --- a/client/internal/lazyconn/listener/listener.go +++ b/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)) { diff --git a/client/internal/lazyconn/listener/manager.go b/client/internal/lazyconn/listener/manager.go index 00493012d..c17ca1472 100644 --- a/client/internal/lazyconn/listener/manager.go +++ b/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)