llmgw/internal/cli/backup.go

package cli

import (
	"archive/tar"
	"compress/gzip"
	"context"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"slices"
	"strings"
	"time"

	"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"
)

// DefaultBackupDir is where backups are written and where the control plane
// looks for the newest one.
const DefaultBackupDir = "/var/backups/llmgw"

func newBackupCmd() *cobra.Command {
	var (
		configPath string
		etcDir     string
		dir        string
		keep       int
	)
	c := &cobra.Command{
		Use:   "backup",
		Short: "Write the store and /etc/llmgw to one archive, and keep the last few",
		Long: "Everything a control plane cannot be rebuilt without: the store — keys,\n" +
			"people, quotas, usage, the models added at runtime, the secret engine\n" +
			"keys derive from — and /etc/llmgw, which holds the configuration, the\n" +
			"store's address and the certificate authority's key.\n\n" +
			"The archive holds that key, so it is written readable by root only. It\n" +
			"is on the same disk as what it backs up: copy it elsewhere, or it does\n" +
			"not survive the loss of this machine.\n\n" +
			"Run daily by llmgw-backup.timer (`llmgw server backup`). The control plane publishes the time of\n" +
			"the newest archive as llmgw_backup_last_success_timestamp_seconds, and\n" +
			"BackupTooOld fires when it is more than a day old.",
		RunE: func(cmd *cobra.Command, _ []string) error {
			cfg, err := config.Load(configPath)
			if err != nil {
				return err
			}
			path, err := writeBackup(cmd.Context(), cfg.Auth.Store, etcDir, dir)
			if err != nil {
				return err
			}
			removed, err := pruneBackups(dir, keep)
			if err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "wrote %s", path)
			if removed > 0 {
				fmt.Fprintf(cmd.OutOrStdout(), "; removed %d older", removed)
			}
			fmt.Fprintln(cmd.OutOrStdout())
			return nil
		},
	}
	c.Flags().StringVar(&configPath, "config", "/etc/llmgw/server.yaml", "the control plane's configuration")
	c.Flags().StringVar(&etcDir, "etc", "/etc/llmgw", "configuration directory to include")
	c.Flags().StringVar(&dir, "dir", DefaultBackupDir, "where archives are written")
	c.Flags().IntVar(&keep, "keep", 14, "how many archives to keep")
	return c
}

const backupPrefix = "llmgw-backup-"

// writeBackup writes one archive and returns its path. Written beside its
// final name and renamed, so a backup interrupted halfway is never the newest
// one the control plane reports.
func writeBackup(ctx context.Context, storeAddr, etcDir, dir string) (string, error) {
	// Listable by anyone, so the control plane — not root — can see when the
	// newest archive was written. The archives themselves are root's alone.
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return "", err
	}
	host, _ := os.Hostname()
	name := backupPrefix + host + "-" + time.Now().UTC().Format("20060102T150405Z") + ".tar.gz"
	final := filepath.Join(dir, name)
	part := final + ".part"

	f, err := os.OpenFile(part, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
	if err != nil {
		return "", err
	}
	ok := false
	defer func() {
		if !ok {
			_ = os.Remove(part)
		}
	}()
	gz := gzip.NewWriter(f)
	tw := tar.NewWriter(gz)

	manifest := fmt.Sprintf("llmgw %s\nhost %s\ntime %s\n", version.Version, host,
		time.Now().UTC().Format(time.RFC3339))
	if err := addBytes(tw, "MANIFEST", []byte(manifest)); err != nil {
		return "", err
	}
	if storeAddr != "" {
		if err := addStore(ctx, tw, storeAddr, dir); err != nil {
			return "", fmt.Errorf("the store: %w", err)
		}
	}
	if err := addTree(tw, etcDir, "etc/llmgw"); err != nil {
		return "", fmt.Errorf("%s: %w", etcDir, err)
	}
	if err := errors.Join(tw.Close(), gz.Close(), f.Sync(), f.Close()); err != nil {
		return "", err
	}
	if err := os.Rename(part, final); err != nil {
		return "", err
	}
	ok = true
	return final, nil
}

// addStore dumps the store into the archive.
func addStore(ctx context.Context, tw *tar.Writer, addr, scratch string) error {
	addr, err := store.ResolveAddress(addr)
	if err != nil {
		return err
	}
	tmp, err := os.CreateTemp(scratch, ".store-*")
	if err != nil {
		return err
	}
	tmpPath := tmp.Name()
	_ = tmp.Close()
	defer os.Remove(tmpPath)

	name := "store.sqlite"
	if store.IsPostgres(addr) {
		name = "store.pgdump"
		if err := pgDump(ctx, addr, tmpPath); err != nil {
			return err
		}
	} else {
		// VACUUM INTO refuses a file that exists.
		_ = os.Remove(tmpPath)
		db, err := store.Open(ctx, addr)
		if err != nil {
			return err
		}
		err = db.SnapshotTo(ctx, tmpPath)
		_ = db.Close()
		if err != nil {
			return err
		}
	}
	return addFile(tw, tmpPath, name)
}

// pgDump runs pg_dump in its custom format, which pg_restore can restore
// selectively. The connection is given through the environment, not the
// command line: a DSN carries the password, and every user on the machine
// can read another process's arguments.
func pgDump(ctx context.Context, dsn, out string) error {
	u, err := url.Parse(dsn)
	if err != nil {
		return fmt.Errorf("the store address is not a URL")
	}
	env := os.Environ()
	set := func(k, v string) {
		if v != "" {
			env = append(env, k+"="+v)
		}
	}
	set("PGHOST", u.Hostname())
	set("PGPORT", u.Port())
	set("PGUSER", u.User.Username())
	if pw, ok := u.User.Password(); ok {
		set("PGPASSWORD", pw)
	}
	set("PGDATABASE", strings.TrimPrefix(u.Path, "/"))
	set("PGSSLMODE", u.Query().Get("sslmode"))
	cmd := exec.CommandContext(ctx, "pg_dump", "--format=custom", "--no-owner", "--file="+out) //nolint:gosec // fixed binary, our own file
	cmd.Env = env
	if b, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("pg_dump: %v: %s", err, strings.TrimSpace(string(b)))
	}
	return nil
}

func addTree(tw *tar.Writer, root, as string) error {
	return filepath.WalkDir(root, func(p string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}
		rel, _ := filepath.Rel(root, p)
		name := filepath.ToSlash(filepath.Join(as, rel))
		info, err := d.Info()
		if err != nil {
			return err
		}
		switch {
		case d.IsDir():
			return tw.WriteHeader(&tar.Header{Name: name + "/", Mode: int64(info.Mode().Perm()),
				Typeflag: tar.TypeDir, ModTime: info.ModTime()})
		case info.Mode().IsRegular():
			return addFile(tw, p, name)
		}
		return nil // sockets, links: not configuration
	})
}

func addFile(tw *tar.Writer, path, name string) error {
	f, err := os.Open(path) //nolint:gosec // files we chose to back up
	if err != nil {
		return err
	}
	defer f.Close()
	info, err := f.Stat()
	if err != nil {
		return err
	}
	if err := tw.WriteHeader(&tar.Header{Name: name, Mode: int64(info.Mode().Perm()),
		Size: info.Size(), ModTime: info.ModTime(), Typeflag: tar.TypeReg}); err != nil {
		return err
	}
	_, err = io.Copy(tw, f)
	return err
}

func addBytes(tw *tar.Writer, name string, b []byte) error {
	if err := tw.WriteHeader(&tar.Header{Name: name, Mode: 0o600, Size: int64(len(b)),
		ModTime: time.Now(), Typeflag: tar.TypeReg}); err != nil {
		return err
	}
	_, err := tw.Write(b)
	return err
}

// pruneBackups keeps the newest keep archives. Only files this command
// wrote are considered, so nothing else in the directory is ever removed.
func pruneBackups(dir string, keep int) (int, error) {
	entries, err := os.ReadDir(dir)
	if err != nil {
		return 0, err
	}
	var names []string
	for _, e := range entries {
		if n := e.Name(); strings.HasPrefix(n, backupPrefix) && strings.HasSuffix(n, ".tar.gz") {
			names = append(names, n)
		}
	}
	// The name carries a UTC timestamp, so name order is age order.
	slices.Sort(names)
	removed := 0
	for len(names) > max(keep, 1) {
		if err := os.Remove(filepath.Join(dir, names[0])); err != nil {
			return removed, err
		}
		names = names[1:]
		removed++
	}
	return removed, nil
}

// NewestBackup is the modification time of the newest complete archive in
// dir, or zero when there is none.
func NewestBackup(dir string) time.Time {
	entries, err := os.ReadDir(dir)
	if err != nil {
		return time.Time{}
	}
	var newest time.Time
	for _, e := range entries {
		if n := e.Name(); !strings.HasPrefix(n, backupPrefix) || !strings.HasSuffix(n, ".tar.gz") {
			continue
		}
		if info, err := e.Info(); err == nil && info.ModTime().After(newest) {
			newest = info.ModTime()
		}
	}
	return newest
}