llmgw/internal/pgha/setup.go

package pgha

import (
	"bytes"
	"context"
	"fmt"
	"os/exec"
	"strings"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/onboard"
	"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
)

// Node is a machine the setup reaches, and how.
type Node struct {
	Member
	Remote onboard.Remote
}

// Setup builds, or finds, a Patroni cluster.
type Setup struct {
	// DB are the database machines; the first is where the cluster starts.
	DB []Node
	// Witness is the third etcd vote when there are only two database
	// machines: the control plane, run locally. Ignored with three or more.
	Witness *Node
	// Clients are the addresses control planes connect from.
	Clients []string
	// Status reads the cluster from Patroni's API; nil asks over HTTP.
	Status func(ctx context.Context, hosts []string) (Cluster, error)
	// Emit reports progress: a step, what happened, and whether anything
	// was changed or found as it should be.
	Emit func(step, detail string, changed bool)
	// Wait bounds each wait for the cluster to come up.
	Wait time.Duration
	// CA signs each member's certificate; nil signs none and turns nothing
	// on. CAFile is where this control plane finds the authority's
	// certificate, for the DSN: used when TLS is on, whoever turned it on.
	CA     *tlsconf.CA
	CAFile string

	// RestPassword is what Patroni's API asks for before a change, once Run
	// has found or made it: for `server add`, which changes the access rules.
	RestPassword string
}

// Local is a Remote for the machine this runs on, which init runs as root.
type Local struct{}

// Run runs script with sh here.
func (Local) Run(ctx context.Context, script string, stdin []byte) ([]byte, error) {
	cmd := exec.CommandContext(ctx, "sh", "-c", script)
	if stdin != nil {
		cmd.Stdin = bytes.NewReader(stdin)
	}
	var out, errb bytes.Buffer
	cmd.Stdout, cmd.Stderr = &out, &errb
	if err := cmd.Run(); err != nil {
		return append(out.Bytes(), errb.Bytes()...), err
	}
	return out.Bytes(), nil
}

func (s *Setup) emit(step string, changed bool, format string, a ...any) {
	if s.Emit != nil {
		s.Emit(step, fmt.Sprintf(format, a...), changed)
	}
}

// etcdMembers are the machines that vote: the database machines, and the
// witness when they are only two.
func (s *Setup) etcdMembers() []Node {
	out := append([]Node{}, s.DB...)
	if len(s.DB) == 2 && s.Witness != nil {
		out = append(out, *s.Witness)
	}
	return out
}

// Run sets the cluster up, or confirms it is, and returns the store address.
func (s *Setup) Run(ctx context.Context) (string, error) {
	if len(s.DB) < 2 {
		return "", fmt.Errorf("a replicated store needs at least two database machines, not %d", len(s.DB))
	}
	if len(s.DB) == 2 && s.Witness == nil {
		return "", fmt.Errorf("two database machines need a third etcd vote, or losing either stops failover")
	}
	if s.Wait == 0 {
		s.Wait = 3 * time.Minute
	}
	if s.Status == nil {
		s.Status = HTTPStatus
	}
	// Members' APIs are verified against the cluster's authority once they
	// speak TLS.
	switch {
	case s.CA != nil:
		if err := setRESTCAPEM(s.CA.PEM); err != nil {
			return "", err
		}
	case s.CAFile != "":
		_ = SetRESTCA(s.CAFile)
	}
	sec, err := s.secrets(ctx)
	if err != nil {
		return "", fmt.Errorf("passwords: %w", err)
	}
	s.RestPassword = sec.RestAPI
	if err := s.etcd(ctx, sec); err != nil {
		return "", fmt.Errorf("etcd: %w", err)
	}
	if err := s.certs(ctx); err != nil {
		return "", fmt.Errorf("certificates: %w", err)
	}
	if err := s.patroni(ctx, sec); err != nil {
		return "", fmt.Errorf("patroni: %w", err)
	}
	// What a cluster set up before them lacks, applied in place; a new one
	// has them already and these find nothing to change.
	if err := s.restAuth(ctx, sec); err != nil {
		return "", fmt.Errorf("patroni api: %w", err)
	}
	if err := s.tlsInShared(ctx, sec); err != nil {
		return "", fmt.Errorf("tls: %w", err)
	}
	if err := s.ensurePatroniTLS(ctx, sec); err != nil {
		return "", fmt.Errorf("patroni api: %w", err)
	}
	if err := s.database(ctx); err != nil {
		return "", fmt.Errorf("database: %w", err)
	}
	members := make([]Member, len(s.DB))
	for i, d := range s.DB {
		members[i] = d.Member
	}
	// TLS when this setup turned it on, or found it on: then the member
	// reached is verified against the authority's certificate, which every
	// control plane has even where its key is not.
	caFile := ""
	if s.CAFile != "" && (s.CA != nil || s.sharedTLS(ctx)) {
		caFile = s.CAFile
	}
	return DSN(members, sec.App, caFile), nil
}

// secretDir holds the passwords on every database machine, readable by root
// alone. Generated on the first, carried to the others.
const secretDir = "/root/llmgw-pg"

var secretNames = []string{"superuser", "replicator", "llmgw", "restapi"}

func (s *Setup) secrets(ctx context.Context) (Secrets, error) {
	script := `set -e
sudo install -d -m 700 ` + secretDir + `
for f in superuser replicator llmgw restapi; do
  sudo test -s ` + secretDir + `/$f.pw || sudo sh -c "umask 077; openssl rand -hex 24 > ` + secretDir + `/$f.pw"
done
for f in superuser replicator llmgw restapi; do printf '%s=' $f; sudo cat ` + secretDir + `/$f.pw; done`
	out, err := s.DB[0].Remote.Run(ctx, script, nil)
	if err != nil {
		return Secrets{}, fmt.Errorf("on %s: %v: %s", s.DB[0].Name, err, strings.TrimSpace(string(out)))
	}
	got := parseSecrets(out)
	if got.Superuser == "" || got.Replicator == "" || got.App == "" || got.RestAPI == "" {
		return Secrets{}, fmt.Errorf("could not read them back from %s", s.DB[0].Name)
	}
	for _, d := range s.DB[1:] {
		have, err := d.Remote.Run(ctx, `for f in superuser replicator llmgw restapi; do
  if sudo test -s `+secretDir+`/$f.pw; then printf '%s=' $f; sudo cat `+secretDir+`/$f.pw; fi
done`, nil)
		if err != nil {
			return Secrets{}, fmt.Errorf("on %s: %v", d.Name, err)
		}
		h := parseSecrets(have)
		if h.Superuser != "" || h.Replicator != "" || h.App != "" {
			if h.Superuser != got.Superuser || h.Replicator != got.Replicator || h.App != got.App {
				// Not overwritten: a member holding other passwords belongs
				// to another cluster, or to an earlier one somebody meant to
				// keep, and replacing them would lock it out of itself.
				return Secrets{}, fmt.Errorf("%s holds passwords in %s that differ from %s's; "+
					"remove them there if that machine is meant to join this cluster", d.Name, secretDir, s.DB[0].Name)
			}
			if h.RestAPI == got.RestAPI {
				continue
			}
			// A cluster from before the API asked for a password: the
			// member is given the one the first holds.
		}
		body := fmt.Sprintf("%s\n%s\n%s\n%s\n", got.Superuser, got.Replicator, got.App, got.RestAPI)
		if out, err := d.Remote.Run(ctx, `set -e
sudo install -d -m 700 `+secretDir+`
read su; read rep; read app; read api
for kv in "superuser $su" "replicator $rep" "llmgw $app" "restapi $api"; do
  set -- $kv
  printf '%s\n' "$2" | sudo sh -c "umask 077; cat > `+secretDir+`/$1.pw"
done`, []byte(body)); err != nil {
			return Secrets{}, fmt.Errorf("on %s: %v: %s", d.Name, err, strings.TrimSpace(string(out)))
		}
		s.emit("passwords", true, "carried to %s", d.Name)
	}
	s.emit("passwords", false, "kept in %s on every database machine", secretDir)
	return got, nil
}

func parseSecrets(out []byte) Secrets {
	var s Secrets
	for _, line := range strings.Split(string(out), "\n") {
		k, v, ok := strings.Cut(strings.TrimSpace(line), "=")
		if !ok {
			continue
		}
		switch k {
		case "superuser":
			s.Superuser = v
		case "replicator":
			s.Replicator = v
		case "llmgw":
			s.App = v
		case "restapi":
			s.RestAPI = v
		}
	}
	return s
}

// etcd brings up the voting members, unless every one of them is already
// healthy in one cluster — in which case nothing is touched, since rewriting
// a member's configuration means emptying its data. With the cluster's
// authority here, members speak TLS; a running cluster that does not yet is
// moved to it (etcdToTLS).
func (s *Setup) etcd(ctx context.Context, sec Secrets) error {
	nodes := s.etcdMembers()
	tls := s.CA != nil
	if s.etcdHealthy(ctx, nodes[0], etcdEndpoints(nodes, tls), len(nodes), tls) {
		s.emit("etcd", false, "%d members healthy%s; left as they are", len(nodes), map[bool]string{true: ", over TLS"}[tls])
		return nil
	}
	if tls && s.etcdHealthy(ctx, nodes[0], etcdEndpoints(nodes, false), len(nodes), false) {
		return s.etcdToTLS(ctx, sec)
	}
	return s.etcdBootstrap(ctx, tls)
}

func etcdEndpoints(nodes []Node, tls bool) []string {
	scheme := "http"
	if tls {
		scheme = "https"
	}
	out := make([]string, len(nodes))
	for i, n := range nodes {
		out[i] = scheme + "://" + n.IP + ":2379"
	}
	return out
}

// etcdInstall puts etcd on every member, stopped if it was not running.
func (s *Setup) etcdInstall(ctx context.Context, nodes []Node) error {
	for _, n := range nodes {
		out, err := n.Remote.Run(ctx, `set -e
if ! command -v etcd >/dev/null 2>&1; then
  sudo DEBIAN_FRONTEND=noninteractive apt-get -qq update >/dev/null
  sudo DEBIAN_FRONTEND=noninteractive apt-get -qq install -y etcd-server etcd-client >/dev/null
fi`, nil)
		if err != nil {
			return fmt.Errorf("on %s: %v: %s", n.Name, err, strings.TrimSpace(string(out)))
		}
	}
	return nil
}

// etcdCerts gives each etcd member a certificate from the cluster's
// authority, for its address, its name and loopback, unless it has one.
func (s *Setup) etcdCerts(ctx context.Context, nodes []Node) error {
	for _, n := range nodes {
		have, _ := n.Remote.Run(ctx, "sudo openssl x509 -in "+EtcdCert+" -noout -ext subjectAltName,extendedKeyUsage 2>/dev/null", nil)
		if strings.Contains(string(have), "IP Address:"+n.IP) && strings.Contains(string(have), "Client Authentication") {
			continue
		}
		leaf, err := s.CA.IssuePeer(n.Name, []string{n.IP, n.Name, "127.0.0.1", "localhost"})
		if err != nil {
			return err
		}
		if err := putTLS(ctx, n, "etcd", EtcdTLSDir, EtcdCA, EtcdCert, EtcdKey, s.CA.PEM, leaf.CertPEM, leaf.KeyPEM); err != nil {
			return err
		}
		s.emit("etcd", true, "certificate for %s (%s)", n.Name, n.IP)
	}
	return nil
}

// putTLS installs an authority's certificate, a leaf and its key on a
// machine, owned by user, the key readable by it alone.
func putTLS(ctx context.Context, n Node, user, dir, caPath, certPath, keyPath string, ca, cert, key []byte) error {
	var body bytes.Buffer
	body.Write(ca)
	body.WriteString("\x00")
	body.Write(cert)
	body.WriteString("\x00")
	body.Write(key)
	out, err := n.Remote.Run(ctx, `set -e
sudo install -d -o `+user+` -g `+user+` -m 750 `+dir+`
t=$(mktemp -d); trap 'rm -rf "$t"' EXIT
cat > "$t/all"
python3 - "$t" <<'PY'
import sys
d = sys.argv[1]
ca, crt, key = open(d + "/all", "rb").read().split(b"\x00")
for name, data in (("ca", ca), ("crt", crt), ("key", key)):
    open(d + "/" + name, "wb").write(data)
PY
sudo install -o `+user+` -g `+user+` -m 644 "$t/ca" `+caPath+`
sudo install -o `+user+` -g `+user+` -m 644 "$t/crt" `+certPath+`
sudo install -o `+user+` -g `+user+` -m 600 "$t/key" `+keyPath, body.Bytes())
	if err != nil {
		return fmt.Errorf("on %s: %v: %s", n.Name, err, strings.TrimSpace(string(out)))
	}
	return nil
}

// etcdBootstrap writes every member's configuration and starts them
// together, emptying a member whose configuration changed.
func (s *Setup) etcdBootstrap(ctx context.Context, tls bool) error {
	nodes := s.etcdMembers()
	var members []Member
	for _, n := range nodes {
		members = append(members, n.Member)
	}
	if err := s.etcdInstall(ctx, nodes); err != nil {
		return err
	}
	if tls {
		if err := s.etcdCerts(ctx, nodes); err != nil {
			return err
		}
	}
	changed := make([]bool, len(nodes))
	for i, n := range nodes {
		out, err := n.Remote.Run(ctx, `set -e
new=$(cat)
if [ "$(sudo cat /etc/default/etcd 2>/dev/null)" != "$new" ]; then
  sudo systemctl stop etcd >/dev/null 2>&1 || true
  sudo rm -rf /var/lib/etcd/default
  printf '%s\n' "$new" | sudo tee /etc/default/etcd >/dev/null
  echo changed
fi
sudo systemctl enable etcd >/dev/null 2>&1`, []byte(Etcd(n.Member, members, tls)))
		if err != nil {
			return fmt.Errorf("on %s: %v: %s", n.Name, err, strings.TrimSpace(string(out)))
		}
		changed[i] = strings.Contains(string(out), "changed")
	}
	// Started together: a new member waits for the others before it says it
	// is up, so starting them one after another would wait for ever on the
	// first.
	var wg sync.WaitGroup
	errs := make([]error, len(nodes))
	for i, n := range nodes {
		wg.Add(1)
		go func(i int, n Node) {
			defer wg.Done()
			verb := "start"
			if changed[i] {
				verb = "restart"
			}
			if out, err := n.Remote.Run(ctx, "sudo systemctl "+verb+" --no-block etcd", nil); err != nil {
				errs[i] = fmt.Errorf("on %s: %v: %s", n.Name, err, strings.TrimSpace(string(out)))
			}
		}(i, n)
	}
	wg.Wait()
	for _, err := range errs {
		if err != nil {
			return err
		}
	}
	deadline := time.Now().Add(s.Wait)
	for !s.etcdHealthy(ctx, nodes[0], etcdEndpoints(nodes, tls), len(nodes), tls) {
		if time.Now().After(deadline) {
			return fmt.Errorf("the %d members did not become healthy within %s; `journalctl -u etcd` on each says why", len(nodes), s.Wait)
		}
		if err := sleep(ctx, 2*time.Second); err != nil {
			return err
		}
	}
	s.emit("etcd", true, "%d members up%s", len(nodes), map[bool]string{true: ", over TLS"}[tls])
	return nil
}

func (s *Setup) etcdHealthy(ctx context.Context, n Node, endpoints []string, want int, tls bool) bool {
	out, err := n.Remote.Run(ctx, etcdctl(tls, endpoints)+" endpoint health 2>&1", nil)
	return err == nil && strings.Count(string(out), "is healthy") == want
}

// ensurePatroniTLS puts a member whose Patroni does not yet reach etcd or
// serve its API over TLS on it, in maintenance mode: an etcd moved to TLS by
// a run that stopped before it got this far.
func (s *Setup) ensurePatroniTLS(ctx context.Context, sec Secrets) error {
	if s.CA == nil {
		return nil
	}
	need := false
	for _, d := range s.DB {
		out, err := d.Remote.Run(ctx, patroniTLSEdit("check"), nil)
		if err != nil {
			return fmt.Errorf("on %s: %v: %s", d.Name, err, strings.TrimSpace(string(out)))
		}
		need = need || strings.TrimSpace(lastLine(out)) == "changed"
	}
	if !need {
		return nil
	}
	if err := s.pause(ctx, sec, true); err != nil {
		return err
	}
	if _, err := s.patroniOnTLS(ctx, true); err != nil {
		return s.stuck(err)
	}
	if err := s.waitFor(ctx, s.dbHosts(), func(c Cluster) bool { return c.Leader() != "" }); err != nil {
		return s.stuck(err)
	}
	if err := s.pause(ctx, sec, false); err != nil {
		return s.stuck(err)
	}
	return s.waitFor(ctx, s.dbHosts(), func(c Cluster) bool { return c.Complete(len(s.DB)) })
}

// patroni installs PostgreSQL and Patroni on every database machine, unless
// the cluster already has a leader and every member streaming.
func (s *Setup) patroni(ctx context.Context, sec Secrets) error {
	hosts := make([]string, len(s.DB))
	for i, d := range s.DB {
		hosts[i] = d.IP
	}
	if c, err := s.Status(ctx, hosts); err == nil && c.Complete(len(s.DB)) {
		s.emit("patroni", false, "%s leads, %d streaming; left as it is", c.Leader(), len(s.DB)-1)
		return nil
	}

	var version string
	for _, d := range s.DB {
		v, err := s.preparePostgres(ctx, d)
		if err != nil {
			return err
		}
		if version != "" && v != version {
			// Streaming replication is between one major version.
			return fmt.Errorf("%s has PostgreSQL %s and %s has %s; they must match", s.DB[0].Name, version, d.Name, v)
		}
		version = v
	}
	// Now that PostgreSQL's user exists, its certificates.
	if err := s.certs(ctx); err != nil {
		return err
	}

	var etcd, db []Member
	for _, n := range s.etcdMembers() {
		etcd = append(etcd, n.Member)
	}
	for _, d := range s.DB {
		db = append(db, d.Member)
	}
	for i, d := range s.DB {
		cfg := Patroni(d.Member, db, etcd, s.Clients, version, sec, s.CA != nil)
		if out, err := d.Remote.Run(ctx, `set -e
sudo install -d -m 755 /etc/patroni
sudo install -d -o postgres -g postgres -m 700 /var/lib/postgresql/`+version+`/llmgw
sudo sh -c 'umask 027; cat > /etc/patroni/config.yml.new'
sudo chown root:postgres /etc/patroni/config.yml.new
sudo chmod 640 /etc/patroni/config.yml.new
sudo mv -f /etc/patroni/config.yml.new /etc/patroni/config.yml
sudo systemctl enable patroni >/dev/null 2>&1
sudo systemctl restart patroni`, []byte(cfg)); err != nil {
			return fmt.Errorf("on %s: %v: %s", d.Name, err, strings.TrimSpace(string(out)))
		}
		// The first starts the cluster; the rest clone it, which they can
		// only do once it has a leader.
		want := func(c Cluster) bool { return c.Leader() != "" }
		if i > 0 {
			want = func(c Cluster) bool { return c.Complete(i + 1) }
		}
		if err := s.waitFor(ctx, hosts, want); err != nil {
			return fmt.Errorf("after starting %s: %w", d.Name, err)
		}
		s.emit("patroni", true, "%s up", d.Name)
	}
	return nil
}

// preparePostgres installs the packages and clears the distribution's own
// cluster, which would hold the port Patroni's needs — unless it holds a
// database somebody made, which is not this setup's to delete.
func (s *Setup) preparePostgres(ctx context.Context, d Node) (string, error) {
	out, err := d.Remote.Run(ctx, `set -e
need=""
command -v patroni >/dev/null 2>&1 || need="$need patroni"
ls /usr/lib/postgresql/*/bin/postgres >/dev/null 2>&1 || need="$need postgresql"
if [ -n "$need" ]; then
  sudo DEBIAN_FRONTEND=noninteractive apt-get -qq update >/dev/null
  sudo DEBIAN_FRONTEND=noninteractive apt-get -qq install -y $need >/dev/null
fi
v=$(ls /usr/lib/postgresql | sort -n | tail -1)
if [ -d /var/lib/postgresql/$v/main ]; then
  if sudo systemctl is-active --quiet postgresql@$v-main; then
    n=$(sudo -u postgres psql -Atc "select count(*) from pg_database where datname not in ('postgres','template0','template1')")
    if [ "$n" != "0" ]; then echo "holds-data"; exit 0; fi
  fi
  sudo systemctl disable --now postgresql@$v-main postgresql >/dev/null 2>&1 || true
  sudo pg_dropcluster --stop $v main >/dev/null 2>&1 || sudo rm -rf /var/lib/postgresql/$v/main
fi
sudo systemctl disable --now postgresql >/dev/null 2>&1 || true
echo "version=$v"`, nil)
	if err != nil {
		return "", fmt.Errorf("on %s: %v: %s", d.Name, err, strings.TrimSpace(string(out)))
	}
	if strings.Contains(string(out), "holds-data") {
		return "", fmt.Errorf("%s runs a PostgreSQL with databases in it; move or drop them first — "+
			"this setup replaces that cluster", d.Name)
	}
	for _, line := range strings.Split(string(out), "\n") {
		if v, ok := strings.CutPrefix(strings.TrimSpace(line), "version="); ok && v != "" {
			return v, nil
		}
	}
	return "", fmt.Errorf("no PostgreSQL found on %s after installing it", d.Name)
}

// database makes the role and database on whichever member leads.
func (s *Setup) database(ctx context.Context) error {
	hosts := make([]string, len(s.DB))
	for i, d := range s.DB {
		hosts[i] = d.IP
	}
	c, err := s.Status(ctx, hosts)
	if err != nil {
		return err
	}
	var leader *Node
	for i := range s.DB {
		if s.DB[i].Name == c.Leader() {
			leader = &s.DB[i]
		}
	}
	if leader == nil {
		return fmt.Errorf("no member leads")
	}
	out, err := leader.Remote.Run(ctx, `set -e
pw=$(sudo cat `+secretDir+`/llmgw.pw)
sudo -u postgres psql -Atc "select 'had=' || (select count(*) from pg_roles where rolname = 'llmgw') + (select count(*) from pg_database where datname = 'llmgw')"
sudo -u postgres psql -q -v ON_ERROR_STOP=1 -v pw="$pw" <<'SQL'
SELECT format('CREATE ROLE %I LOGIN', 'llmgw') WHERE NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'llmgw') \gexec
ALTER ROLE llmgw PASSWORD :'pw';
SELECT format('CREATE DATABASE %I OWNER %I', 'llmgw', 'llmgw') WHERE NOT EXISTS (SELECT 1 FROM pg_database WHERE datname = 'llmgw') \gexec
SQL`, nil)
	if err != nil {
		return fmt.Errorf("on %s: %v: %s", leader.Name, err, strings.TrimSpace(string(out)))
	}
	// Said as a change only when something was made: a rerun that finds
	// both is not a reason to restart the gateway.
	s.emit("database", !strings.Contains(string(out), "had=2"), "role and database %s on %s", Database, leader.Name)
	return nil
}

func (s *Setup) waitFor(ctx context.Context, hosts []string, ok func(Cluster) bool) error {
	deadline := time.Now().Add(s.Wait)
	var last error
	for {
		c, err := s.Status(ctx, hosts)
		if err == nil && ok(c) {
			return nil
		}
		if err != nil {
			last = err
		} else {
			last = fmt.Errorf("the cluster is %s", c)
		}
		if time.Now().After(deadline) {
			return fmt.Errorf("not up within %s: %v; `journalctl -u patroni` says why", s.Wait, last)
		}
		if err := sleep(ctx, 3*time.Second); err != nil {
			return err
		}
	}
}

func sleep(ctx context.Context, d time.Duration) error {
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-time.After(d):
		return nil
	}
}