Add peer monitor

peermonitor/peermonitor.go

@@ -0,0 +1,232 @@
+package peermonitor
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"sync"
+	"time"
+
+	"github.com/fosrl/olm/wgtester"
+)
+
+// PeerMonitorCallback is the function type for connection status change callbacks
+type PeerMonitorCallback func(siteID int, connected bool, rtt time.Duration)
+
+// PeerMonitor handles monitoring the connection status to multiple WireGuard peers
+type PeerMonitor struct {
+	monitors    map[int]*wgtester.Client
+	callback    PeerMonitorCallback
+	mutex       sync.Mutex
+	running     bool
+	interval    time.Duration
+	timeout     time.Duration
+	maxAttempts int
+}
+
+// NewPeerMonitor creates a new peer monitor with the given callback
+func NewPeerMonitor(callback PeerMonitorCallback) *PeerMonitor {
+	return &PeerMonitor{
+		monitors:    make(map[int]*wgtester.Client),
+		callback:    callback,
+		interval:    5 * time.Second, // Default check interval
+		timeout:     500 * time.Millisecond,
+		maxAttempts: 3,
+	}
+}
+
+// SetInterval changes how frequently peers are checked
+func (pm *PeerMonitor) SetInterval(interval time.Duration) {
+	pm.mutex.Lock()
+	defer pm.mutex.Unlock()
+
+	pm.interval = interval
+
+	// Update interval for all existing monitors
+	for _, client := range pm.monitors {
+		client.SetPacketInterval(interval)
+	}
+}
+
+// SetTimeout changes the timeout for waiting for responses
+func (pm *PeerMonitor) SetTimeout(timeout time.Duration) {
+	pm.mutex.Lock()
+	defer pm.mutex.Unlock()
+
+	pm.timeout = timeout
+
+	// Update timeout for all existing monitors
+	for _, client := range pm.monitors {
+		client.SetTimeout(timeout)
+	}
+}
+
+// SetMaxAttempts changes the maximum number of attempts for TestConnection
+func (pm *PeerMonitor) SetMaxAttempts(attempts int) {
+	pm.mutex.Lock()
+	defer pm.mutex.Unlock()
+
+	pm.maxAttempts = attempts
+
+	// Update max attempts for all existing monitors
+	for _, client := range pm.monitors {
+		client.SetMaxAttempts(attempts)
+	}
+}
+
+// AddPeer adds a new peer to monitor
+func (pm *PeerMonitor) AddPeer(siteID int, endpoint string) error {
+	pm.mutex.Lock()
+	defer pm.mutex.Unlock()
+
+	// Check if we're already monitoring this peer
+	if _, exists := pm.monitors[siteID]; exists {
+		// Update the endpoint instead of creating a new monitor
+		pm.RemovePeer(siteID)
+	}
+
+	// Add UDP port if not present, assuming default WireGuard port
+	if _, _, err := net.SplitHostPort(endpoint); err != nil {
+		endpoint = endpoint + ":51820" // Default WireGuard port
+	}
+
+	client, err := wgtester.NewClient(endpoint)
+	if err != nil {
+		return err
+	}
+
+	// Configure the client with our settings
+	client.SetPacketInterval(pm.interval)
+	client.SetTimeout(pm.timeout)
+	client.SetMaxAttempts(pm.maxAttempts)
+
+	// Store the client
+	pm.monitors[siteID] = client
+
+	// If monitor is already running, start monitoring this peer
+	if pm.running {
+		siteIDCopy := siteID // Create a copy for the closure
+		err = client.StartMonitor(func(status wgtester.ConnectionStatus) {
+			pm.callback(siteIDCopy, status.Connected, status.RTT)
+		})
+	}
+
+	return err
+}
+
+// RemovePeer stops monitoring a peer and removes it from the monitor
+func (pm *PeerMonitor) RemovePeer(siteID int) {
+	pm.mutex.Lock()
+	defer pm.mutex.Unlock()
+
+	client, exists := pm.monitors[siteID]
+	if !exists {
+		return
+	}
+
+	client.StopMonitor()
+	client.Close()
+	delete(pm.monitors, siteID)
+}
+
+// Start begins monitoring all peers
+func (pm *PeerMonitor) Start() {
+	pm.mutex.Lock()
+	defer pm.mutex.Unlock()
+
+	if pm.running {
+		return // Already running
+	}
+
+	pm.running = true
+
+	// Start monitoring all peers
+	for siteID, client := range pm.monitors {
+		siteIDCopy := siteID // Create a copy for the closure
+		client.StartMonitor(func(status wgtester.ConnectionStatus) {
+			pm.callback(siteIDCopy, status.Connected, status.RTT)
+		})
+	}
+}
+
+// Stop stops monitoring all peers
+func (pm *PeerMonitor) Stop() {
+	pm.mutex.Lock()
+	defer pm.mutex.Unlock()
+
+	if !pm.running {
+		return
+	}
+
+	pm.running = false
+
+	// Stop all monitors
+	for _, client := range pm.monitors {
+		client.StopMonitor()
+	}
+}
+
+// Close stops monitoring and cleans up resources
+func (pm *PeerMonitor) Close() {
+	pm.mutex.Lock()
+	defer pm.mutex.Unlock()
+
+	// Stop and close all clients
+	for siteID, client := range pm.monitors {
+		client.StopMonitor()
+		client.Close()
+		delete(pm.monitors, siteID)
+	}
+
+	pm.running = false
+}
+
+// TestPeer tests connectivity to a specific peer
+func (pm *PeerMonitor) TestPeer(siteID int) (bool, time.Duration, error) {
+	pm.mutex.Lock()
+	client, exists := pm.monitors[siteID]
+	pm.mutex.Unlock()
+
+	if !exists {
+		return false, 0, fmt.Errorf("peer with siteID %d not found", siteID)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), pm.timeout*time.Duration(pm.maxAttempts))
+	defer cancel()
+
+	connected, rtt := client.TestConnection(ctx)
+	return connected, rtt, nil
+}
+
+// TestAllPeers tests connectivity to all peers
+func (pm *PeerMonitor) TestAllPeers() map[int]struct {
+	Connected bool
+	RTT       time.Duration
+} {
+	pm.mutex.Lock()
+	peers := make(map[int]*wgtester.Client, len(pm.monitors))
+	for siteID, client := range pm.monitors {
+		peers[siteID] = client
+	}
+	pm.mutex.Unlock()
+
+	results := make(map[int]struct {
+		Connected bool
+		RTT       time.Duration
+	})
+	for siteID, client := range peers {
+		ctx, cancel := context.WithTimeout(context.Background(), pm.timeout*time.Duration(pm.maxAttempts))
+		connected, rtt := client.TestConnection(ctx)
+		cancel()
+
+		results[siteID] = struct {
+			Connected bool
+			RTT       time.Duration
+		}{
+			Connected: connected,
+			RTT:       rtt,
+		}
+	}
+
+	return results
+}