Create conn mgr

This commit is contained in:
Zoltán Papp
2025-02-25 23:50:25 +01:00
parent 0457251d09
commit 8542674a83
15 changed files with 184 additions and 125 deletions
@@ -4,20 +4,18 @@ import (
"net"
"sync"
"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
log "github.com/sirupsen/logrus"
)
type Listener struct {
peerID wgtypes.Key
peerID string
conn *net.UDPConn
// todo is not thread safe. If you start the ReadPackets in upper layer in a Go thread then wait for Close() there too
wg sync.WaitGroup
}
func NewListener(peerID wgtypes.Key, addr *net.UDPAddr) (*Listener, error) {
func NewListener(peerID string, addr *net.UDPAddr) (*Listener, error) {
conn, err := net.ListenUDP("udp", addr)
if err != nil {
return nil, err
@@ -30,7 +28,7 @@ func NewListener(peerID wgtypes.Key, addr *net.UDPAddr) (*Listener, error) {
return d, nil
}
func (d *Listener) ReadPackets(trigger func(peerID wgtypes.Key)) {
func (d *Listener) ReadPackets(trigger func(peerID string)) {
d.wg.Add(1)
defer d.wg.Done()
+20 -16
View File
@@ -5,7 +5,6 @@ import (
"net"
log "github.com/sirupsen/logrus"
"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
"github.com/netbirdio/netbird/client/internal/lazyconn"
)
@@ -30,27 +29,28 @@ func (p *portGenerator) nextPort() int {
}
type Manager struct {
TrafficStartChan chan wgtypes.Key
TrafficStartChan chan string
wgIface lazyconn.WGIface
portGenerator *portGenerator
peers map[wgtypes.Key]*Listener
done chan struct{}
// todo peers add/remove is not thread safe because of the callback function
peers map[string]*Listener
done chan struct{}
}
func NewManager(wgIface lazyconn.WGIface) *Manager {
m := &Manager{
TrafficStartChan: make(chan wgtypes.Key, 1),
TrafficStartChan: make(chan string, 1),
wgIface: wgIface,
portGenerator: newPortGenerator(),
peers: make(map[wgtypes.Key]*Listener),
peers: make(map[string]*Listener),
done: make(chan struct{}),
}
return m
}
func (m *Manager) CreateFakePeers(peerCfg lazyconn.PeerConfig) error {
func (m *Manager) CreateFakePeer(peerCfg lazyconn.PeerConfig) error {
if _, ok := m.peers[peerCfg.PublicKey]; ok {
return nil
}
@@ -58,11 +58,11 @@ func (m *Manager) CreateFakePeers(peerCfg lazyconn.PeerConfig) error {
if err := m.createFakePeer(peerCfg); err != nil {
return err
}
log.Debugf("created lazy connection listener for: %s", peerCfg.PublicKey.String())
log.Debugf("created lazy connection listener for: %s", peerCfg.PublicKey)
return nil
}
func (m *Manager) RemovePeer(peerID wgtypes.Key) {
func (m *Manager) RemovePeer(peerID string) {
listener, ok := m.peers[peerID]
if !ok {
return
@@ -70,7 +70,7 @@ func (m *Manager) RemovePeer(peerID wgtypes.Key) {
listener.Close()
if err := m.wgIface.RemovePeer(peerID.String()); err != nil {
if err := m.wgIface.RemovePeer(peerID); err != nil {
log.Warnf("failed to remove fake peer: %v", err)
}
@@ -104,11 +104,11 @@ func (m *Manager) createFakePeer(peerCfg lazyconn.PeerConfig) error {
}
if listener == nil {
return fmt.Errorf("failed to allocate lazy connection port for: %s", peerCfg.PublicKey.String())
return fmt.Errorf("failed to allocate lazy connection port for: %s", peerCfg.PublicKey)
}
if err := m.createEndpoint(peerCfg, addr); err != nil {
log.Errorf("failed to create endpoint for %s: %v", peerCfg.PublicKey.String(), err)
log.Errorf("failed to create endpoint for %s: %v", peerCfg.PublicKey, err)
listener.Close()
return err
}
@@ -119,10 +119,14 @@ func (m *Manager) createFakePeer(peerCfg lazyconn.PeerConfig) error {
return nil
}
func (m *Manager) onTrigger(peerID wgtypes.Key) {
// todo: it is not thread safe, but it is ok if we protect from upper layer
func (m *Manager) onTrigger(peerID string) {
log.Debugf("peer started to send traffic, remove lazy endpoint: %s", peerID)
// todo: it is not thread safe, but it is ok if we protect from upper layer
m.RemovePeer(peerID)
if err := m.wgIface.RemovePeer(peerID); err != nil {
log.Warnf("failed to remove fake peer: %v", err)
}
delete(m.peers, peerID)
select {
case <-m.done:
@@ -131,5 +135,5 @@ func (m *Manager) onTrigger(peerID wgtypes.Key) {
}
func (m *Manager) createEndpoint(peerCfg lazyconn.PeerConfig, endpoint *net.UDPAddr) error {
return m.wgIface.UpdatePeer(peerCfg.PublicKey.String(), peerCfg.AllowedIPs, 0, endpoint, nil)
return m.wgIface.UpdatePeer(peerCfg.PublicKey, peerCfg.AllowedIPs, 0, endpoint, nil)
}
+9 -10
View File
@@ -5,7 +5,6 @@ import (
"sync"
log "github.com/sirupsen/logrus"
"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
"github.com/netbirdio/netbird/client/internal/lazyconn"
"github.com/netbirdio/netbird/client/internal/lazyconn/listener"
@@ -13,11 +12,11 @@ import (
)
type Manager struct {
PeerActivityChan chan wgtypes.Key
PeerActivityChan chan string
watcher *watcher.Watcher
listenerMgr *listener.Manager
managedPeers map[wgtypes.Key]lazyconn.PeerConfig
managedPeers map[string]lazyconn.PeerConfig
watcherWG sync.WaitGroup
mu sync.Mutex
@@ -25,10 +24,10 @@ type Manager struct {
func NewManager(wgIface lazyconn.WGIface) *Manager {
m := &Manager{
PeerActivityChan: make(chan wgtypes.Key, 1),
PeerActivityChan: make(chan string, 1),
watcher: watcher.NewWatcher(wgIface),
listenerMgr: listener.NewManager(wgIface),
managedPeers: make(map[wgtypes.Key]lazyconn.PeerConfig),
managedPeers: make(map[string]lazyconn.PeerConfig),
}
return m
}
@@ -59,7 +58,7 @@ func (m *Manager) Start() {
continue
}
if err := m.listenerMgr.CreateFakePeers(cfg); err != nil {
if err := m.listenerMgr.CreateFakePeer(cfg); err != nil {
log.Errorf("failed to start watch lazy connection tries: %s", err)
}
m.mu.Unlock()
@@ -90,7 +89,7 @@ func (m *Manager) AddPeer(peer lazyconn.PeerConfig) error {
return nil
}
if err := m.listenerMgr.CreateFakePeers(peer); err != nil {
if err := m.listenerMgr.CreateFakePeer(peer); err != nil {
return err
}
@@ -98,7 +97,7 @@ func (m *Manager) AddPeer(peer lazyconn.PeerConfig) error {
return nil
}
func (m *Manager) RemovePeer(peerID wgtypes.Key) bool {
func (m *Manager) RemovePeer(peerID string) bool {
m.mu.Lock()
defer m.mu.Unlock()
@@ -120,10 +119,10 @@ func (m *Manager) Close() {
m.listenerMgr.Close()
m.watcherWG.Wait()
m.managedPeers = make(map[wgtypes.Key]lazyconn.PeerConfig)
m.managedPeers = make(map[string]lazyconn.PeerConfig)
}
func (m *Manager) notifyPeerAction(ctx context.Context, peerID wgtypes.Key) {
func (m *Manager) notifyPeerAction(ctx context.Context, peerID string) {
select {
case <-ctx.Done():
case m.PeerActivityChan <- peerID:
+1 -3
View File
@@ -2,11 +2,9 @@ package lazyconn
import (
"net/netip"
"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
)
type PeerConfig struct {
PublicKey wgtypes.Key
PublicKey string
AllowedIPs []netip.Prefix
}
+7 -8
View File
@@ -6,7 +6,6 @@ import (
"time"
log "github.com/sirupsen/logrus"
"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
"github.com/netbirdio/netbird/client/iface/configurer"
"github.com/netbirdio/netbird/client/internal/lazyconn"
@@ -22,19 +21,19 @@ type rxHistory struct {
}
type Watcher struct {
PeerTimedOutChan chan wgtypes.Key
PeerTimedOutChan chan string
wgIface lazyconn.WGIface
peers map[wgtypes.Key]*rxHistory
peers map[string]*rxHistory
peersMu sync.Mutex
}
func NewWatcher(wgIface lazyconn.WGIface) *Watcher {
return &Watcher{
PeerTimedOutChan: make(chan wgtypes.Key, 1),
PeerTimedOutChan: make(chan string, 1),
wgIface: wgIface,
peers: make(map[wgtypes.Key]*rxHistory),
peers: make(map[string]*rxHistory),
}
}
@@ -57,14 +56,14 @@ func (m *Watcher) Watch(ctx context.Context) {
}
}
func (m *Watcher) AddPeer(peerID wgtypes.Key) {
func (m *Watcher) AddPeer(peerID string) {
m.peersMu.Lock()
defer m.peersMu.Unlock()
m.peers[peerID] = &rxHistory{}
}
func (m *Watcher) RemovePeer(id wgtypes.Key) {
func (m *Watcher) RemovePeer(id string) {
m.peersMu.Lock()
defer m.peersMu.Unlock()
@@ -72,7 +71,7 @@ func (m *Watcher) RemovePeer(id wgtypes.Key) {
}
// Todo: this is a naive implementation, we must to finish it
func (m *Watcher) checkTimeouts(ctx context.Context, allPeersStats map[wgtypes.Key]configurer.WGStats) {
func (m *Watcher) checkTimeouts(ctx context.Context, allPeersStats map[string]configurer.WGStats) {
m.peersMu.Lock()
defer m.peersMu.Unlock()
+1 -1
View File
@@ -11,7 +11,7 @@ import (
)
type WGIface interface {
Transfers() (map[wgtypes.Key]configurer.WGStats, error)
Transfers() (map[string]configurer.WGStats, error)
RemovePeer(peerKey string) error
UpdatePeer(peerKey string, allowedIps []netip.Prefix, keepAlive time.Duration, endpoint *net.UDPAddr, preSharedKey *wgtypes.Key) error
}