nmon/internal/backup/restore.go

package backup

import (
	"archive/tar"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"

	"github.com/v10networks/nmon/internal/clock"
	"github.com/v10networks/nmon/internal/config"
	"github.com/v10networks/nmon/internal/store"
)

// Restore errors.
var (
	// ErrNotEmpty means the target data directory already holds files; Restore never merges.
	ErrNotEmpty = errors.New("restore: data directory is not empty")
	// ErrTruncated means the archive ended early or does not match its manifest.
	ErrTruncated = errors.New("restore: archive is truncated or damaged")
	// ErrVerify means the restored data does not match the manifest.
	ErrVerify = errors.New("restore: verification failed")
)

const stagingDir = ".restore-staging"

// RestoreOptions configures Restore.
type RestoreOptions struct {
	// File is the archive; DataDir the clean directory to restore into (absent or empty).
	File, DataDir string
	// KeyFile is where an included key file is written (mode 0600); empty means <DataDir>/secret.key.
	KeyFile string
	// Store receives a JSON-lines inventory (Postgres/MySQL archives): an open, migrated, empty database.
	// SQLite archives restore to <DataDir>/inventory.db and ignore it.
	Store store.Store
	// Verifier checks the restored metric stores' sampled series; nil uses NewVMVerifier with defaults.
	Verifier SeriesVerifier
	Clock    clock.Clock
	Log      io.Writer
}

// RestoreResult reports what was restored and verified.
type RestoreResult struct {
	Manifest *Manifest
	// KeyRestoredTo is empty when the archive had no key.
	KeyRestoredTo string
	// RowsVerified and SeriesVerified count what matched the manifest.
	RowsVerified   int64
	SeriesVerified int
}

// Restore unpacks file into a clean data directory and verifies the result before returning success.
// Any failure removes what it created, so a retry starts clean.
func Restore(ctx context.Context, o RestoreOptions) (res *RestoreResult, err error) {
	if o.Clock == nil {
		o.Clock = clock.Real()
	}
	if o.Log == nil {
		o.Log = io.Discard
	}
	if err := requireClean(o.DataDir); err != nil {
		return nil, err
	}
	defer func() {
		if err != nil {
			err = errors.Join(err, emptyDir(o.DataDir))
		}
	}()
	if err := os.MkdirAll(filepath.Join(o.DataDir, stagingDir), 0o700); err != nil {
		return nil, err
	}
	m, err := extract(o.File, filepath.Join(o.DataDir, stagingDir))
	if err != nil {
		return nil, err
	}
	logf(o.Log, "restore: archive complete (%d files, backup of %s, key %s)\n", len(m.Files),
		m.CreatedAt.Format(time.RFC3339), map[bool]string{true: "included", false: "not included"}[m.KeyIncluded])
	res = &RestoreResult{Manifest: m}
	if err := place(&o, m, res); err != nil {
		return nil, err
	}
	if err := verify(ctx, &o, m, res); err != nil {
		return nil, fmt.Errorf("%w: %w", ErrVerify, err)
	}
	logf(o.Log, "restore: verified %d rows and %d series\n", res.RowsVerified, res.SeriesVerified)
	return res, os.RemoveAll(filepath.Join(o.DataDir, stagingDir))
}

func requireClean(dir string) error {
	ents, err := os.ReadDir(dir)
	if errors.Is(err, os.ErrNotExist) {
		return os.MkdirAll(dir, 0o700)
	}
	if err != nil {
		return err
	}
	if len(ents) > 0 {
		return fmt.Errorf("%w: %s", ErrNotEmpty, dir)
	}
	return nil
}

// emptyDir removes the children of dir (which Restore verified empty on entry).
func emptyDir(dir string) error {
	ents, err := os.ReadDir(dir)
	if err != nil {
		return err
	}
	var errs []error
	for _, e := range ents {
		errs = append(errs, os.RemoveAll(filepath.Join(dir, e.Name())))
	}
	return errors.Join(errs...)
}

// extract unpacks every entry into staging, hashing as it goes, then proves the archive is whole:
// the last entry is the manifest and every other entry matches it.
func extract(file, staging string) (*Manifest, error) {
	f, err := os.Open(file)
	if err != nil {
		return nil, err
	}
	defer func() { _ = f.Close() }() // read-only handle
	tr := tar.NewReader(f)
	seen := map[string]FileEntry{}
	for {
		hdr, err := tr.Next()
		if errors.Is(err, io.EOF) {
			return nil, fmt.Errorf("%w: no manifest", ErrTruncated)
		}
		if err != nil {
			return nil, fmt.Errorf("%w: %w", ErrTruncated, err)
		}
		if hdr.Name == manifestName {
			return finishExtract(tr, f, seen)
		}
		fe, err := extractEntry(tr, hdr, staging)
		if err != nil {
			return nil, err
		}
		seen[fe.Name] = fe
	}
}

func finishExtract(tr *tar.Reader, f *os.File, seen map[string]FileEntry) (*Manifest, error) {
	b, err := io.ReadAll(tr)
	if err != nil {
		return nil, fmt.Errorf("%w: manifest: %w", ErrTruncated, err)
	}
	var m Manifest
	if err := json.Unmarshal(b, &m); err != nil || m.Format != manifestFormat {
		return nil, fmt.Errorf("%w: unreadable manifest (format %d): %w", ErrTruncated, m.Format, err)
	}
	if _, err := tr.Next(); !errors.Is(err, io.EOF) {
		return nil, fmt.Errorf("%w: data after the manifest", ErrTruncated)
	}
	if err := requireTrailer(f); err != nil {
		return nil, err
	}
	if len(m.Files) != len(seen) {
		return nil, fmt.Errorf("%w: manifest lists %d files, archive has %d", ErrTruncated, len(m.Files), len(seen))
	}
	for _, want := range m.Files {
		if got, ok := seen[want.Name]; !ok || got != want {
			return nil, fmt.Errorf("%w: entry %s does not match the manifest", ErrTruncated, want.Name)
		}
	}
	return &m, nil
}

func extractEntry(tr *tar.Reader, hdr *tar.Header, staging string) (FileEntry, error) {
	name := path.Clean(hdr.Name)
	if hdr.Typeflag != tar.TypeReg || strings.HasPrefix(name, "../") || name == ".." || path.IsAbs(name) {
		return FileEntry{}, fmt.Errorf("%w: unsafe entry %q", ErrTruncated, hdr.Name)
	}
	dst := filepath.Join(staging, filepath.FromSlash(name))
	if err := os.MkdirAll(filepath.Dir(dst), 0o700); err != nil {
		return FileEntry{}, err
	}
	f, err := os.OpenFile(dst, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
	if err != nil {
		return FileEntry{}, err
	}
	h := sha256.New()
	n, err := io.Copy(io.MultiWriter(f, h), tr)
	if cerr := f.Close(); err == nil {
		err = cerr
	}
	if err != nil {
		return FileEntry{}, fmt.Errorf("%w: %s: %w", ErrTruncated, name, err)
	}
	return FileEntry{Name: name, Size: n, SHA256: hex.EncodeToString(h.Sum(nil))}, nil
}

// place moves the staged entries to their final homes.
func place(o *RestoreOptions, m *Manifest, res *RestoreResult) error {
	stg := filepath.Join(o.DataDir, stagingDir)
	moves := map[string]string{
		configEntry:                             config.File(o.DataDir),
		strings.TrimSuffix(profilesPrefix, "/"): config.ProfilesDir(o.DataDir),
		strings.TrimSuffix(ouiPrefix, "/"):      config.OUIDir(o.DataDir),
		"hot":                                   config.HotDir(o.DataDir),
		"long":                                  config.LongDir(o.DataDir),
	}
	if m.Inventory == inventoryDB {
		moves[inventoryDB] = config.InventoryDB(o.DataDir)
	}
	for src, dst := range moves {
		if _, err := os.Stat(filepath.Join(stg, src)); err != nil {
			continue
		}
		if err := os.Rename(filepath.Join(stg, src), dst); err != nil {
			return fmt.Errorf("restore: place %s: %w", src, err)
		}
	}
	if m.KeyIncluded {
		kf := o.KeyFile
		if kf == "" {
			kf = filepath.Join(o.DataDir, keyEntry)
		}
		if err := moveFile(filepath.Join(stg, keyEntry), kf); err != nil {
			return fmt.Errorf("restore: place key file: %w", err)
		}
		res.KeyRestoredTo = kf
	}
	return nil
}

// moveFile renames src to dst, copying when they are on different filesystems (the key file usually is).
func moveFile(src, dst string) error {
	if err := os.Rename(src, dst); err == nil {
		return nil
	}
	b, err := os.ReadFile(src)
	if err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Dir(dst), 0o700); err != nil {
		return err
	}
	return os.WriteFile(dst, b, 0o600)
}

// requireTrailer checks the archive ends with tar's two zero blocks, so a copy cut anywhere (even just
// after the manifest) is refused.
func requireTrailer(f *os.File) error {
	st, err := f.Stat()
	if err != nil {
		return err
	}
	if st.Size()%512 != 0 || st.Size() < 1024 {
		return fmt.Errorf("%w: size %d is not a whole number of tar blocks", ErrTruncated, st.Size())
	}
	tail := make([]byte, 1024)
	if _, err := f.ReadAt(tail, st.Size()-1024); err != nil {
		return fmt.Errorf("%w: %w", ErrTruncated, err)
	}
	for _, b := range tail {
		if b != 0 {
			return fmt.Errorf("%w: missing end-of-archive marker", ErrTruncated)
		}
	}
	return nil
}