Maven/internal/update/snapshot.go

package update

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"time"
)

// A snapshot is a byte-for-byte copy of the deployed artifacts plus a manifest
// of their sha256 sums, taken before an install.
//
// It is copies, not hardlinks and not a git stash, for one reason: the restore
// path must work when everything else is broken. A hardlink into the install dir
// would be clobbered by the very install it exists to undo, and a git-based
// undo needs a toolchain, a clean tree, and a rebuild — three things a failed
// update is likely to have taken away. Copying two dozen megabytes of Go
// binaries costs a second and needs nothing but the filesystem.
//
// The sums are what make a restore verifiable rather than hopeful: Restore
// re-hashes every file it writes, so "the old bytes are back" is checked, not
// assumed.

// FileRec — one file in a snapshot.
type FileRec struct {
	Name   string      `json:"name"`   // relative name inside the install dir
	SHA256 string      `json:"sha256"` // of the snapshotted bytes
	Mode   os.FileMode `json:"mode"`
	Size   int64       `json:"size"`
}

// Snapshot — the manifest. Written last, so a directory without a readable
// manifest.json is an aborted snapshot and is never offered as a rollback target.
type Snapshot struct {
	ID        string    `json:"id"` // sortable timestamp, also the directory name
	CreatedAt time.Time `json:"created_at"`
	Commit    string    `json:"commit,omitempty"` // git HEAD of the tree that produced it, when known
	Note      string    `json:"note,omitempty"`
	Files     []FileRec `json:"files"`

	dir string // absolute path, filled in by List/Load
}

// Dir — where this snapshot's file copies live.
func (s Snapshot) Dir() string { return s.dir }

const manifestName = "manifest.json"

// Store is a directory of snapshots.
type Store struct {
	Dir string
	now func() time.Time
}

func (st *Store) clock() time.Time {
	if st.now != nil {
		return st.now()
	}
	return time.Now()
}

// Save copies names (relative to srcDir) into a new snapshot and writes the
// manifest. A name that does not exist is skipped rather than fatal: the first
// ever run happens on a box where some artifact may legitimately be missing, and
// refusing to snapshot then would mean refusing to update.
func (st *Store) Save(srcDir string, names []string, commit, note string) (Snapshot, error) {
	ts := st.clock().UTC()
	snap := Snapshot{
		ID:        ts.Format("20060102-150405"),
		CreatedAt: ts,
		Commit:    commit,
		Note:      note,
	}
	snap.dir = filepath.Join(st.Dir, snap.ID)
	if err := os.MkdirAll(snap.dir, 0o700); err != nil {
		return Snapshot{}, fmt.Errorf("update: snapshot dir: %w", err)
	}
	for _, name := range names {
		src := filepath.Join(srcDir, name)
		fi, err := os.Stat(src)
		if err != nil {
			if errors.Is(err, os.ErrNotExist) {
				continue
			}
			return Snapshot{}, fmt.Errorf("update: snapshot %s: %w", name, err)
		}
		if fi.IsDir() {
			return Snapshot{}, fmt.Errorf("update: snapshot %s: is a directory (only files are deployable artifacts)", name)
		}
		dst := filepath.Join(snap.dir, name)
		if err := os.MkdirAll(filepath.Dir(dst), 0o700); err != nil {
			return Snapshot{}, err
		}
		sum, err := copyFile(src, dst, fi.Mode().Perm())
		if err != nil {
			return Snapshot{}, fmt.Errorf("update: snapshot %s: %w", name, err)
		}
		snap.Files = append(snap.Files, FileRec{Name: name, SHA256: sum, Mode: fi.Mode().Perm(), Size: fi.Size()})
	}
	if len(snap.Files) == 0 {
		os.RemoveAll(snap.dir)
		return Snapshot{}, fmt.Errorf("update: snapshot of %s is empty — none of the listed artifacts exist", srcDir)
	}
	// Manifest last: its presence is what makes the snapshot usable.
	blob, err := json.MarshalIndent(snap, "", "  ")
	if err != nil {
		return Snapshot{}, err
	}
	if err := os.WriteFile(filepath.Join(snap.dir, manifestName), blob, 0o600); err != nil {
		return Snapshot{}, fmt.Errorf("update: snapshot manifest: %w", err)
	}
	return snap, nil
}

// List returns the complete snapshots, newest first.
func (st *Store) List() ([]Snapshot, error) {
	ents, err := os.ReadDir(st.Dir)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return nil, nil
		}
		return nil, err
	}
	var out []Snapshot
	for _, e := range ents {
		if !e.IsDir() {
			continue
		}
		s, err := st.Load(e.Name())
		if err != nil {
			continue // aborted or hand-mangled: not a rollback target
		}
		out = append(out, s)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].ID > out[j].ID })
	return out, nil
}

// Load reads one snapshot's manifest.
func (st *Store) Load(id string) (Snapshot, error) {
	dir := filepath.Join(st.Dir, id)
	blob, err := os.ReadFile(filepath.Join(dir, manifestName))
	if err != nil {
		return Snapshot{}, err
	}
	var s Snapshot
	if err := json.Unmarshal(blob, &s); err != nil {
		return Snapshot{}, fmt.Errorf("update: manifest %s: %w", id, err)
	}
	s.dir = dir
	return s, nil
}

// Restore copies a snapshot's files back over dstDir and verifies every write
// against the manifest sum. Only the named files are touched; anything else in
// dstDir is left alone.
//
// This is the function the whole package exists to be able to run. It uses the
// filesystem and nothing else — no toolchain, no build, no cooperation from the
// code being replaced.
func (s Snapshot) Restore(dstDir string) error { return s.RestoreOnly(dstDir, nil) }

// RestoreOnly is Restore limited to the named files. A nil list means all of
// them. The caller uses it to put binaries back without putting config back:
// see Updater.restoreBinaries.
func (s Snapshot) RestoreOnly(dstDir string, names []string) error {
	if s.dir == "" {
		return errors.New("update: snapshot has no directory (load it through the store)")
	}
	var want map[string]bool
	if names != nil {
		want = make(map[string]bool, len(names))
		for _, n := range names {
			want[n] = true
		}
	}
	for _, f := range s.Files {
		if want != nil && !want[f.Name] {
			continue
		}
		src := filepath.Join(s.dir, f.Name)
		sum, err := hashFile(src)
		if err != nil {
			return fmt.Errorf("update: restore %s: %w", f.Name, err)
		}
		if sum != f.SHA256 {
			return fmt.Errorf("update: restore %s: snapshot is corrupt (sha256 %s, manifest says %s)", f.Name, sum, f.SHA256)
		}
		dst := filepath.Join(dstDir, f.Name)
		if err := os.MkdirAll(filepath.Dir(dst), 0o755); err != nil {
			return err
		}
		got, err := copyFile(src, dst, f.Mode)
		if err != nil {
			return fmt.Errorf("update: restore %s: %w", f.Name, err)
		}
		if got != f.SHA256 {
			return fmt.Errorf("update: restore %s: wrote the wrong bytes (sha256 %s)", f.Name, got)
		}
	}
	return nil
}

// Prune keeps the newest keep snapshots and removes the rest. The newest is
// never pruned regardless of keep — it is the rollback target.
func (st *Store) Prune(keep int) error {
	if keep < 1 {
		keep = 1
	}
	snaps, err := st.List()
	if err != nil {
		return err
	}
	for _, s := range snaps[min(keep, len(snaps)):] {
		if err := os.RemoveAll(s.dir); err != nil {
			return err
		}
	}
	return nil
}

// copyFile writes src to dst atomically (temp + rename, so a reader never sees a
// half file and an interrupted copy leaves the old one intact) and returns the
// sha256 of what was written.
func copyFile(src, dst string, mode os.FileMode) (string, error) {
	in, err := os.Open(src)
	if err != nil {
		return "", err
	}
	defer in.Close()
	if mode == 0 {
		mode = 0o644
	}
	tmp, err := os.CreateTemp(filepath.Dir(dst), ".update-*")
	if err != nil {
		return "", err
	}
	tmpName := tmp.Name()
	defer os.Remove(tmpName) // no-op once the rename succeeds
	h := sha256.New()
	if _, err := io.Copy(io.MultiWriter(tmp, h), in); err != nil {
		tmp.Close()
		return "", err
	}
	// fsync before the rename: a binary that is renamed into place but whose
	// bytes are still in the page cache is exactly the file a power cut turns
	// into an unbootable daemon.
	if err := tmp.Sync(); err != nil {
		tmp.Close()
		return "", err
	}
	if err := tmp.Chmod(mode); err != nil {
		tmp.Close()
		return "", err
	}
	if err := tmp.Close(); err != nil {
		return "", err
	}
	if err := os.Rename(tmpName, dst); err != nil {
		return "", err
	}
	return hex.EncodeToString(h.Sum(nil)), nil
}

func hashFile(p string) (string, error) {
	f, err := os.Open(p)
	if err != nil {
		return "", err
	}
	defer f.Close()
	h := sha256.New()
	if _, err := io.Copy(h, f); err != nil {
		return "", err
	}
	return hex.EncodeToString(h.Sum(nil)), nil
}