llmgw/internal/cli/restore.go

package cli

import (
	"archive/tar"
	"bufio"
	"bytes"
	"compress/gzip"
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	agecrypt "filippo.io/age"
	"filippo.io/age/agessh"
	"github.com/spf13/cobra"

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

// restore is `llmgw server restore`: an archive `server backup` wrote, put
// back on this machine.
type restore struct {
	archive    string
	identities []string
	etcDir     string
	storeAddr  string // overrides the restored server.yaml's
	dryRun     bool
	force      bool
	// service is the unit stopped before and started after; "" touches no
	// service, for tests and machines without systemd.
	service string
	out     io.Writer
	now     func() time.Time
}

func newRestoreCmd() *cobra.Command {
	rs := restore{etcDir: "/etc/llmgw", service: "llmgw-server", now: time.Now}
	c := &cobra.Command{
		Use:   "restore ARCHIVE",
		Short: "Put a backup back: /etc/llmgw and the store, on this machine",
		Long: "Reads an archive `llmgw server backup` wrote — opening an encrypted one\n" +
			"with --identity, an age identity file or an SSH private key — and checks\n" +
			"it before anything is touched: that it is whole, that this llmgw is not\n" +
			"older than the one that wrote it, and where the store goes. --dry-run\n" +
			"stops there and says what it would do.\n\n" +
			"Then: the server is stopped, /etc/llmgw is moved aside to\n" +
			"/etc/llmgw.before-restore-<time> and the archive's put in its place, the\n" +
			"store is restored — a SQLite file replaced, the old one kept beside it;\n" +
			"PostgreSQL by pg_restore --clean into the database the address names,\n" +
			"which must exist — and the server is started and asked whether it\n" +
			"answers.\n\n" +
			"Running engines keep serving throughout: their agents reconnect, and the\n" +
			"engine-key secret came back with the store. For a whole drill, see\n" +
			"\"Restoring\" in docs/install.md.",
		Example: "  sudo llmgw server restore /var/backups/llmgw-server/llmgw-backup-ctl-20261004T020000Z.tar.gz --dry-run\n" +
			"  sudo llmgw server restore backup.tar.gz.age --identity ~/llmgw-backup-key.txt",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			rs.archive, rs.out = args[0], cmd.OutOrStdout()
			return rs.run(cmd.Context())
		},
	}
	f := c.Flags()
	f.StringArrayVar(&rs.identities, "identity", nil, "a key that opens an encrypted archive: an age identity file or an SSH private key; repeatable")
	f.StringVar(&rs.etcDir, "etc", rs.etcDir, "the configuration directory to restore into")
	f.StringVar(&rs.storeAddr, "store", "", "restore the store here rather than where the archive's server.yaml says: a SQLite path or a postgres:// address")
	f.BoolVar(&rs.dryRun, "dry-run", false, "check the archive and say what would be done; change nothing")
	f.BoolVar(&rs.force, "force", false, "restore an archive written by a newer llmgw")
	f.StringVar(&rs.service, "service", rs.service, "the unit stopped before and started after; empty touches none")
	return c
}

func (rs *restore) say(format string, a ...any) { fmt.Fprintf(rs.out, format+"\n", a...) }

func (rs *restore) run(ctx context.Context) error {
	tmp, err := os.MkdirTemp("", "llmgw-restore-")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tmp)

	// 1. The archive, opened and unpacked where only this user can read it.
	man, err := rs.unpack(tmp)
	if err != nil {
		return err
	}
	rs.say("archive   written by llmgw %s on %s at %s", man.version, man.host, man.time)
	if newer(man.version, version.Version) && !rs.force {
		return fmt.Errorf("the archive was written by llmgw %s, newer than this one (%s): its store may have "+
			"changes this one does not know. Install %s or newer first, or pass --force", man.version, version.Version, man.version)
	}
	restoredEtc := filepath.Join(tmp, "etc", "llmgw")
	if _, err := os.Stat(filepath.Join(restoredEtc, "server.yaml")); err != nil {
		return errors.New("the archive has no etc/llmgw/server.yaml: not a control plane's backup")
	}
	n := 0
	_ = filepath.WalkDir(restoredEtc, func(_ string, d os.DirEntry, err error) error {
		if err == nil && !d.IsDir() {
			n++
		}
		return nil
	})
	rs.say("config    %d files for %s", n, rs.etcDir)

	// 2. Where the store goes, and what it is.
	dump, kind := filepath.Join(tmp, "store.sqlite"), "sqlite"
	if _, err := os.Stat(dump); err != nil {
		dump, kind = filepath.Join(tmp, "store.pgdump"), "postgres"
		if _, err := os.Stat(dump); err != nil {
			dump, kind = "", ""
		}
	}
	target, err := rs.storeTarget(restoredEtc)
	if err != nil {
		return err
	}
	switch {
	case kind == "":
		rs.say("store     none in the archive; nothing to restore")
	case kind == "postgres" && !store.IsPostgres(target):
		return fmt.Errorf("the archive holds a PostgreSQL dump and the store goes to %s: "+
			"give --store a postgres:// address", target)
	case kind == "sqlite" && store.IsPostgres(target):
		return fmt.Errorf("the archive holds a SQLite store and the store goes to PostgreSQL: give --store a file, " +
			"then `llmgw keys move-store` copies it into PostgreSQL")
	case kind == "postgres":
		if _, err := exec.LookPath("pg_restore"); err != nil {
			return errors.New("the store is a PostgreSQL dump and pg_restore is not installed here (apt-get install postgresql-client)")
		}
		rs.say("store     PostgreSQL dump into %s, with pg_restore --clean", redactDSN(target))
	default:
		rs.say("store     SQLite file to %s%s", target, aside(target, rs.now()))
	}
	if rs.service != "" {
		if exec.Command("systemctl", "is-active", "--quiet", rs.service).Run() == nil { //nolint:gosec // our unit
			rs.say("service   %s is running: stopped first, started after", rs.service)
		} else {
			rs.say("service   %s is not running: started after", rs.service)
		}
	}
	if _, err := os.Stat(rs.etcDir); err == nil {
		rs.say("current   %s moved aside to %s", rs.etcDir, rs.etcDir+".before-restore-"+stamp(rs.now()))
	}
	if rs.dryRun {
		rs.say("\ndry run: nothing was changed")
		return nil
	}

	// 3. Done, in an order that can be undone by hand at every step.
	if rs.service != "" {
		_ = exec.CommandContext(ctx, "systemctl", "stop", rs.service).Run() //nolint:gosec // our unit
	}
	if _, err := os.Stat(rs.etcDir); err == nil {
		if err := os.Rename(rs.etcDir, rs.etcDir+".before-restore-"+stamp(rs.now())); err != nil {
			return err
		}
	}
	if err := copyTree(restoredEtc, rs.etcDir); err != nil {
		return fmt.Errorf("putting back %s: %w", rs.etcDir, err)
	}
	rs.say("ok        %s restored", rs.etcDir)
	switch kind {
	case "sqlite":
		if err := restoreSQLite(dump, target, rs.now()); err != nil {
			return fmt.Errorf("the store: %w", err)
		}
		rs.say("ok        store restored to %s", target)
	case "postgres":
		if err := pgRestore(ctx, target, dump); err != nil {
			return fmt.Errorf("the store: %w", err)
		}
		rs.say("ok        store restored into %s", redactDSN(target))
	}
	if rs.service != "" {
		if err := exec.CommandContext(ctx, "systemctl", "start", rs.service).Run(); err != nil { //nolint:gosec // our unit
			return fmt.Errorf("starting %s: %w; `journalctl -u %s` says why", rs.service, err, rs.service)
		}
		if err := rs.waitHealthy(ctx); err != nil {
			return err
		}
		rs.say("ok        %s answers", rs.service)
	}
	rs.say("\nRestored. `llmgw ctl status` shows the cluster as it comes back; agents reconnect by themselves.")
	return nil
}

type manifest struct{ version, host, time string }

// unpack opens the archive — decrypting it when it is sealed — and writes
// its files under dir, keeping their modes and, by name, their owners.
func (rs *restore) unpack(dir string) (manifest, error) {
	f, err := os.Open(rs.archive)
	if err != nil {
		return manifest{}, err
	}
	defer f.Close()
	var r io.Reader = bufio.NewReader(f)
	if strings.HasSuffix(rs.archive, ".age") {
		if len(rs.identities) == 0 {
			return manifest{}, errors.New("the archive is encrypted: give the key that opens it with --identity")
		}
		ids, err := readIdentities(rs.identities)
		if err != nil {
			return manifest{}, err
		}
		if r, err = agecrypt.Decrypt(r, ids...); err != nil {
			return manifest{}, fmt.Errorf("opening the archive: %w", err)
		}
	}
	gz, err := gzip.NewReader(r)
	if err != nil {
		return manifest{}, fmt.Errorf("not an archive server backup wrote: %w", err)
	}
	tr := tar.NewReader(gz)
	var man manifest
	for {
		h, err := tr.Next()
		if errors.Is(err, io.EOF) {
			break
		}
		if err != nil {
			return manifest{}, fmt.Errorf("the archive is damaged: %w", err)
		}
		name := filepath.Clean(h.Name)
		if filepath.IsAbs(name) || strings.HasPrefix(name, "..") {
			return manifest{}, fmt.Errorf("the archive names %q, outside itself", h.Name)
		}
		p := filepath.Join(dir, name)
		switch h.Typeflag {
		case tar.TypeDir:
			if err := os.MkdirAll(p, 0o700); err != nil {
				return manifest{}, err
			}
		case tar.TypeReg:
			if err := os.MkdirAll(filepath.Dir(p), 0o700); err != nil {
				return manifest{}, err
			}
			out, err := os.OpenFile(p, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
			if err != nil {
				return manifest{}, err
			}
			_, err = io.Copy(out, tr)
			if cerr := out.Close(); err == nil {
				err = cerr
			}
			if err != nil {
				return manifest{}, err
			}
			if name == "MANIFEST" {
				b, _ := os.ReadFile(p)
				man = parseManifest(b)
			}
		default:
			continue
		}
		// The mode and owner, applied when the tree is copied into place.
		if err := writeMeta(p, h); err != nil {
			return manifest{}, err
		}
	}
	if man.version == "" {
		return manifest{}, errors.New("the archive has no MANIFEST: not one server backup wrote, or cut short")
	}
	return man, nil
}

func parseManifest(b []byte) manifest {
	var m manifest
	for _, l := range strings.Split(string(b), "\n") {
		k, v, _ := strings.Cut(l, " ")
		switch k {
		case "llmgw":
			m.version = v
		case "host":
			m.host = v
		case "time":
			m.time = v
		}
	}
	return m
}

// meta is a restored file's mode and owner, kept beside it until it is
// copied into place.
const metaSuffix = ".llmgw-restore-meta"

func writeMeta(p string, h *tar.Header) error {
	return os.WriteFile(p+metaSuffix, fmt.Appendf(nil, "%o %s %s %d %d", h.Mode&0o7777, h.Uname, h.Gname, h.Uid, h.Gid), 0o600)
}

// copyTree puts the unpacked tree in place with each file's mode and owner.
func copyTree(from, to string) error {
	return filepath.WalkDir(from, func(p string, d os.DirEntry, err error) error {
		if err != nil || strings.HasSuffix(p, metaSuffix) {
			return err
		}
		rel, _ := filepath.Rel(from, p)
		dst := filepath.Join(to, rel)
		mode, uid, gid := readMeta(p, d.IsDir())
		if d.IsDir() {
			if err := os.MkdirAll(dst, mode); err != nil {
				return err
			}
		} else {
			b, err := os.ReadFile(p) //nolint:gosec // our own temporary copy
			if err != nil {
				return err
			}
			if err := os.WriteFile(dst, b, mode); err != nil {
				return err
			}
		}
		if err := os.Chmod(dst, mode); err != nil {
			return err
		}
		if os.Geteuid() == 0 && uid >= 0 {
			return os.Lchown(dst, uid, gid)
		}
		return nil
	})
}

// readMeta is a file's mode, and its owner resolved by name on this machine,
// falling back to the number; -1 leaves it as it is.
func readMeta(p string, dir bool) (os.FileMode, int, int) {
	mode := os.FileMode(0o640)
	if dir {
		mode = 0o750
	}
	b, err := os.ReadFile(p + metaSuffix)
	if err != nil {
		return mode, -1, -1
	}
	var m uint32
	var un, gn string
	var uid, gid int
	if _, err := fmt.Sscanf(string(b), "%o %s %s %d %d", &m, &un, &gn, &uid, &gid); err != nil {
		// Names may be empty: read what is there.
		f := strings.Fields(string(b))
		if len(f) == 0 {
			return mode, -1, -1
		}
		v, _ := strconv.ParseUint(f[0], 8, 32)
		return os.FileMode(v), -1, -1
	}
	if u, err := user.Lookup(un); err == nil {
		uid, _ = strconv.Atoi(u.Uid)
	}
	if g, err := user.LookupGroup(gn); err == nil {
		gid, _ = strconv.Atoi(g.Gid)
	}
	return os.FileMode(m), uid, gid
}

func readIdentities(paths []string) ([]agecrypt.Identity, error) {
	var out []agecrypt.Identity
	for _, p := range paths {
		b, err := os.ReadFile(p) //nolint:gosec // the operator's key file
		if err != nil {
			return nil, err
		}
		if bytes.Contains(b, []byte("PRIVATE KEY")) {
			id, err := agessh.ParseIdentity(b)
			if err != nil {
				return nil, fmt.Errorf("%s: %w", p, err)
			}
			out = append(out, id)
			continue
		}
		ids, err := agecrypt.ParseIdentities(bytes.NewReader(b))
		if err != nil {
			return nil, fmt.Errorf("%s is neither an age identity file nor an SSH private key: %w", p, err)
		}
		out = append(out, ids...)
	}
	return out, nil
}

// storeTarget is where the store is restored: --store, or the restored
// server.yaml's auth.store, read from where it was unpacked.
func (rs *restore) storeTarget(restoredEtc string) (string, error) {
	if rs.storeAddr != "" {
		return store.ResolveAddress(rs.storeAddr)
	}
	cfg, err := config.Load(filepath.Join(restoredEtc, "server.yaml"))
	if err != nil {
		return "", fmt.Errorf("the archive's server.yaml: %w", err)
	}
	addr := cfg.Auth.Store
	if addr == "" {
		return "/var/lib/llmgw/llmgw.db", nil
	}
	// "@/etc/llmgw/secrets/…" names a file the archive brought back.
	if f, ok := strings.CutPrefix(addr, "@"); ok && strings.HasPrefix(f, rs.etcDir+"/") {
		addr = "@" + filepath.Join(restoredEtc, strings.TrimPrefix(f, rs.etcDir+"/"))
	}
	return store.ResolveAddress(addr)
}

func restoreSQLite(dump, target string, now time.Time) error {
	if err := os.MkdirAll(filepath.Dir(target), 0o750); err != nil {
		return err
	}
	if _, err := os.Stat(target); err == nil {
		if err := os.Rename(target, target+".before-restore-"+stamp(now)); err != nil {
			return err
		}
		for _, side := range []string{"-wal", "-shm"} {
			_ = os.Remove(target + side)
		}
	}
	b, err := os.ReadFile(dump) //nolint:gosec // our own temporary copy
	if err != nil {
		return err
	}
	if err := os.WriteFile(target, b, 0o600); err != nil {
		return err
	}
	if os.Geteuid() == 0 {
		if u, err := user.Lookup("llmgw"); err == nil {
			uid, _ := strconv.Atoi(u.Uid)
			gid, _ := strconv.Atoi(u.Gid)
			return os.Chown(target, uid, gid)
		}
	}
	return nil
}

// pgRestore loads a dump into the database the address names, as its
// owner, on the primary.
func pgRestore(ctx context.Context, dsn, dump string) error {
	env, err := pgEnv(dsn)
	if err != nil {
		return err
	}
	for i, e := range env {
		if strings.HasPrefix(e, "PGTARGETSESSIONATTRS=") {
			env[i] = "PGTARGETSESSIONATTRS=read-write"
		}
	}
	env = append(env, "PGTARGETSESSIONATTRS=read-write")
	db := ""
	for _, e := range env {
		if v, ok := strings.CutPrefix(e, "PGDATABASE="); ok {
			db = v
		}
	}
	if db == "" {
		return errors.New("the store's address names no database")
	}
	// Without --dbname pg_restore writes SQL out rather than restoring; the
	// rest of the connection comes from the environment, the password with it.
	cmd := exec.CommandContext(ctx, "pg_restore", "--clean", "--if-exists", "--no-owner", "--single-transaction", //nolint:gosec // fixed binary, our own file
		"--exit-on-error", "--dbname="+db, dump)
	cmd.Env = append(os.Environ(), env...)
	if b, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("pg_restore: %v: %s", err, strings.TrimSpace(string(b)))
	}
	return nil
}

// waitHealthy waits for the restored server to answer, for a minute.
func (rs *restore) waitHealthy(ctx context.Context) error {
	cfg, err := config.Load(filepath.Join(rs.etcDir, "server.yaml"))
	if err != nil {
		return err
	}
	c := newCtlClient(cfg.AdminListen, adminTokenPath(filepath.Join(rs.etcDir, "server.yaml"), cfg))
	deadline := time.Now().Add(time.Minute)
	for {
		var v map[string]any
		if err := c.get(ctx, "/api/v1/summary", &v); err == nil {
			return nil
		} else if time.Now().After(deadline) {
			return fmt.Errorf("%s started but does not answer: %w", rs.service, err)
		}
		time.Sleep(2 * time.Second)
	}
}

func stamp(t time.Time) string { return t.UTC().Format("20060102T150405Z") }

func aside(target string, now time.Time) string {
	if _, err := os.Stat(target); err == nil {
		return "; the current one kept as " + filepath.Base(target) + ".before-restore-" + stamp(now)
	}
	return ""
}

// redactDSN is an address without its password, for saying where.
func redactDSN(dsn string) string {
	if i := strings.Index(dsn, "://"); i >= 0 {
		if j := strings.Index(dsn[i+3:], "@"); j >= 0 {
			cred := dsn[i+3 : i+3+j]
			if k := strings.IndexByte(cred, ':'); k >= 0 {
				return dsn[:i+3] + cred[:k] + ":…" + dsn[i+3+j:]
			}
		}
	}
	return dsn
}

// newer says whether version a is newer than b, by their leading
// major.minor.patch; anything unreadable is not newer.
func newer(a, b string) bool {
	pa, oka := semverOf(a)
	pb, okb := semverOf(b)
	if !oka || !okb {
		return false
	}
	for i := range pa {
		if pa[i] != pb[i] {
			return pa[i] > pb[i]
		}
	}
	return false
}

func semverOf(v string) ([3]int, bool) {
	v = strings.TrimPrefix(v, "v")
	if i := strings.IndexAny(v, "-+"); i >= 0 {
		v = v[:i]
	}
	parts := strings.Split(v, ".")
	var out [3]int
	if len(parts) != 3 {
		return out, false
	}
	for i, p := range parts {
		n, err := strconv.Atoi(p)
		if err != nil {
			return out, false
		}
		out[i] = n
	}
	return out, true
}