llmgw/internal/agent/firewall.go

package agent

import (
	"context"
	"net"
	"net/netip"
	"net/url"
	"strings"
	"time"

	llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
	"gitea.mynodes.xyz/ilker/llmgw/internal/firewall"
	"gitea.mynodes.xyz/ilker/llmgw/internal/slurm"
)

// slurmPorts are Slurm's on a node that runs any of it: slurmctld's and
// slurmd's, and srun's range for the steps it starts.
var slurmPorts = []firewall.Port{
	{From: slurm.CtldPort, To: slurm.SlurmdPort, What: "slurm: slurmctld and slurmd"},
	{From: slurm.SrunPortsFrom, To: slurm.SrunPortsTo, What: "slurm: srun, for the steps it starts"},
}

// applyFirewall keeps llmgw's firewall on this node as the cluster's
// settings say: the engines' ports and the node exporter's for the networks
// the node reaches its control planes on, and those named, or no table at
// all. Off the receive loop, one at a time.
func (a *Agent) applyFirewall(m *llmgwv1.Firewall) {
	a.fwMu.Lock()
	defer a.fwMu.Unlock()
	a.fwLast = m
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
	defer cancel()
	r := a.fwRunner
	if r == nil {
		r = firewall.Local{}
	}
	if !m.GetEnabled() {
		had, err := firewall.Disable(ctx, r)
		switch {
		case err != nil:
			a.log.Warn("could not take llmgw's firewall off this node", "err", err)
		case had:
			a.log.Info("llmgw's firewall is off on this node, as the cluster's settings say")
		}
		return
	}
	a.mu.Lock()
	servers := append([]string(nil), a.serverList...)
	a.mu.Unlock()
	nets, err := firewall.Networks(ctx, r, controlPlaneIPs(ctx, servers))
	if err != nil {
		a.log.Warn("could not tell which networks reach the control planes; the firewall is left as it is", "err", err)
		return
	}
	for _, n := range m.GetNetworks() {
		if p, err := netip.ParsePrefix(n); err == nil {
			nets = append(nets, p)
		} else if ad, err := netip.ParseAddr(n); err == nil {
			nets = append(nets, netip.PrefixFrom(ad, ad.BitLen()))
		}
	}
	from, to := int(m.GetEnginePortsFrom()), int(m.GetEnginePortsTo())
	ports := []firewall.Port{firewall.NodeExporter}
	if from > 0 && to >= from {
		ports = append(ports, firewall.Engines(from, to))
	}
	// Once there, they stay: the ruleset keeps every port it held, so a
	// node that leaves Slurm keeps them closed to everyone but the cluster.
	if a.slurm != nil && a.slurm.HasPart() {
		ports = append(ports, slurmPorts...)
	}
	changed, err := firewall.Install(ctx, r, firewall.Spec{Ports: ports, Allow: nets})
	switch {
	case err != nil:
		a.log.Warn("could not put llmgw's firewall on this node", "err", err)
	case changed:
		a.log.Info("llmgw's firewall is on this node", "networks", len(nets))
	}
}

// slurmMessage is the leader's word on this node's part in Slurm. A node
// that gets one has the firewall applied again, with Slurm's ports.
func (a *Agent) slurmMessage(m *llmgwv1.Slurm) {
	had := a.slurm.HasPart()
	a.slurm.Apply(m)
	if had || !a.slurm.HasPart() {
		return
	}
	a.fwMu.Lock()
	last := a.fwLast
	a.fwMu.Unlock()
	if last != nil && last.GetEnabled() {
		go a.applyFirewall(last)
	}
}

// controlPlaneIPs are the addresses of the control planes an agent dials.
func controlPlaneIPs(ctx context.Context, servers []string) []string {
	var out []string
	for _, s := range servers {
		host := s
		if u, err := url.Parse(s); err == nil && u.Host != "" {
			host = u.Hostname()
		}
		host = strings.Trim(host, "[]")
		if a, err := netip.ParseAddr(host); err == nil {
			out = append(out, a.String())
			continue
		}
		ips, err := net.DefaultResolver.LookupNetIP(ctx, "ip", host)
		if err != nil {
			continue
		}
		for _, ip := range ips {
			out = append(out, ip.Unmap().String())
		}
	}
	return out
}