llmgw/internal/cli/serverinit.go

package cli

import (
	"context"
	"crypto/rand"
	"encoding/base32"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"runtime"
	"slices"
	"strconv"
	"strings"
	"time"

	"github.com/spf13/cobra"

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

// serverUnit is deploy/systemd/llmgw-server.service. A test holds the two
// together, because a machine set up by init and one set up from the docs
// should not differ in a way nobody chose.
const serverUnit = `[Unit]
Description=llmgw control plane
After=network-online.target
Wants=network-online.target

[Service]
ExecStart=/usr/local/bin/llmgw server --config /etc/llmgw/server.yaml
User=llmgw
Group=llmgw
Restart=on-failure
RestartSec=2s
NoNewPrivileges=true
ProtectSystem=strict
ProtectHome=true
PrivateTmp=true
StateDirectory=llmgw

[Install]
WantedBy=multi-user.target
`

// serverInit is everything a control plane needs before it can start, done in
// the order it has to be done in. Every step looks before it acts, so running
// it again finishes an install that stopped halfway and changes nothing on
// one that did not.
type serverInit struct {
	etcDir   string // /etc/llmgw
	stateDir string // /var/lib/llmgw
	pkiDir   string // where the authority lives
	unitPath string // "" leaves systemd alone

	// user owns the state and reads the key; "" means no user is created and
	// nothing is chowned, which is what a test wants.
	user  string
	hosts []string
	admin string
	// postgres keeps the store in a PostgreSQL on this machine, installed if
	// need be, moving whatever SQLite already holds.
	postgres bool
	// patroni are user@host of the machines to keep the store on under
	// Patroni, replicated; this machine is the third etcd vote for two.
	patroni []string
	// noMirror, noNode and noModel turn off the parts of a first install that
	// make it serve something: a mirror, this machine as a node, and a small
	// model on it.
	noMirror, noNode, noModel bool
	// mirrorURL is a mirror on another machine to use instead of running one.
	mirrorURL string
	// pgSetup prepares PostgreSQL and returns the store address; a seam for
	// tests. Nil does it for real.
	pgSetup func(context.Context) (string, error)
	pg      pgServer

	out io.Writer

	uid, gid int
	changed  bool
}

func newServerInitCmd() *cobra.Command {
	in := serverInit{
		etcDir:   "/etc/llmgw",
		stateDir: "/var/lib/llmgw",
		pkiDir:   defaultPKIDir,
		unitPath: "/etc/systemd/system/llmgw-server.service",
		user:     "llmgw",
	}
	var plain bool
	cmd := &cobra.Command{
		Use:   "init",
		Short: "Set this machine up as a control plane and start it",
		Long: "Creates the llmgw user and its directories, a certificate authority, the\n" +
			"control plane's certificate for every address this machine has, a\n" +
			"configuration, an administrator account, and the systemd unit — then\n" +
			"starts it and waits until it answers.\n\n" +
			"Then it makes the machine serve something: a model mirror beside the\n" +
			"control plane, this machine as a node of it (on its GPUs when nvidia-smi\n" +
			"finds any, on CPU otherwise), and " + starterName + " fetched into the\n" +
			"mirror and served on this node. --no-mirror, --no-node and --no-model\n" +
			"leave those out.\n\n" +
			"With --postgres the store is a PostgreSQL database on this machine:\n" +
			"installed if there is none, with its own role and a generated password\n" +
			"kept in /etc/llmgw/secrets. An existing SQLite store is copied into it,\n" +
			"so accounts and keys survive, and server.yaml is pointed at it.\n\n" +
			"With --patroni, given two or more machines, the store is PostgreSQL\n" +
			"replicated across them under Patroni, which promotes a standby when the\n" +
			"primary is lost. etcd holds its leader lock on those machines — and on\n" +
			"this one, when they are two, as the third vote. Installed over ssh with\n" +
			"passwordless sudo, as 'llmgw node add' does; the store already in use,\n" +
			"SQLite or PostgreSQL, is moved in. Run again, it changes nothing.\n\n" +
			"Each step is skipped when its result already exists, so it is safe to run\n" +
			"again: to finish an install that stopped, or with another --host to\n" +
			"reissue the certificate. It never overwrites server.yaml or the authority.\n\n" +
			"The authority's key stays in " + defaultPKIDir + " on this machine. Move it\n" +
			"somewhere safer once the nodes are onboarded if this machine is exposed.",
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			if runtime.GOOS != "linux" {
				return fmt.Errorf("init sets up a Linux machine; this is %s", runtime.GOOS)
			}
			if os.Geteuid() != 0 {
				return errors.New("init creates a user, writes /etc and starts a service: run it as root")
			}
			in.out = cmd.OutOrStdout()
			if !plain && isTerminal(in.out) {
				return in.runInitTUI(cmd.Context(), in.out)
			}
			return in.run(cmd.Context())
		},
	}
	f := cmd.Flags()
	f.BoolVar(&plain, "plain", false, "plain output, even in a terminal")
	f.StringArrayVar(&in.hosts, "host", nil,
		"an address agents dial besides the ones found on this machine; repeatable")
	f.StringVar(&in.admin, "admin", "", "email of the first administrator; a password is generated and printed once")
	f.StringVar(&in.pkiDir, "pki", in.pkiDir, "directory holding the certificate authority")
	f.BoolVar(&in.noMirror, "no-mirror", false, "do not run a model mirror on this machine")
	f.StringVar(&in.mirrorURL, "mirror", "",
		"use the mirror on another machine, set up there by 'llmgw mirror init': its management API, http://host:8082")
	f.BoolVar(&in.noNode, "no-node", false, "do not make this machine a node (and serve no model on it)")
	f.BoolVar(&in.noModel, "no-model", false, "do not fetch and serve the starter model ("+starterName+")")
	f.BoolVar(&in.postgres, "postgres", false,
		"keep the store in PostgreSQL on this machine, installing it and moving an existing SQLite store into it")
	f.StringArrayVar(&in.patroni, "patroni", nil,
		"user@host of a machine to keep the store on, replicated under Patroni; give two or more. "+
			"Installs PostgreSQL, Patroni and etcd there over ssh, and moves the existing store in")
	return cmd
}

func (in *serverInit) run(ctx context.Context) error {
	steps := []struct {
		name string
		fn   func(context.Context) error
	}{
		{"user", in.ensureUser},
		{"directories", in.ensureDirs},
		{"certificates", in.ensureCerts},
		{"configuration", in.ensureConfig},
		{"database", in.ensureDatabase},
		{"store", in.ensureStore},
	}
	for _, s := range steps {
		if err := s.fn(ctx); err != nil {
			return fmt.Errorf("%s: %w", s.name, err)
		}
	}
	if in.unitPath == "" {
		return nil
	}
	hasSystemd := true
	if _, err := exec.LookPath("systemctl"); err != nil {
		hasSystemd = false
	}
	// Before the service starts, because the mirror goes into server.yaml.
	switch {
	case in.mirrorURL != "":
		if err := in.useMirror(ctx); err != nil {
			return fmt.Errorf("mirror: %w", err)
		}
	case hasSystemd && !in.noMirror:
		if err := in.ensureMirror(ctx); err != nil {
			return fmt.Errorf("mirror: %w", err)
		}
	}
	if err := in.ensureService(ctx); err != nil {
		return fmt.Errorf("service: %w", err)
	}
	if hasSystemd && !in.noNode {
		if _, err := in.ensureNode(ctx); err != nil {
			return fmt.Errorf("node: %w", err)
		}
	}
	switch {
	case !hasSystemd || in.noModel:
	case in.noMirror && in.mirrorURL == "":
		fmt.Fprintln(in.out, "  skip  no starter model without the mirror")
	default:
		if err := in.ensureStarter(ctx); err != nil {
			return fmt.Errorf("starter model: %w", err)
		}
	}
	in.nextSteps()
	return nil
}

func (in *serverInit) done(format string, a ...any) {
	fmt.Fprintf(in.out, "  ok    "+format+"\n", a...)
	in.changed = true
}

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

func (in *serverInit) ensureUser(context.Context) error {
	if in.user == "" {
		return nil
	}
	u, err := user.Lookup(in.user)
	if err != nil {
		var unknown user.UnknownUserError
		if !errors.As(err, &unknown) {
			return err
		}
		out, err := exec.Command("useradd", "--system", "--home", in.stateDir,
			"--shell", "/usr/sbin/nologin", in.user).CombinedOutput()
		if err != nil {
			return fmt.Errorf("useradd: %v: %s", err, strings.TrimSpace(string(out)))
		}
		if u, err = user.Lookup(in.user); err != nil {
			return err
		}
		in.done("system user %s", in.user)
	} else {
		in.kept("system user %s", in.user)
	}
	in.uid, _ = strconv.Atoi(u.Uid)
	in.gid, _ = strconv.Atoi(u.Gid)
	return nil
}

func (in *serverInit) ensureDirs(context.Context) error {
	for _, d := range []struct {
		path  string
		mode  os.FileMode
		owned bool
	}{
		{in.etcDir, 0o755, false},
		{in.pkiDir, 0o700, false},
		{in.stateDir, 0o750, true},
	} {
		if err := os.MkdirAll(d.path, d.mode); err != nil {
			return err
		}
		if d.owned {
			if err := in.chown(d.path, true); err != nil {
				return err
			}
		}
	}
	return nil
}

// chown gives a path to the service user, or only its group when the user
// should read and not write.
func (in *serverInit) chown(path string, own bool) error {
	if in.user == "" {
		return nil
	}
	uid := 0
	if own {
		uid = in.uid
	}
	return os.Chown(path, uid, in.gid)
}

func (in *serverInit) ensureCerts(context.Context) error {
	certPath := filepath.Join(in.etcDir, "server.pem")
	_, keyErr := os.Stat(filepath.Join(in.pkiDir, "ca-key.pem"))
	_, caErr := os.Stat(filepath.Join(in.etcDir, "ca.pem"))
	_, certErr := os.Stat(certPath)
	if keyErr != nil && caErr == nil && certErr == nil {
		// Issued by an authority kept elsewhere, which is the arrangement for
		// anyone who will not have the CA key on the control plane. Making a
		// new authority here would replace their certificate with one no node
		// trusts.
		if len(in.hosts) > 0 {
			return fmt.Errorf("%s was issued by an authority that is not in %s; "+
				"reissue it there with the new --host", certPath, in.pkiDir)
		}
		in.kept("certificates in %s, issued elsewhere", in.etcDir)
		return nil
	}

	var (
		ca  *tlsconf.CA
		err error
	)
	if keyErr == nil {
		if ca, err = loadCA(in.pkiDir); err != nil {
			return err
		}
		in.kept("certificate authority in %s", in.pkiDir)
	} else {
		if ca, err = tlsconf.NewCA("llmgw cluster CA"); err != nil {
			return err
		}
		if err := ca.Write(in.pkiDir); err != nil {
			return err
		}
		in.done("certificate authority in %s (ca-key.pem signs every node: back it up)", in.pkiDir)
	}
	if err := os.WriteFile(filepath.Join(in.etcDir, "ca.pem"), ca.PEM, 0o644); err != nil {
		return err
	}

	want := in.certHosts()
	if cert, err := readCert(certPath); err == nil {
		have := strings.Split(sanNames(cert), ", ")
		var missing []string
		for _, h := range want {
			if !slices.Contains(have, h) {
				missing = append(missing, h)
			}
		}
		// A certificate from this authority naming every address is left
		// alone. One naming fewer is reissued with the union, so an address
		// that worked yesterday is never dropped by adding one today.
		if len(missing) == 0 && cert.CheckSignatureFrom(ca.Cert) == nil {
			in.kept("server certificate for %s", strings.Join(have, ", "))
			return nil
		}
		for _, h := range have {
			if h != "" && !slices.Contains(want, h) {
				want = append(want, h)
			}
		}
	}

	name, _ := os.Hostname()
	if name == "" {
		name = "llmgw-control"
	}
	leaf, err := ca.IssueServer(name, want)
	if err != nil {
		return err
	}
	if err := leaf.Write(in.etcDir, "server"); err != nil {
		return err
	}
	// The service runs as llmgw and must read its key; nothing else should.
	keyPath := filepath.Join(in.etcDir, "server-key.pem")
	if err := in.chown(keyPath, false); err != nil {
		return err
	}
	if in.user != "" {
		if err := os.Chmod(keyPath, 0o640); err != nil {
			return err
		}
	}
	in.done("server certificate for %s", strings.Join(want, ", "))
	return nil
}

// certHosts is every address an agent could plausibly dial: the hostname and
// each address on this machine, plus whatever --host added. Asking the
// operator to list them is how a certificate ends up naming the one address
// nobody uses, and that is found out by the first node refusing it.
func (in *serverInit) certHosts() []string {
	var out []string
	add := func(h string) {
		if h != "" && !slices.Contains(out, h) {
			out = append(out, h)
		}
	}
	for _, h := range in.hosts {
		add(h)
	}
	if name, err := os.Hostname(); err == nil {
		add(name)
	}
	for _, ip := range localAddrs() {
		add(ip.String())
	}
	return out
}

// localAddrs are this machine's addresses another machine could reach:
// loopback and link-local are left out.
func localAddrs() []net.IP {
	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return nil
	}
	var out []net.IP
	for _, a := range addrs {
		n, ok := a.(*net.IPNet)
		if !ok || !n.IP.IsGlobalUnicast() {
			continue
		}
		out = append(out, n.IP)
	}
	return out
}

func (in *serverInit) configPath() string { return filepath.Join(in.etcDir, "server.yaml") }

func (in *serverInit) ensureConfig(context.Context) error {
	path := in.configPath()
	if _, err := os.Stat(path); err == nil {
		in.kept("%s (never overwritten)", path)
	} else {
		contact := in.admin
		if contact == "" {
			contact = "platform@example.com"
		}
		body := fmt.Sprintf(serverConfigTemplate, in.etcDir, in.etcDir, in.etcDir,
			in.storeAddress(), contact)
		if err := os.WriteFile(path, []byte(body), 0o640); err != nil {
			return err
		}
		if err := in.chown(path, false); err != nil {
			return err
		}
		in.done("%s", path)
	}

	// Checked whether written now or years ago: a configuration that will not
	// load is better found here than as a unit restarting every two seconds.
	cfg, err := config.Load(path)
	if err != nil {
		return err
	}
	if err := cfg.Validate(); err != nil {
		return fmt.Errorf("%s: %w", path, err)
	}
	return nil
}

// serverConfigTemplate is the smallest file that is not a toy: mutual TLS,
// authentication, default limits and a price, and a control plane with an
// empty catalog.
const serverConfigTemplate = `# Written by 'llmgw server init'. Every key: docs/reference.md
listen: ":8080"
admin_listen: "127.0.0.1:9090"   # serves pprof, which can dump API keys: keep it local
agent_listen: ":7443"

agent_tls:
  ca: %s/ca.pem
  cert: %s/server.pem
  key: %s/server-key.pem

# A control plane with nothing in its catalog yet. Models are added while it
# runs, with 'llmgw ctl model add' or from /admin/mirror.
control:
  enabled: true

auth:
  enabled: true
  store: "%s"   # a SQLite path, or @file holding a PostgreSQL address
  signup_open: true
  # Applied to any key no scope gives a limit to. Without them, a project
  # created without limits is metered by nothing at all.
  default_rpm: 60
  default_tpm: 200000
  # Without a price every request costs zero and no budget ever binds. 1000
  # everywhere turns a budget into a plain token quota.
  default_price:
    prompt_per_1k: 1000
    completion_per_1k: 1000

landing:
  contact: "%s"

# The engines every node installs for itself. Omitted, each gets the llama.cpp
# build this release was checked against.
# engines:
#   llamacpp: b10991

# A model mirror, once there is one (docs/install.md, step 3):
# mirror:
#   api: https://models.internal:8082
#   tls: { ca: /etc/llmgw/ca.pem, cert: /etc/llmgw/server.pem, key: /etc/llmgw/server-key.pem }
`

func (in *serverInit) ensureStore(ctx context.Context) error {
	cfg, err := config.Load(in.configPath())
	if err != nil {
		return err
	}
	address := cfg.Auth.Store
	if address == "" {
		address = filepath.Join(in.stateDir, "llmgw.db")
	}
	s, err := store.Open(ctx, address)
	if err != nil {
		return fmt.Errorf("opening %s: %w", address, err)
	}
	// Opened as root, so a SQLite file created just now belongs to root and
	// the service could not write it. Given back before anything else runs.
	defer in.chownSQLite(address)
	defer s.Close()

	// The tree an approved account hangs off. Approving into a project that
	// does not exist succeeds, and the first key that person makes does not.
	if err := s.PutOrg(ctx, store.Org{ID: "default", Name: "default"}); err != nil {
		return err
	}
	if err := s.PutTeam(ctx, store.Team{ID: "default", OrgID: "default", Name: "default"}); err != nil {
		return err
	}
	if err := s.PutProject(ctx, store.Project{ID: "default", TeamID: "default", Name: "default"}); err != nil {
		return err
	}

	if in.admin == "" {
		in.kept("store at %s (no --admin: create an account with 'llmgw users create')", address)
		return nil
	}
	if u, err := s.UserByEmail(ctx, in.admin); err == nil {
		if !u.IsAdmin {
			if err := s.SetAdmin(ctx, u.ID, true); err != nil {
				return err
			}
		}
		in.kept("administrator %s (%s); its password is unchanged", in.admin, u.ID)
		return nil
	} else if !errors.Is(err, store.ErrNotFound) {
		return err
	}

	password := randomPassword()
	hash, err := account.HashPassword(password)
	if err != nil {
		return err
	}
	id := randomID("usr")
	now := time.Now()
	if err := s.CreateUser(ctx, store.User{
		ID: id, Email: in.admin, Name: in.admin,
		PasswordHash: hash, State: store.UserPending, CreatedAt: now,
	}); err != nil {
		return err
	}
	if err := s.ApproveUser(ctx, id, "default", "init", now); err != nil {
		return err
	}
	if err := s.SetAdmin(ctx, id, true); err != nil {
		return err
	}
	in.done("administrator %s (%s)", in.admin, id)
	fmt.Fprintf(in.out, "\n        password: %s\n        shown once — sign in and change it from /profile\n\n", password)
	return nil
}

func (in *serverInit) chownSQLite(address string) {
	if strings.Contains(address, "://") || strings.HasPrefix(address, "@") {
		return
	}
	for _, suffix := range []string{"", "-wal", "-shm"} {
		if _, err := os.Stat(address + suffix); err == nil {
			_ = in.chown(address+suffix, true)
		}
	}
}

func randomPassword() string {
	var buf [15]byte
	_, _ = rand.Read(buf[:])
	return strings.ToLower(base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(buf[:]))
}

func (in *serverInit) ensureService(ctx context.Context) error {
	if _, err := exec.LookPath("systemctl"); err != nil {
		fmt.Fprintf(in.out, "  skip  no systemd here; start it with:\n"+
			"        sudo -u %s llmgw server --config %s\n", in.user, in.configPath())
		return nil
	}
	if b, err := os.ReadFile(in.unitPath); err == nil && string(b) == serverUnit {
		in.kept("%s", in.unitPath)
	} else {
		if err == nil {
			// Somebody's edits. Kept beside it, because the unit is also where
			// people put environment overrides and those are not ours to lose.
			_ = os.WriteFile(in.unitPath+".bak", b, 0o644)
		}
		if err := os.WriteFile(in.unitPath, []byte(serverUnit), 0o644); err != nil {
			return err
		}
		in.done("%s", in.unitPath)
	}

	systemctl := func(args ...string) error {
		out, err := exec.CommandContext(ctx, "systemctl", args...).CombinedOutput()
		if err != nil {
			return fmt.Errorf("systemctl %s: %v: %s", strings.Join(args, " "), err, strings.TrimSpace(string(out)))
		}
		return nil
	}
	if err := systemctl("daemon-reload"); err != nil {
		return err
	}
	active := exec.CommandContext(ctx, "systemctl", "is-active", "--quiet", "llmgw-server").Run() == nil
	switch {
	case !active:
		if err := systemctl("enable", "--now", "llmgw-server"); err != nil {
			return err
		}
		in.done("llmgw-server started")
	case in.changed:
		// A certificate or configuration it has not read is one it is not
		// using, and the difference shows up as a node that cannot connect.
		if err := systemctl("restart", "llmgw-server"); err != nil {
			return err
		}
		in.done("llmgw-server restarted to pick up the changes")
	default:
		in.kept("llmgw-server running")
	}
	return in.waitHealthy(ctx)
}

func (in *serverInit) waitHealthy(ctx context.Context) error {
	cfg, err := config.Load(in.configPath())
	if err != nil {
		return err
	}
	_, port, err := net.SplitHostPort(cfg.Listen)
	if err != nil {
		return fmt.Errorf("listen %q: %w", cfg.Listen, err)
	}
	url := "http://127.0.0.1:" + port + "/health"
	deadline := time.Now().Add(30 * time.Second)
	for {
		req, _ := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
		resp, err := http.DefaultClient.Do(req)
		if err == nil {
			resp.Body.Close()
			if resp.StatusCode == http.StatusOK {
				in.kept("answering on %s", url)
				return nil
			}
		}
		if time.Now().After(deadline) {
			return fmt.Errorf("not answering on %s after 30s; 'journalctl -u llmgw-server -n 50' says why", url)
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(time.Second):
		}
	}
}

func (in *serverInit) nextSteps() {
	fmt.Fprint(in.out, `
The control plane is up. Next:

  llmgw ctl status                      # the fleet in one answer

  # each node, from this machine (it holds the authority):
  sudo llmgw node add <user>@<node address>

  # then a model: /admin/mirror, "Explore models", or
  llmgw ctl model add -f model.yaml

The web UI is on port 8080; /admin once you are signed in as the administrator.
`)
}