llmgw/internal/cli/firewallpub.go

package cli

import (
	"context"
	"encoding/json"
	"log/slog"
	"net"
	"net/netip"
	"net/url"
	"slices"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// firewallStateKey is where, in the store, the firewall's state is kept
// for the store's members to read.
const firewallStateKey = "firewall"

// firewallPublisher tells the machines that cannot ask a control plane
// what the cluster's settings say of their firewall: the mirror, over its
// API, and the store's members, through the store they hold a copy of.
// Nodes are told over their agents' streams and control planes by their
// own servers. On the leader only, at start, at every change to the
// firewall's settings, and every few minutes for nodes come and gone.
type firewallPublisher struct {
	settings func() config.Server
	st       *store.DB
	nodes    func() []string
	leader   func() bool
	log      *slog.Logger

	mu     sync.Mutex
	mirror *mirror.Client
	kick   chan struct{}
}

func newFirewallPublisher(settings func() config.Server, st *store.DB, nodes func() []string,
	leader func() bool, log *slog.Logger) *firewallPublisher {
	return &firewallPublisher{settings: settings, st: st, nodes: nodes, leader: leader, log: log,
		kick: make(chan struct{}, 1)}
}

// setMirror has the mirror told too.
func (p *firewallPublisher) setMirror(m *mirror.Client) {
	p.mu.Lock()
	p.mirror = m
	p.mu.Unlock()
	p.Kick()
}

// Kick has the state worked out and told again now.
func (p *firewallPublisher) Kick() {
	select {
	case p.kick <- struct{}{}:
	default:
	}
}

func (p *firewallPublisher) run(ctx context.Context, every time.Duration) {
	t := time.NewTicker(every)
	defer t.Stop()
	for {
		p.once(ctx)
		select {
		case <-ctx.Done():
			return
		case <-t.C:
		case <-p.kick:
		}
	}
}

// state is what the settings say now, with the cluster's machines.
func (p *firewallPublisher) state(ctx context.Context) mirror.FirewallState {
	cfg := p.settings()
	st := mirror.FirewallState{Enabled: cfg.Firewall.Enabled, Networks: slices.Clone(cfg.FirewallNetworks)}
	var toward []string
	if p.nodes != nil {
		toward = append(toward, p.nodes()...)
	}
	if p.st != nil {
		if cps, err := p.st.ControlPlanes(ctx); err == nil {
			for _, c := range cps {
				toward = append(toward, c.URL)
			}
		}
	}
	for _, t := range toward {
		if a, ok := hostAddr(t); ok {
			st.Toward = append(st.Toward, a)
		}
	}
	slices.Sort(st.Toward)
	st.Toward = slices.Compact(st.Toward)
	return st
}

// hostAddr is the address in an address, host:port or URL; names are not
// looked up here.
func hostAddr(s string) (string, bool) {
	if u, err := url.Parse(s); err == nil && u.Host != "" {
		s = u.Hostname()
	} else if h, _, err := net.SplitHostPort(s); err == nil {
		s = h
	}
	a, err := netip.ParseAddr(s)
	if err != nil {
		return "", false
	}
	return a.Unmap().String(), true
}

func (p *firewallPublisher) once(ctx context.Context) {
	if p.leader != nil && !p.leader() {
		return
	}
	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()
	st := p.state(ctx)
	p.mu.Lock()
	m := p.mirror
	p.mu.Unlock()
	if m != nil {
		changed, err := m.SetFirewall(ctx, st)
		switch {
		case err != nil:
			p.log.Warn("could not tell the mirror what the settings say of its firewall", "err", err)
		case changed:
			p.log.Info("the mirror's firewall follows the cluster's settings", "enabled", st.Enabled)
		}
	}
	if p.st != nil {
		b, _ := json.Marshal(st)
		changed, err := p.st.PutClusterState(ctx, firewallStateKey, string(b))
		switch {
		case err != nil:
			p.log.Warn("could not record the firewall's state for the store's members", "err", err)
		case changed:
			p.log.Info("the store's members' firewall follows the cluster's settings", "enabled", st.Enabled)
		}
	}
}