felhom-agent/internal/selfupdate/prune_test.go

package selfupdate

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"testing"
)

// R-1368 part 2 — the self-update folder keeps what it needs and no more (prune.go).
//
// RED-PROOFS (each observed 2026-11-10, then restored):
//   - drop the pruneStaged call from MaybeCommit → TestPrune_AfterNUpdatesTheFolderHoldsOneBinary fails with
//     "after update 1 (0.170.0) the folder holds [… felhom-agent-0.169.0 felhom-agent-0.170.0 …]", and
//     TestPrune_RealRollbackStillWorksAfterNUpdates fails too;
//   - drop the pending-marker refusal from pruneStaged → TestPrune_RefusesWhileAnUpdateIsPending fails with
//     "removed [...] while an update was pending";
//   - drop the IsRegular check → TestPrune_AfterNUpdatesTheFolderHoldsOneBinary fails with "a symlink with a
//     staged-binary name was removed".

// committingWrapper is the guarded wrapper's `commit` as far as this package can see it: it clears the pending marker.
type committingWrapper struct {
	stateDir string
	calls    int
}

func (w *committingWrapper) Run(_ context.Context, _ string, args ...string) ([]byte, []byte, error) {
	if len(args) == 1 && args[0] == "commit" {
		w.calls++
		if err := os.Remove(pendingPath(w.stateDir)); err != nil {
			return nil, []byte(err.Error()), err
		}
	}
	return []byte("ok"), nil, nil
}

func stage(t *testing.T, stateDir, version string) {
	t.Helper()
	if err := os.MkdirAll(stagingDir(stateDir), 0o750); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(stagingDir(stateDir), "felhom-agent-"+version), []byte("binary "+version), 0o755); err != nil {
		t.Fatal(err)
	}
}

func folder(t *testing.T, stateDir string) []string {
	t.Helper()
	ents, err := os.ReadDir(stagingDir(stateDir))
	if err != nil {
		t.Fatal(err)
	}
	var out []string
	for _, e := range ents {
		out = append(out, e.Name())
	}
	sort.Strings(out)
	return out
}

// The consequence: after N good updates the folder holds ONE staged binary — the running version's — and every file
// that is not a staged binary is still there.
func TestPrune_AfterNUpdatesTheFolderHoldsOneBinary(t *testing.T) {
	stateDir := t.TempDir()
	dir := stagingDir(stateDir)
	if err := os.MkdirAll(dir, 0o750); err != nil {
		t.Fatal(err)
	}
	// Bystanders: names this code does not know, and two things that carry a staged binary's NAME but are not files.
	outside := filepath.Join(t.TempDir(), "precious")
	if err := os.WriteFile(outside, []byte("not ours"), 0o644); err != nil {
		t.Fatal(err)
	}
	for _, n := range []string{"notes.txt", "felhom-agent-0.1.0.partial", "felhom-agent.bak-0.139.0"} {
		if err := os.WriteFile(filepath.Join(dir, n), []byte("x"), 0o644); err != nil {
			t.Fatal(err)
		}
	}
	if err := os.Mkdir(filepath.Join(dir, "felhom-agent-9.9.9"), 0o750); err != nil {
		t.Fatal(err)
	}
	if err := os.Symlink(outside, filepath.Join(dir, "felhom-agent-8.8.8")); err != nil {
		t.Fatal(err)
	}
	bystanders := []string{"felhom-agent-0.1.0.partial", "felhom-agent-8.8.8", "felhom-agent-9.9.9", "felhom-agent.bak-0.139.0", "notes.txt"}

	versions := []string{"0.170.0", "0.171.0", "0.172.0", "0.172.1", "0.173.0-rc.1", "0.173.0"}
	prev := "0.169.0"
	stage(t, stateDir, prev) // what an old agent left behind
	for i, v := range versions {
		stage(t, stateDir, v)                                                                                   // Execute downloaded it
		writePending(t, stateDir, PendingMarker{OldVersion: prev, NewVersion: v, SHA256: "ab", AppliedAt: "t"}) // the wrapper flipped
		wrap := &committingWrapper{stateDir: stateDir}
		m := newMgr(t, stateDir, v, wrap, nil)
		m.MaybeCommit(context.Background())
		if wrap.calls != 1 {
			t.Fatalf("update %d (%s): commit calls = %d, want 1", i+1, v, wrap.calls)
		}
		want := append([]string{"felhom-agent-" + v}, bystanders...)
		sort.Strings(want)
		got := folder(t, stateDir)
		if strings.Join(got, " ") != strings.Join(want, " ") {
			if _, err := os.Lstat(filepath.Join(dir, "felhom-agent-8.8.8")); err != nil {
				t.Fatalf("a symlink with a staged-binary name was removed: %v", err)
			}
			t.Fatalf("after update %d (%s) the folder holds %v, want %v", i+1, v, got, want)
		}
		prev = v
	}
	if b, err := os.ReadFile(outside); err != nil || string(b) != "not ours" {
		t.Fatalf("the file behind the symlink changed: %q %v", b, err)
	}
}

// While an update is pending — the only time a rollback can fire — nothing is removed.
func TestPrune_RefusesWhileAnUpdateIsPending(t *testing.T) {
	stateDir := t.TempDir()
	for _, v := range []string{"0.170.0", "0.171.0", "0.172.0"} {
		stage(t, stateDir, v)
	}
	writePending(t, stateDir, PendingMarker{OldVersion: "0.171.0", NewVersion: "0.172.0", SHA256: "ab", AppliedAt: "t"})
	before := folder(t, stateDir)

	removed, err := pruneStaged(stateDir, "0.172.0", nil)
	if err == nil || len(removed) != 0 {
		t.Fatalf("removed %v while an update was pending (err=%v) — want a refusal", removed, err)
	}
	if got := folder(t, stateDir); strings.Join(got, " ") != strings.Join(before, " ") {
		t.Fatalf("removed %v while an update was pending; folder now %v", before, got)
	}

	// The same through the commit path: a `commit` that reports success but leaves the marker (fakeWrapper does) must
	// not lead to a prune.
	m := newMgr(t, stateDir, "0.172.0", &fakeWrapper{}, nil)
	m.MaybeCommit(context.Background())
	if got := folder(t, stateDir); strings.Join(got, " ") != strings.Join(before, " ") {
		t.Fatalf("the commit path pruned with the marker still there; folder now %v", got)
	}

	// A marker nobody can read is a refusal too.
	if err := os.WriteFile(pendingPath(stateDir), []byte("{not json"), 0o644); err != nil {
		t.Fatal(err)
	}
	if removed, err := pruneStaged(stateDir, "0.172.0", nil); err == nil || len(removed) != 0 {
		t.Fatalf("an unreadable marker must refuse; removed %v err=%v", removed, err)
	}
}

// Only a commit prunes: a plain start, a version mismatch and a failed commit remove nothing — so the binary of a
// version that was rolled back stays for the operator.
func TestPrune_OnlyAHealthyCommitPrunes(t *testing.T) {
	cases := []struct {
		name    string
		pending *PendingMarker
		running string
		wrap    WrapperRunner
	}{
		{"plain start, no marker (also the state after a rollback)", nil, "0.171.0", &fakeWrapper{}},
		{"marker names another version", &PendingMarker{OldVersion: "0.171.0", NewVersion: "0.172.0"}, "0.171.0", &fakeWrapper{}},
		{"commit fails", &PendingMarker{OldVersion: "0.171.0", NewVersion: "0.172.0"}, "0.172.0", &fakeWrapper{err: fmt.Errorf("sudo: no")}},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			stateDir := t.TempDir()
			for _, v := range []string{"0.170.0", "0.171.0", "0.172.0"} {
				stage(t, stateDir, v)
			}
			if c.pending != nil {
				writePending(t, stateDir, *c.pending)
			}
			before := folder(t, stateDir)
			newMgr(t, stateDir, c.running, c.wrap, nil).MaybeCommit(context.Background())
			if got := folder(t, stateDir); strings.Join(got, " ") != strings.Join(before, " ") {
				t.Fatalf("folder changed without a healthy commit: %v → %v", before, got)
			}
		})
	}
	// A running version this code cannot name a file for: refuse, do not guess.
	stateDir := t.TempDir()
	stage(t, stateDir, "0.170.0")
	if removed, err := pruneStaged(stateDir, "dev", nil); err == nil || len(removed) != 0 {
		t.Fatalf("running version \"dev\" must refuse; removed %v err=%v", removed, err)
	}
}

// The name rule and the version rule say the same thing: every name Execute can write is a name the prune knows, and
// nothing else is.
func TestPrune_NameShapeFollowsTheVersionRule(t *testing.T) {
	for _, v := range []string{"0.70.1", "0.172.1", "10.0.0", "0.70.2-crash", "0.173.0-rc.1"} {
		if !versionRe.MatchString(v) {
			t.Fatalf("test input %q is not a version versionRe accepts", v)
		}
		if !stagedNameRe.MatchString("felhom-agent-" + v) {
			t.Errorf("Execute can stage felhom-agent-%s but the prune does not know that name", v)
		}
	}
	for _, n := range []string{"pending.json", "felhom-agent", "felhom-agent-", "felhom-agent-1.2", "felhom-agent-0.1.0.partial",
		"felhom-agent.bak-0.139.0", "felhom-agent.prev", "felhom-agent-0.1.0/x", "../felhom-agent-0.1.0", "felhom-agent-0.1.0 ", "xfelhom-agent-0.1.0"} {
		if stagedNameRe.MatchString(n) {
			t.Errorf("the prune would treat %q as a staged binary", n)
		}
	}
	m := stagedNameRe.FindStringSubmatch("felhom-agent-0.173.0-rc.1")
	if m == nil || !versionRe.MatchString(m[1]) {
		t.Fatalf("the version part of a staged name must satisfy versionRe; got %v", m)
	}
}

func wrapperText(t *testing.T) string {
	t.Helper()
	b, err := os.ReadFile(filepath.Join("..", "..", "configs", "felhom-selfupdate-guarded"))
	if err != nil {
		t.Fatalf("the guarded wrapper is the other half of this proof and must be readable: %v", err)
	}
	return string(b)
}

// The stated dependency, pinned: the wrapper's `rollback` and `commit` read the pending marker and `.prev`, and no
// staged binary. If either verb starts to read the staging folder, this fails and prune.go's rule must be looked at.
func TestPrune_RollbackReadsNoStagedBinary(t *testing.T) {
	s := wrapperText(t)
	for _, want := range []string{"\nBIN=/usr/local/bin/felhom-agent\n", "\nPREV=$BIN.prev\n", "\nSTAGING=/var/lib/felhom-agent/selfupdate\n", "\nPENDING=$STAGING/pending.json\n"} {
		if !strings.Contains(s, want) {
			t.Fatalf("the wrapper no longer holds the line %q — the folder's rule (prune.go) was written against it", strings.TrimSpace(want))
		}
	}
	for _, verb := range []string{"commit", "rollback"} {
		i := strings.Index(s, "\n"+verb+")\n")
		if i < 0 {
			t.Fatalf("no %s) block in the wrapper", verb)
		}
		block := s[i:]
		j := strings.Index(block, "\n    ;;\n")
		if j < 0 {
			t.Fatalf("no end of the %s) block", verb)
		}
		block = block[:j]
		if !strings.Contains(block, "$PENDING") {
			t.Fatalf("%s) does not read $PENDING any more", verb)
		}
		code := regexp.MustCompile(`(?m)^\s*#.*$`).ReplaceAllString(block, "")
		for _, bad := range []string{"STAGING", "staged", "felhom-agent-"} {
			if strings.Contains(code, bad) {
				t.Fatalf("%s) now mentions %q — it may read a staged binary, and the prune removes those", verb, bad)
			}
		}
	}
	if stagedNameRe.MatchString("pending.json") || stagedNameRe.MatchString("felhom-agent.prev") {
		t.Fatal("the prune would remove a file the rollback needs")
	}
}

// The consequence, with the wrapper's own code: after N good updates (each committed by the REAL `commit` verb and
// pruned), update N+1 goes bad and the REAL `rollback` verb puts the previous binary back. The wrapper runs from a copy
// whose four fixed paths point into a temp folder; systemctl, logger and chown are stubs.
func TestPrune_RealRollbackStillWorksAfterNUpdates(t *testing.T) {
	sh, err := exec.LookPath("sh")
	if err != nil {
		t.Skip("no sh on this machine")
	}
	root := t.TempDir()
	stateDir := filepath.Join(root, "state")
	binDir := filepath.Join(root, "bin")
	stubs := filepath.Join(root, "stubs")
	for _, d := range []string{stagingDir(stateDir), binDir, stubs, filepath.Join(root, "rootstaging")} {
		if err := os.MkdirAll(d, 0o750); err != nil {
			t.Fatal(err)
		}
	}
	live := filepath.Join(binDir, "felhom-agent")
	s := wrapperText(t)
	for old, repl := range map[string]string{
		"\nBIN=/usr/local/bin/felhom-agent\n":                    "\nBIN=" + live + "\n",
		"\nSTAGING=/var/lib/felhom-agent/selfupdate\n":           "\nSTAGING=" + stagingDir(stateDir) + "\n",
		"\nROOT_STAGING=/var/lib/felhom-os-apply/agent-update\n": "\nROOT_STAGING=" + filepath.Join(root, "rootstaging") + "\n",
	} {
		if strings.Count(s, old) != 1 {
			t.Fatalf("the wrapper's path line %q is not there exactly once — this test cannot sandbox it", strings.TrimSpace(old))
		}
		s = strings.Replace(s, old, repl, 1)
	}
	// Safety: no assignment in the copy may still point at a real system path (PREV and PENDING are built from BIN and
	// STAGING, so they follow).
	for _, line := range strings.Split(s, "\n") {
		if regexp.MustCompile(`^[A-Z_]+=/`).MatchString(line) && !strings.Contains(line, root) {
			t.Fatalf("the sandboxed wrapper still holds a real path: %q", line)
		}
	}
	script := filepath.Join(root, "wrapper.sh")
	if err := os.WriteFile(script, []byte(s), 0o755); err != nil {
		t.Fatal(err)
	}
	stubLog := filepath.Join(root, "stub.log")
	for _, n := range []string{"systemctl", "logger", "chown", "systemd-run"} {
		body := "#!/bin/sh\necho \"" + n + " $*\" >> '" + stubLog + "'\nexit 0\n"
		if err := os.WriteFile(filepath.Join(stubs, n), []byte(body), 0o755); err != nil {
			t.Fatal(err)
		}
	}
	run := func(verb string) (string, error) {
		cmd := exec.Command(sh, script, verb)
		cmd.Env = []string{"PATH=" + stubs + ":/usr/bin:/bin"}
		out, err := cmd.CombinedOutput()
		return string(out), err
	}
	// applied is what the wrapper's `apply` leaves behind (it needs root, so the test writes the same four things).
	applied := func(oldV, newV string) {
		cur, err := os.ReadFile(live)
		if err != nil {
			t.Fatal(err)
		}
		if err := os.WriteFile(live+".prev", cur, 0o755); err != nil {
			t.Fatal(err)
		}
		if err := os.WriteFile(live, []byte("binary "+newV), 0o755); err != nil {
			t.Fatal(err)
		}
		writePending(t, stateDir, PendingMarker{OldVersion: oldV, NewVersion: newV, SHA256: "ab", AppliedAt: "t"})
	}
	if err := os.WriteFile(live, []byte("binary 0.169.0"), 0o755); err != nil {
		t.Fatal(err)
	}

	prev := "0.169.0"
	for _, v := range []string{"0.170.0", "0.171.0", "0.172.0", "0.172.1"} {
		stage(t, stateDir, v)
		applied(prev, v)
		m := newMgr(t, stateDir, v, runnerFunc(func(args ...string) ([]byte, []byte, error) {
			out, err := run(args[0])
			return []byte(out), nil, err
		}), nil)
		m.MaybeCommit(context.Background()) // the REAL commit verb, then the prune
		if got := folder(t, stateDir); len(got) != 1 || got[0] != "felhom-agent-"+v {
			t.Fatalf("after the good update to %s the folder holds %v", v, got)
		}
		prev = v
	}

	// Update N+1 is applied and goes bad: systemd's OnFailure runs the wrapper's rollback.
	stage(t, stateDir, "0.173.0")
	applied("0.172.1", "0.173.0")
	out, err := run("rollback")
	if err != nil {
		t.Fatalf("the real rollback verb failed after the prunes: %v\n%s", err, out)
	}
	if b, _ := os.ReadFile(live); string(b) != "binary 0.172.1" {
		t.Fatalf("after the rollback the live binary is %q, want the previous version's", b)
	}
	if p, _ := readPending(stateDir); p != nil {
		t.Fatal("the rollback left the pending marker")
	}
	if b, _ := os.ReadFile(stubLog); !bytes.Contains(b, []byte("systemctl restart felhom-agent.service")) {
		t.Fatalf("the rollback did not restart the agent; stub log:\n%s", b)
	}
	// The rolled-back agent starts: no marker, so no commit and no prune — the bad version's binary stays to be read.
	newMgr(t, stateDir, "0.172.1", &fakeWrapper{}, nil).MaybeCommit(context.Background())
	if got := folder(t, stateDir); strings.Join(got, " ") != "felhom-agent-0.172.1 felhom-agent-0.173.0" {
		t.Fatalf("after the rollback the folder holds %v, want the running version and the one that failed", got)
	}
}

// runnerFunc adapts a function to WrapperRunner; it gets the verb arguments only.
type runnerFunc func(args ...string) ([]byte, []byte, error)

func (f runnerFunc) Run(_ context.Context, _ string, args ...string) ([]byte, []byte, error) {
	return f(args...)
}