llmgw/internal/cli/firewallcmd.go

package cli

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/netip"
	"os"
	"strings"
	"time"

	"github.com/spf13/cobra"

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

// firewallStateFile is where a control plane writes what the cluster's
// settings say of the firewall, for llmgw-firewall-sync to apply as root:
// the server itself runs as llmgw and cannot change the ruleset.
const firewallStateFile = "/var/lib/llmgw/firewall.json"

// firewallState is what is written there: the same as the mirror is told
// and keeps beside its models.
type firewallState = mirror.FirewallState

// writeFirewallState records the settings' firewall for this machine's sync
// unit. Without the state directory — a server not run by systemd — there
// is no unit to read it, and nothing is written.
func writeFirewallState(cfg config.Server) error {
	if _, err := os.Stat("/var/lib/llmgw"); err != nil {
		return nil //nolint:nilerr // not a control plane set up by init
	}
	b, _ := json.Marshal(firewallState{Enabled: cfg.Firewall.Enabled, Networks: cfg.FirewallNetworks})
	if old, err := os.ReadFile(firewallStateFile); err == nil && string(old) == string(b) {
		return nil
	}
	tmp := firewallStateFile + ".new"
	if err := os.WriteFile(tmp, b, 0o644); err != nil {
		return err
	}
	return os.Rename(tmp, firewallStateFile)
}

// The units that apply it: a path unit watching the file, and the oneshot it
// starts.
const firewallSyncPathUnit = "llmgw-firewall-sync.path"

// firewallSyncUnits are the two units for a state file: on a control plane
// the one its server writes, on a mirror the one the control plane has the
// mirror write.
func firewallSyncUnits(state string) (service, pathUnit string) {
	arg := ""
	if state != firewallStateFile {
		arg = " --state " + state
	}
	service = `[Unit]
Description=llmgw firewall: apply the cluster's settings on this machine

[Service]
Type=oneshot
ExecStart=/usr/local/bin/llmgw firewall sync` + arg + `
`
	pathUnit = `[Unit]
Description=llmgw firewall: follow the cluster's settings on this machine

[Path]
PathChanged=` + state + `

[Install]
WantedBy=multi-user.target
`
	return service, pathUnit
}

func newFirewallCmd() *cobra.Command {
	c := &cobra.Command{
		Use:   "firewall",
		Short: "llmgw's firewall on this machine",
		Long: "Whether llmgw keeps a firewall is the cluster's setting firewall.enabled\n" +
			"(llmgw ctl config set firewall.enabled false). Nodes are told by the\n" +
			"control plane; a control plane writes it to " + firewallStateFile + ",\n" +
			"which " + firewallSyncPathUnit + " applies with 'llmgw firewall sync'.",
	}
	var install bool
	state := firewallStateFile
	sync := &cobra.Command{
		Use:   "sync",
		Short: "Apply what the cluster's settings say of the firewall, here (root)",
		Args:  cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			if os.Geteuid() != 0 {
				return errors.New("the firewall is changed as root: run it with sudo")
			}
			if install {
				if err := installFirewallSync(cmd.Context(), firewall.Local{}, state); err != nil {
					return err
				}
				fmt.Fprintln(cmd.OutOrStdout(), firewallSyncPathUnit+" follows the cluster's settings here")
				if _, err := os.Stat(state); err != nil {
					return nil // written when the settings are first told
				}
			}
			b, err := os.ReadFile(state)
			if err != nil {
				return err
			}
			var st firewallState
			if err := json.Unmarshal(b, &st); err != nil {
				return fmt.Errorf("%s: %w", state, err)
			}
			ctx, cancel := context.WithTimeout(cmd.Context(), 2*time.Minute)
			defer cancel()
			msg, err := syncFirewall(ctx, firewall.Local{}, st)
			if err != nil {
				return err
			}
			fmt.Fprintln(cmd.OutOrStdout(), msg)
			return nil
		},
	}
	sync.Flags().BoolVar(&install, "install", false, "first install the units that run this whenever the settings change")
	sync.Flags().StringVar(&state, "state", firewallStateFile,
		"the file the settings are written to: a control plane's, or "+mirror.FirewallFile+" in a mirror's directory")
	c.AddCommand(sync)
	return c
}

// installFirewallSync writes the sync units on a machine and has the path
// unit follow the state file.
func installFirewallSync(ctx context.Context, r firewall.Runner, state string) error {
	service, pathUnit := firewallSyncUnits(state)
	body := service + "\x00" + pathUnit
	out, err := r.Run(ctx, `set -e
t=$(mktemp -d); trap 'rm -rf "$t"' EXIT
cat > "$t/all"
python3 - "$t" <<'PY'
import sys
d = sys.argv[1]
a, b = open(d + "/all", "rb").read().split(b"\x00")
open(d + "/service", "wb").write(a)
open(d + "/path", "wb").write(b)
PY
sudo install -m 644 "$t/service" /etc/systemd/system/llmgw-firewall-sync.service
sudo install -m 644 "$t/path" /etc/systemd/system/`+firewallSyncPathUnit+`
sudo systemctl daemon-reload
sudo systemctl enable --now `+firewallSyncPathUnit+` >/dev/null 2>&1`, []byte(body))
	if err != nil {
		return fmt.Errorf("installing %s: %v: %s", firewallSyncPathUnit, err, string(out))
	}
	return nil
}

// syncFirewall puts a firewall state in force on this machine.
func syncFirewall(ctx context.Context, r firewall.Runner, st firewallState) (string, error) {
	if !st.Enabled {
		had, err := firewall.Disable(ctx, r)
		if err != nil {
			return "", err
		}
		if had {
			return "llmgw's firewall is off here", nil
		}
		return "llmgw's firewall was already off here", nil
	}
	loaded, err := firewall.Enable(ctx, r)
	if err != nil {
		return "", err
	}
	if !loaded {
		return "no llmgw rules on this machine to put back; 'llmgw server init' writes them", nil
	}
	var extra []netip.Prefix
	for _, n := range st.Networks {
		if p, err := netip.ParsePrefix(n); err == nil {
			extra = append(extra, p)
		} else if a, err := netip.ParseAddr(n); err == nil {
			extra = append(extra, netip.PrefixFrom(a, a.BitLen()))
		}
	}
	if len(extra) > 0 {
		if _, err := firewall.Install(ctx, r, firewall.Spec{Allow: extra}); err != nil {
			return "", err
		}
	}
	// The cluster's machines: a node on a network this one has not let in
	// yet, which the mirror's files are for.
	if len(st.Toward) > 0 {
		added, err := firewall.Admit(ctx, r, st.Toward)
		if err != nil {
			return "", err
		}
		if len(added) > 0 {
			shown := make([]string, len(added))
			for i, p := range added {
				shown[i] = p.String()
			}
			return "llmgw's firewall is on here, and now lets in " + strings.Join(shown, ", "), nil
		}
	}
	return "llmgw's firewall is on here", nil
}