Refuse mutating commands during a try and scope reverts to the try ID

apply, flush and purge take the try lock and fail with ErrPending, which now names 'tomswall revert' as the recovery after a failed automatic revert. Each try gets an ID passed to the timer's 'revert --id', so a stale timer cannot revert a newer try. Adds tests for Revert success/failure/stale ID and for restoring a present table at the engine level.
2026-10-03 20:56:52 +10:00
parent 6ac03e1012
commit 0b92f1c2f3
6 changed files with 306 additions and 46 deletions
@@ -3,6 +3,8 @@
 package tryapply

 import (
+	"crypto/rand"
+	"encoding/hex"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -30,12 +32,21 @@ var (
 		}
 		return nil
 	}
+	// restore rolls the live table back to a snapshot; replaced in tests.
+	restore = func(s *nftables.Snapshot) error {
+		engine, err := nftables.NewEngine(&config.Config{Settings: config.Settings{TableName: s.Table}})
+		if err != nil {
+			return err
+		}
+		return engine.Restore(s)
+	}
 )

 // ErrPending means a try awaits confirmation; nothing else may apply meanwhile.
-var ErrPending = errors.New("a 'tomswall try' is pending; run 'tomswall confirm' or 'tomswall revert'")
+var ErrPending = errors.New("a 'tomswall try' is pending; run 'tomswall confirm' to keep it or 'tomswall revert' to restore the previous ruleset (also the recovery if an automatic revert failed)")

 type pending struct {
+	ID       string             `json:"id"`
 	PID      int                `json:"pid"`
 	Snapshot *nftables.Snapshot `json:"snapshot"`
 }
@@ -70,43 +81,49 @@ func lock() (func(), error) {
 	return func() { f.Close() }, nil
 }

-// Arm persists snap and schedules an out-of-process revert after delay.
+// Arm persists snap and schedules an out-of-process revert after delay,
+// returning the try ID that scopes later reverts to this try.
 // The caller must hold the lock from Acquire.
-func Arm(snap *nftables.Snapshot, pid int, delay time.Duration) error {
-	b, err := json.Marshal(pending{PID: pid, Snapshot: snap})
+func Arm(snap *nftables.Snapshot, pid int, delay time.Duration) (string, error) {
+	raw := make([]byte, 8)
+	if _, err := rand.Read(raw); err != nil {
+		return "", err
+	}
+	id := hex.EncodeToString(raw)
+	b, err := json.Marshal(pending{ID: id, PID: pid, Snapshot: snap})
 	if err != nil {
-		return err
+		return "", err
 	}
 	f, err := os.CreateTemp(Dir, ".try-snapshot-*")
 	if err != nil {
-		return err
+		return "", err
 	}
 	defer os.Remove(f.Name())
 	if _, err := f.Write(b); err != nil {
 		f.Close()
-		return err
+		return "", err
 	}
 	if err := f.Sync(); err != nil {
 		f.Close()
-		return err
+		return "", err
 	}
 	if err := f.Close(); err != nil {
-		return err
+		return "", err
 	}
 	if err := os.Rename(f.Name(), snapshotPath()); err != nil {
-		return err
+		return "", err
 	}

 	exe, err := os.Executable()
 	if err != nil {
-		return discardWith(err)
+		return "", discardWith(err)
 	}
 	_ = disarm() // a leftover timer from an earlier try would block the unit name
 	if err := run("systemd-run", "--quiet", "--collect", "--unit", Unit,
-		fmt.Sprintf("--on-active=%ds", int(delay.Round(time.Second).Seconds())), exe, "revert"); err != nil {
-		return discardWith(fmt.Errorf("arming revert timer: %w", err))
+		fmt.Sprintf("--on-active=%ds", int(delay.Round(time.Second).Seconds())), exe, "revert", "--id", id); err != nil {
+		return "", discardWith(fmt.Errorf("arming revert timer: %w", err))
 	}
-	return nil
+	return id, nil
 }

 // Discard drops the pending snapshot and timer without restoring. The caller must hold the lock.
@@ -144,23 +161,21 @@ func Confirm() (pid int, ok bool, err error) {
 	return p.PID, true, Discard()
 }

-// Revert restores the pending snapshot. reverted is false when nothing was
-// pending (already confirmed or reverted).
-func Revert() (reverted bool, err error) {
+// Revert restores the pending snapshot. A non-empty id only reverts that try,
+// so a stale timer cannot undo a newer one. reverted is false when nothing
+// matching was pending (already confirmed or reverted). A failed restore keeps
+// the snapshot so 'tomswall revert' can retry.
+func Revert(id string) (reverted bool, err error) {
 	unlock, err := lock()
 	if err != nil {
 		return false, err
 	}
 	defer unlock()
 	p, err := load()
-	if err != nil || p == nil {
+	if err != nil || p == nil || (id != "" && p.ID != id) {
 		return false, err
 	}
-	engine, err := nftables.NewEngine(&config.Config{Settings: config.Settings{TableName: p.Snapshot.Table}})
-	if err != nil {
-		return false, err
-	}
-	if err := engine.Restore(p.Snapshot); err != nil {
+	if err := restore(p.Snapshot); err != nil {
 		return false, fmt.Errorf("restoring snapshot: %w", err)
 	}
 	return true, Discard()
@@ -24,23 +24,25 @@ func setup(t *testing.T) *[]string {
 	return &cmds
 }

-func arm(t *testing.T, snap *nftables.Snapshot) {
+func arm(t *testing.T, snap *nftables.Snapshot) string {
 	t.Helper()
 	unlock, err := Acquire()
 	if err != nil {
 		t.Fatal(err)
 	}
 	defer unlock()
-	if err := Arm(snap, 4242, 90*time.Second); err != nil {
+	id, err := Arm(snap, 4242, 90*time.Second)
+	if err != nil {
 		t.Fatal(err)
 	}
+	return id
 }

 func TestArmPersistsSnapshotAndTimer(t *testing.T) {
 	cmds := setup(t)
 	snap := &nftables.Snapshot{Table: "tomswall", Present: true,
 		Rules: map[string][]nftables.SnapshotRule{"input": {{Tag: "ssh", Exprs: [][]byte{{1, 2, 3}}}}}}
-	arm(t, snap)
+	id := arm(t, snap)

 	info, err := os.Stat(snapshotPath())
 	if err != nil {
@@ -53,12 +55,12 @@ func TestArmPersistsSnapshotAndTimer(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	if p.PID != 4242 || !reflect.DeepEqual(p.Snapshot, snap) {
+	if p.ID != id || id == "" || p.PID != 4242 || !reflect.DeepEqual(p.Snapshot, snap) {
 		t.Errorf("round trip mismatch: %+v", p)
 	}
 	last := (*cmds)[len(*cmds)-1]
 	if !strings.HasPrefix(last, "systemd-run ") || !strings.Contains(last, "--unit "+Unit) ||
-		!strings.Contains(last, "--on-active=90s") || !strings.HasSuffix(last, " revert") {
+		!strings.Contains(last, "--on-active=90s") || !strings.HasSuffix(last, " revert --id "+id) {
 		t.Errorf("unexpected arm command %q", last)
 	}
 }
@@ -78,9 +80,13 @@ func TestAbsentTableSnapshotRoundTrip(t *testing.T) {
 func TestAcquireRefusesWhilePending(t *testing.T) {
 	setup(t)
 	arm(t, &nftables.Snapshot{Table: "tomswall"})
-	if _, err := Acquire(); !errors.Is(err, ErrPending) {
+	_, err := Acquire()
+	if !errors.Is(err, ErrPending) {
 		t.Fatalf("second try: got %v, want ErrPending", err)
 	}
+	if !strings.Contains(err.Error(), "tomswall revert") {
+		t.Errorf("ErrPending does not name the recovery: %v", err)
+	}
 }

 func TestArmFailureDiscardsSnapshot(t *testing.T) {
@@ -96,7 +102,7 @@ func TestArmFailureDiscardsSnapshot(t *testing.T) {
 		t.Fatal(err)
 	}
 	defer unlock()
-	if err := Arm(&nftables.Snapshot{Table: "tomswall"}, 1, time.Minute); err == nil {
+	if _, err := Arm(&nftables.Snapshot{Table: "tomswall"}, 1, time.Minute); err == nil {
 		t.Fatal("expected arm error")
 	}
 	if _, err := os.Stat(snapshotPath()); !os.IsNotExist(err) {
@@ -130,8 +136,76 @@ func TestConfirmAfterRevertFails(t *testing.T) {
 	if err != nil || ok {
 		t.Fatalf("Confirm with nothing pending = %v, %v; want not ok", ok, err)
 	}
-	reverted, err := Revert()
+	reverted, err := Revert("")
 	if err != nil || reverted {
 		t.Fatalf("Revert with nothing pending = %v, %v", reverted, err)
 	}
 }
+
+func stubRestore(t *testing.T, err error) *[]*nftables.Snapshot {
+	t.Helper()
+	var got []*nftables.Snapshot
+	orig := restore
+	restore = func(s *nftables.Snapshot) error {
+		got = append(got, s)
+		return err
+	}
+	t.Cleanup(func() { restore = orig })
+	return &got
+}
+
+func TestRevertPendingRestoresAndDisarms(t *testing.T) {
+	cmds := setup(t)
+	restored := stubRestore(t, nil)
+	snap := &nftables.Snapshot{Table: "tomswall", Present: true,
+		Rules: map[string][]nftables.SnapshotRule{"input": {{Tag: "ssh", Exprs: [][]byte{{1}}}}}}
+	id := arm(t, snap)
+
+	reverted, err := Revert(id)
+	if err != nil || !reverted {
+		t.Fatalf("Revert = %v, %v", reverted, err)
+	}
+	if len(*restored) != 1 || !reflect.DeepEqual((*restored)[0], snap) {
+		t.Errorf("restored %+v, want the armed snapshot", *restored)
+	}
+	if _, err := os.Stat(snapshotPath()); !os.IsNotExist(err) {
+		t.Error("snapshot not removed")
+	}
+	if last := (*cmds)[len(*cmds)-1]; last != "systemctl stop "+Unit+".timer" {
+		t.Errorf("timer not stopped, last command %q", last)
+	}
+}
+
+func TestRevertFailureKeepsSnapshot(t *testing.T) {
+	setup(t)
+	boom := errors.New("netlink down")
+	stubRestore(t, boom)
+	id := arm(t, &nftables.Snapshot{Table: "tomswall"})
+
+	if _, err := Revert(id); !errors.Is(err, boom) {
+		t.Fatalf("Revert error = %v, want %v", err, boom)
+	}
+	if _, err := os.Stat(snapshotPath()); err != nil {
+		t.Fatalf("snapshot gone after failed revert: %v", err)
+	}
+	if _, err := Acquire(); !errors.Is(err, ErrPending) {
+		t.Errorf("failed revert must keep the try pending, got %v", err)
+	}
+}
+
+func TestRevertStaleIDIgnored(t *testing.T) {
+	setup(t)
+	restored := stubRestore(t, nil)
+	arm(t, &nftables.Snapshot{Table: "tomswall"})
+
+	reverted, err := Revert("stale-try")
+	if err != nil || reverted {
+		t.Fatalf("stale Revert = %v, %v; want no-op", reverted, err)
+	}
+	if len(*restored) != 0 {
+		t.Error("stale timer restored a newer try's snapshot")
+	}
+	if _, err := os.Stat(snapshotPath()); err != nil {
+		t.Errorf("newer try's snapshot removed: %v", err)
+	}
+}