kedge/internal/mesh/manager.go

package mesh

import (
	"context"
	"fmt"
	"sync"
	"time"

	"go.uber.org/zap"
	"golang.zx2c4.com/wireguard/wgctrl"
	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"

	"github.com/guildhouse-co/kedge/internal/config"
	"github.com/guildhouse-co/kedge/internal/quartermaster"
)

// Manager manages WireGuard tunnel lifecycle for overlay mode.
type Manager struct {
	cfg    config.MeshConfig
	qm     *quartermaster.Client
	log    *zap.SugaredLogger
	client *wgctrl.Client

	mu    sync.RWMutex
	peers map[string]*Peer // keyed by public key
}

// NewManager creates a new WireGuard mesh manager.
func NewManager(cfg config.MeshConfig, qm *quartermaster.Client, log *zap.SugaredLogger) *Manager {
	return &Manager{
		cfg:   cfg,
		qm:    qm,
		log:   log,
		peers: make(map[string]*Peer),
	}
}

// Run starts the mesh manager loop. It initializes the WireGuard interface,
// configures initial peers, and monitors tunnel health.
func (m *Manager) Run(ctx context.Context) error {
	if !m.cfg.Enabled {
		m.log.Info("overlay mode disabled, mesh manager idle")
		<-ctx.Done()
		return nil
	}

	var err error
	m.client, err = wgctrl.New()
	if err != nil {
		return fmt.Errorf("failed to create wireguard client: %w", err)
	}
	defer m.client.Close()

	// Ensure WireGuard interface exists.
	if err := m.ensureInterface(); err != nil {
		return fmt.Errorf("failed to ensure wg interface: %w", err)
	}

	// Configure initial peers from mesh config.
	if err := m.configureInitialPeers(); err != nil {
		return fmt.Errorf("failed to configure initial peers: %w", err)
	}

	// Health monitoring loop.
	ticker := time.NewTicker(m.cfg.HealthInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			m.log.Info("mesh manager shutting down")
			return nil
		case <-ticker.C:
			m.checkPeerHealth()
		}
	}
}

// PeerCount returns the number of active peers.
func (m *Manager) PeerCount() int {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return len(m.peers)
}

// AddPeer adds or updates a WireGuard peer.
func (m *Manager) AddPeer(pubKey string, endpoint string, allowedIPs []string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	key, err := wgtypes.ParseKey(pubKey)
	if err != nil {
		return fmt.Errorf("invalid public key: %w", err)
	}

	peer := &Peer{
		PublicKey:  pubKey,
		Endpoint:   endpoint,
		AllowedIPs: allowedIPs,
		State:      PeerStateActive,
		LastSeen:   time.Now(),
	}
	m.peers[pubKey] = peer

	_ = key // TODO: Apply to wgctrl config.
	m.log.Infow("peer added", "pubkey", pubKey[:8]+"...", "endpoint", endpoint)
	return nil
}

func (m *Manager) ensureInterface() error {
	// TODO: Create WireGuard interface if it doesn't exist.
	// Use netlink to create the interface, then wgctrl to configure it.
	m.log.Infof("ensuring wireguard interface %s", m.cfg.InterfaceName)
	return nil
}

func (m *Manager) configureInitialPeers() error {
	for _, p := range m.cfg.InitialPeers {
		if err := m.AddPeer(p.PublicKey, p.Endpoint, p.AllowedIPs); err != nil {
			m.log.Warnw("failed to add initial peer", "pubkey", p.PublicKey[:8]+"...", "error", err)
		}
	}
	return nil
}

func (m *Manager) checkPeerHealth() {
	m.mu.Lock()
	defer m.mu.Unlock()

	device, err := m.client.Device(m.cfg.InterfaceName)
	if err != nil {
		m.log.Warnw("failed to get wg device", "error", err)
		return
	}

	for _, wgPeer := range device.Peers {
		pubKey := wgPeer.PublicKey.String()
		peer, ok := m.peers[pubKey]
		if !ok {
			continue
		}

		if !wgPeer.LastHandshakeTime.IsZero() {
			peer.LastSeen = wgPeer.LastHandshakeTime
			peer.State = PeerStateActive
		}

		if time.Since(peer.LastSeen) > m.cfg.DeadPeerTimeout {
			peer.State = PeerStateDead
			m.log.Warnw("peer dead", "pubkey", pubKey[:8]+"...", "last_seen", peer.LastSeen)
		}
	}
}