llmgw/internal/cli/serveradd.go

package cli

import (
	"bytes"
	"context"
	"crypto/rand"
	"crypto/x509"
	"encoding/hex"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"time"

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/onboard"
	"gitea.mynodes.xyz/ilker/llmgw/internal/pgha"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
	"gitea.mynodes.xyz/ilker/llmgw/internal/version"
)

// serverAdd makes another machine a control plane of this cluster: a standby
// that takes over when this one is lost. Run on a control plane that holds
// the authority, as root, like `node add`.
type serverAdd struct {
	etcDir, pkiDir string
	// hosts are more names to put in the new control plane's certificate.
	hosts []string
	// vip is a floating address, "10.0.0.160/24", to keep on whichever of the
	// two leads; "" leaves client traffic to whatever is in front.
	vip string

	remote onboard.Remote
	// addr is the machine's address as given; ip what the others reach it at.
	addr, ip, name string
	out            io.Writer
}

func newServerAddCmd() *cobra.Command {
	sa := serverAdd{etcDir: "/etc/llmgw", pkiDir: defaultPKIDir}
	cmd := &cobra.Command{
		Use:   "add [user@]address",
		Short: "Make another machine a control plane of this cluster: a standby that takes over",
		Long: "Run on a control plane, as root. It logs in to the machine with ssh, as\n" +
			"node add does, and gives it this llmgw, a certificate from this\n" +
			"cluster's authority, this control plane's configuration and the store's\n" +
			"address, lets it into the store, and starts it. It stands by while this\n" +
			"one leads, and takes over within seconds when this one is lost; every\n" +
			"node is told of it and fails over to it without being touched.\n\n" +
			"The store must be one both can reach — replicated with 'server init\n" +
			"--patroni' — not a database on this machine.\n\n" +
			"--vip keeps a floating address on whichever control plane leads, with\n" +
			"keepalived on both: the address clients use for the gateway.\n\n" +
			"Run it again at any time: it changes only what is missing.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if os.Geteuid() != 0 {
				return errors.New("run it as root: it reads the certificate authority and the store's address")
			}
			user, host, ok := strings.Cut(args[0], "@")
			if !ok {
				user, host = "", args[0]
			}
			sa.remote = &onboard.SSH{Addr: host, User: user}
			sa.addr = host
			sa.out = cmd.OutOrStdout()
			return sa.run(cmd.Context())
		},
	}
	f := cmd.Flags()
	f.StringArrayVar(&sa.hosts, "host", nil, "another name for the new control plane's certificate; repeatable")
	f.StringVar(&sa.vip, "vip", "", "a floating address for whichever control plane leads, as 10.0.0.160/24")
	f.StringVar(&sa.pkiDir, "pki", sa.pkiDir, "directory holding the certificate authority")
	return cmd
}

func (sa *serverAdd) done(format string, a ...any) {
	fmt.Fprintf(sa.out, "  ok    "+format+"\n", a...)
}

func (sa *serverAdd) kept(format string, a ...any) {
	fmt.Fprintf(sa.out, "  kept  "+format+"\n", a...)
}

func (sa *serverAdd) sh(ctx context.Context, script string, stdin []byte) (string, error) {
	out, err := sa.remote.Run(ctx, script, stdin)
	if err != nil {
		return string(out), fmt.Errorf("on %s: %v: %s", sa.addr, err, strings.TrimSpace(string(out)))
	}
	return string(out), nil
}

// put writes a file there, through standard input, with its mode and owner.
func (sa *serverAdd) put(ctx context.Context, path, mode, owner string, body []byte) error {
	_, err := sa.sh(ctx, fmt.Sprintf(`set -e
sudo install -d -m 0755 %[1]s
sudo sh -c 'umask 077; cat > %[2]s.new'
sudo chmod %[3]s %[2]s.new
sudo chown %[4]s %[2]s.new
sudo mv -f %[2]s.new %[2]s`, filepath.Dir(path), path, mode, owner), body)
	return err
}

func (sa *serverAdd) run(ctx context.Context) error {
	steps := []struct {
		name string
		fn   func(context.Context) error
	}{
		{"connect", sa.connect},
		{"binary", sa.binary},
		{"directories", sa.dirs},
		{"certificate", sa.certificate},
		{"configuration", sa.configuration},
		{"store", sa.storeAccess},
		{"service", sa.service},
	}
	if sa.vip != "" {
		steps = append(steps, struct {
			name string
			fn   func(context.Context) error
		}{"floating address", sa.floating})
	}
	for _, s := range steps {
		if err := s.fn(ctx); err != nil {
			return fmt.Errorf("%s: %w", s.name, err)
		}
	}
	fmt.Fprintf(sa.out, "\n%s is a control plane of this cluster, standing by while this one leads.\n", sa.name)
	return nil
}

func (sa *serverAdd) connect(ctx context.Context) error {
	out, err := sa.sh(ctx, "hostname; sudo -n true && echo sudo-ok", nil)
	if err != nil {
		return fmt.Errorf("%w — it needs ssh and sudo without a password", err)
	}
	lines := strings.Split(strings.TrimSpace(out), "\n")
	sa.name = strings.TrimSpace(lines[0])
	if sa.ip == "" {
		if sa.ip, err = ipv4Of(sa.addr); err != nil {
			return err
		}
	}
	sa.done("%s (%s)", sa.name, sa.ip)
	return nil
}

// binary gives it this very llmgw, so the two control planes never differ.
func (sa *serverAdd) binary(ctx context.Context) error {
	have, _ := sa.sh(ctx, "llmgw --version 2>/dev/null || true", nil)
	if strings.TrimSpace(have) == "llmgw version "+version.Version {
		sa.kept("llmgw %s", version.Version)
		return nil
	}
	exe, err := os.Executable()
	if err != nil {
		return err
	}
	b, err := os.ReadFile(exe)
	if err != nil {
		return err
	}
	if _, err := sa.sh(ctx, `set -e; t=$(mktemp); cat > "$t"; sudo install -m 0755 "$t" /usr/local/bin/llmgw; rm -f "$t"`, b); err != nil {
		return err
	}
	sa.done("llmgw %s, copied from this machine", version.Version)
	return nil
}

func (sa *serverAdd) dirs(ctx context.Context) error {
	_, err := sa.sh(ctx, `set -e
id llmgw >/dev/null 2>&1 || sudo useradd --system --home /var/lib/llmgw --shell /usr/sbin/nologin llmgw
sudo install -d -m 0755 /etc/llmgw
sudo install -d -o llmgw -g llmgw -m 0700 /etc/llmgw/secrets
sudo install -d -o llmgw -g llmgw -m 0750 /var/lib/llmgw`, nil)
	if err == nil {
		sa.done("user llmgw, /etc/llmgw and /var/lib/llmgw")
	}
	return err
}

// certificate issues it a server certificate from this cluster's authority,
// for its name and addresses, unless it holds one that already names them.
func (sa *serverAdd) certificate(ctx context.Context) error {
	ca, err := loadCA(sa.pkiDir)
	if err != nil {
		return fmt.Errorf("the authority's key is not in %s: run this where it is (%v)", sa.pkiDir, err)
	}
	ips, _ := sa.sh(ctx, "hostname -I", nil)
	want := []string{sa.name, sa.ip}
	for _, ip := range strings.Fields(ips) {
		if net.ParseIP(ip).To4() != nil && !slices.Contains(want, ip) {
			want = append(want, ip)
		}
	}
	want = append(want, sa.hosts...)

	if have, _ := sa.sh(ctx, "sudo cat /etc/llmgw/server.pem 2>/dev/null || true", nil); have != "" {
		if blk, _ := pem.Decode([]byte(have)); blk != nil {
			if cert, err := x509.ParseCertificate(blk.Bytes); err == nil && cert.CheckSignatureFrom(ca.Cert) == nil {
				names := strings.Split(sanNames(cert), ", ")
				missing := false
				for _, w := range want {
					missing = missing || !slices.Contains(names, w)
				}
				if !missing {
					sa.kept("certificate for %s", strings.Join(names, ", "))
					return sa.put(ctx, "/etc/llmgw/ca.pem", "0644", "root:root", ca.PEM)
				}
			}
		}
	}
	leaf, err := ca.IssueServer(sa.name, want)
	if err != nil {
		return err
	}
	if err := sa.put(ctx, "/etc/llmgw/ca.pem", "0644", "root:root", ca.PEM); err != nil {
		return err
	}
	if err := sa.put(ctx, "/etc/llmgw/server.pem", "0644", "root:root", leaf.CertPEM); err != nil {
		return err
	}
	if err := sa.put(ctx, "/etc/llmgw/server-key.pem", "0640", "root:llmgw", leaf.KeyPEM); err != nil {
		return err
	}
	sa.done("certificate for %s, from this cluster's authority", strings.Join(want, ", "))
	return nil
}

// configuration copies this control plane's, and the secrets beside it: the
// two must agree on everything, the store above all.
func (sa *serverAdd) configuration(ctx context.Context) error {
	if have, _ := sa.sh(ctx, "sudo test -f /etc/llmgw/server.yaml && echo yes || true", nil); strings.TrimSpace(have) == "yes" {
		sa.kept("/etc/llmgw/server.yaml (never overwritten)")
	} else {
		b, err := os.ReadFile(filepath.Join(sa.etcDir, "server.yaml"))
		if err != nil {
			return err
		}
		if err := sa.put(ctx, "/etc/llmgw/server.yaml", "0644", "root:root", b); err != nil {
			return err
		}
		sa.done("/etc/llmgw/server.yaml, as this control plane's")
	}
	secrets, err := os.ReadDir(filepath.Join(sa.etcDir, "secrets"))
	if err != nil && !os.IsNotExist(err) {
		return err
	}
	for _, e := range secrets {
		if e.IsDir() || strings.Contains(e.Name(), ".pre-") || strings.HasSuffix(e.Name(), ".new") {
			continue
		}
		path := "/etc/llmgw/secrets/" + e.Name()
		mine, err := os.ReadFile(filepath.Join(sa.etcDir, "secrets", e.Name()))
		if err != nil {
			return err
		}
		theirs, _ := sa.sh(ctx, "sudo cat "+path+" 2>/dev/null || true", nil)
		switch {
		case theirs == string(mine):
			sa.kept("%s", path)
		case theirs != "":
			// Two control planes on two stores are two clusters.
			return fmt.Errorf("%s there differs from this one's; remove it there if that machine is to join this cluster", path)
		default:
			if err := sa.put(ctx, path, "0600", "llmgw:llmgw", mine); err != nil {
				return err
			}
			sa.done("%s", path)
		}
	}
	return nil
}

// storeAccess lets it into the store, and records both control planes for
// the nodes to be told.
func (sa *serverAdd) storeAccess(ctx context.Context) error {
	cfg, err := config.Load(filepath.Join(sa.etcDir, "server.yaml"))
	if err != nil {
		return err
	}
	addr, err := store.ResolveAddress(cfg.Auth.Store)
	if err != nil {
		return err
	}
	addr = strings.TrimSpace(addr)
	if !store.IsPostgres(addr) {
		return errors.New("the store is SQLite on this machine; a second control plane needs one both reach — run 'llmgw server init --patroni' first")
	}
	hosts, _, _, err := store.SplitPostgres(addr)
	if err != nil {
		return err
	}
	for _, h := range hosts {
		if ip := net.ParseIP(h); ip != nil && ip.IsLoopback() || h == "localhost" {
			return fmt.Errorf("the store is at %s, on this machine only; replicate it first with 'llmgw server init --patroni'", h)
		}
	}
	switch {
	case slices.Contains(hosts, sa.ip):
		sa.kept("store: %s is one of its members", sa.ip)
	case len(hosts) > 1:
		changed, err := pgha.AllowClient(ctx, hosts, sa.ip)
		if err != nil {
			return err
		}
		if changed {
			sa.done("store: %s let in, through Patroni", sa.ip)
		} else {
			sa.kept("store: %s already let in", sa.ip)
		}
	default:
		sa.kept("store at %s: make sure it lets %s in", hosts[0], sa.ip)
	}

	st, err := store.Open(ctx, addr)
	if err != nil {
		return err
	}
	defer st.Close()
	_, port, _ := net.SplitHostPort(cfg.AgentListen)
	if port == "" {
		port = "7443"
	}
	self, err := sourceIPToward(sa.ip)
	if err != nil {
		return err
	}
	me, _ := os.Hostname()
	for _, cp := range []store.ControlPlane{
		{URL: "https://" + net.JoinHostPort(self, port), Name: me},
		{URL: "https://" + net.JoinHostPort(sa.ip, port), Name: sa.name, AddedAt: time.Now().Add(time.Second)},
	} {
		if err := st.PutControlPlane(ctx, cp); err != nil {
			return err
		}
	}
	sa.done("nodes will be told of both control planes within half a minute")
	return nil
}

// service installs the unit and starts it, then waits for it to answer.
func (sa *serverAdd) service(ctx context.Context) error {
	if err := sa.put(ctx, "/etc/systemd/system/llmgw-server.service", "0644", "root:root", []byte(serverUnit)); err != nil {
		return err
	}
	if _, err := sa.sh(ctx, "sudo systemctl daemon-reload && sudo systemctl enable llmgw-server >/dev/null 2>&1 && sudo systemctl restart llmgw-server", nil); err != nil {
		return err
	}
	deadline := time.Now().Add(90 * time.Second)
	for {
		code, _ := sa.sh(ctx, "curl -s -o /dev/null -m 2 -w '%{http_code}' http://127.0.0.1:8080/leader || true", nil)
		switch strings.TrimSpace(code) {
		case "200":
			sa.done("llmgw-server running, and leading")
			return nil
		case "503":
			sa.done("llmgw-server running, standing by")
			return nil
		}
		if time.Now().After(deadline) {
			logs, _ := sa.sh(ctx, "sudo journalctl -u llmgw-server -n 15 --no-pager", nil)
			return fmt.Errorf("it did not answer within 90s:\n%s", logs)
		}
		time.Sleep(2 * time.Second)
	}
}

// floating keeps an address on whichever of the two leads, with keepalived
// on both, following /leader.
func (sa *serverAdd) floating(ctx context.Context) error {
	vip, nw, err := net.ParseCIDR(sa.vip)
	if err != nil || vip.To4() == nil {
		return fmt.Errorf("--vip %q is not an IPv4 address with its prefix, like 10.0.0.160/24", sa.vip)
	}
	prefix, _ := nw.Mask.Size()
	pass, err := vrrpPassword(sa.etcDir)
	if err != nil {
		return err
	}
	here := pgha.Local{}
	for _, m := range []struct {
		r        onboard.Remote
		name     string
		priority int
	}{{here, "this machine", 110}, {sa.remote, sa.name, 100}} {
		out, err := m.r.Run(ctx, fmt.Sprintf(`set -e
command -v keepalived >/dev/null 2>&1 || { sudo DEBIAN_FRONTEND=noninteractive apt-get -qq update >/dev/null; sudo DEBIAN_FRONTEND=noninteractive apt-get -qq install -y keepalived >/dev/null; }
ip -o route get %s | awk '{for(i=1;i<=NF;i++) if($i=="dev") {print $(i+1); exit}}'`, vip), nil)
		if err != nil {
			return fmt.Errorf("on %s: %v: %s", m.name, err, strings.TrimSpace(string(out)))
		}
		iface := strings.TrimSpace(string(out))
		conf := keepalivedConf(iface, vip.String(), prefix, m.priority, pass)
		if out, err := m.r.Run(ctx, `set -e
sudo install -d -m 0755 /etc/keepalived
sudo sh -c 'umask 077; cat > /etc/keepalived/keepalived.conf'
sudo systemctl enable keepalived >/dev/null 2>&1
sudo systemctl restart keepalived`, []byte(conf)); err != nil {
			return fmt.Errorf("on %s: %v: %s", m.name, err, strings.TrimSpace(string(out)))
		}
	}
	sa.done("%s follows whichever control plane leads (keepalived on both)", sa.vip)
	return nil
}

// vrrpPassword is the shared secret of the two keepalived, made once and
// kept with the other secrets, so running this again does not split them.
func vrrpPassword(etcDir string) (string, error) {
	p := filepath.Join(etcDir, "secrets", "vrrp")
	if b, err := os.ReadFile(p); err == nil && len(bytes.TrimSpace(b)) > 0 {
		return string(bytes.TrimSpace(b)), nil
	}
	b := make([]byte, 4)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	pass := hex.EncodeToString(b) // keepalived reads eight characters
	if err := os.WriteFile(p, []byte(pass+"\n"), 0o600); err != nil {
		return "", err
	}
	return pass, nil
}

func keepalivedConf(iface, vip string, prefix, priority int, pass string) string {
	return fmt.Sprintf(`# llmgw: the gateway's floating address, on whichever control plane leads.
vrrp_script llmgw_leader {
  script "/usr/bin/curl -fsS -o /dev/null -m 2 http://127.0.0.1:8080/leader"
  interval 2
  fall 2
  rise 2
}

vrrp_instance llmgw {
  state BACKUP
  interface %s
  virtual_router_id 51
  priority %d
  advert_int 1
  authentication {
    auth_type PASS
    auth_pass %s
  }
  virtual_ipaddress {
    %s/%d
  }
  track_script {
    llmgw_leader
  }
}
`, iface, priority, pass, vip, prefix)
}