zeromesh/internal/controller/controller.go

package controller

import (
	"context"
	"fmt"
	"log/slog"
	"net"
	"sync"
	"time"

	"zeromesh/internal/config"
	"zeromesh/internal/database"
	"zeromesh/internal/identity"
	"zeromesh/internal/model"
	"zeromesh/internal/repository"
	"zeromesh/internal/vl1"
)

type Controller struct {
	cfg       *config.Config
	db        *database.DB
	identity  *identity.Identity
	peers     *vl1.PeerManager
	transport *vl1.Transport
	netRepo   *repository.NetworkRepo
	nodeRepo  *repository.NodeRepo
	wsClients map[string]net.Conn
	mu        sync.RWMutex
	log       *slog.Logger
	ctx       context.Context
	cancel    context.CancelFunc
	wg        sync.WaitGroup
}

func New(cfg *config.Config, db *database.DB, log *slog.Logger) (*Controller, error) {
	id, err := identity.LoadOrGenerate("./data/controller.identity")
	if err != nil {
		return nil, fmt.Errorf("load controller identity: %w", err)
	}
	ctx, cancel := context.WithCancel(context.Background())

	return &Controller{
		cfg:       cfg,
		db:        db,
		identity:  id,
		peers:     vl1.NewPeerManager(log),
		netRepo:   repository.NewNetworkRepo(db),
		nodeRepo:  repository.NewNodeRepo(db),
		wsClients: make(map[string]net.Conn),
		log:       log.With("component", "controller"),
		ctx:       ctx,
		cancel:    cancel,
	}, nil
}

func (c *Controller) Start() error {
	transport, err := vl1.NewTransport(c.cfg.Controller.ListenPort, c.log)
	if err != nil {
		return err
	}
	c.transport = transport
	c.log.Info("controller started",
		"address", c.identity.Address,
		"port", c.transport.Port(),
	)

	c.wg.Add(1)
	go c.udpReadLoop()

	return nil
}

func (c *Controller) Stop() {
	c.cancel()
	if c.transport != nil {
		c.transport.Close()
	}
	c.wg.Wait()
}

func (c *Controller) udpReadLoop() {
	defer c.wg.Done()
	buf := make([]byte, vl1.MaxPacketSize)
	for {
		select {
		case <-c.ctx.Done():
			return
		default:
		}
		n, remoteAddr, err := c.transport.ReadFrom(buf)
		if err != nil {
			if c.ctx.Err() != nil {
				return
			}
			c.log.Error("UDP read error", "err", err)
			time.Sleep(time.Millisecond)
			continue
		}
		c.handlePacket(buf[:n], remoteAddr)
	}
}

func (c *Controller) handlePacket(data []byte, from *net.UDPAddr) {
	pkt, err := vl1.DecodePacket(data)
	if err != nil {
		c.log.Debug("decode packet", "err", err)
		return
	}
	switch pkt.Header.Type {
	case vl1.PacketTypeHandshake:
		c.handleHandshake(pkt.Payload, from)
	case vl1.PacketTypeKeepalive:
		if peer := c.peers.GetPeerByEndpoint(from); peer != nil {
			peer.Touch()
			nodeID := peer.Address.String()
			_ = c.nodeRepo.SetOnline(nodeID, true)
		}
	case vl1.PacketTypeData:
		// Controller does not forward data; data is P2P between agents
	}
}

func (c *Controller) handleHandshake(payload []byte, from *net.UDPAddr) {
	if len(payload) < 32 {
		return
	}
	var pubKey [32]byte
	copy(pubKey[:], payload[:32])
	addr := identity.AddressFromPublicKey(pubKey[:])
	nodeID := addr.String()

	peer := c.peers.GetPeer(addr)
	if peer == nil {
		peer = c.peers.AddPeer(addr, pubKey, from)
		c.log.Info("new peer connected", "node_id", nodeID, "addr", from)
	} else {
		c.peers.UpdatePeerEndpoint(addr, from)
		peer.Touch()
	}
	_ = c.HandleNodeHello(nodeID, fmt.Sprintf("%x", pubKey[:]), from)
}

func (c *Controller) Identity() *identity.Identity {
	return c.identity
}

func (c *Controller) HandleNodeHello(nodeID string, publicKey string, addr *net.UDPAddr) error {
	node, err := c.nodeRepo.FindByNodeID(nodeID)
	if err != nil {
		// New node: register it
		node = &model.Node{
			NodeID:    nodeID,
			PublicKey: publicKey,
			IPAddress: addr.IP.String(),
			Port:      addr.Port,
			Online:    true,
			Version:   "1.0.0",
		}
		if err := c.nodeRepo.Create(node); err != nil {
			return err
		}
		c.log.Info("new node registered", "node_id", nodeID, "addr", addr)
	} else {
		c.nodeRepo.SetOnline(nodeID, true)
		c.log.Debug("node hello", "node_id", nodeID, "addr", addr)
	}
	return nil
}

func (c *Controller) GetNodeList() ([]model.Node, error) {
	return c.nodeRepo.ListOnline()
}

func (c *Controller) GetNetworkConfig(networkID uint32) (*model.Network, error) {
	return c.netRepo.FindByNetworkID(networkID)
}

func (c *Controller) AuthorizeNode(networkID uint32, nodeID string) error {
	_, err := c.netRepo.FindByNetworkID(networkID)
	if err != nil {
		return err
	}
	_, err = c.nodeRepo.FindByNodeID(nodeID)
	if err != nil {
		return err
	}
	return c.netRepo.AddMember(&model.NetworkMember{
		NetworkID:  networkID,
		NodeID:     nodeID,
		Authorized: true,
	})
}