prx/internal/helper/helper_test.go

// SPDX-License-Identifier: MIT

package helper

import (
	"context"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"prx/internal/execx"
	"prx/internal/proto"
)

func basePlan(t *testing.T) proto.Plan {
	t.Helper()
	return proto.Plan{
		Action:        proto.ActionApply,
		Mark:          0x100,
		TableV4:       100,
		TableV6:       101,
		PriorityV4:    100,
		PriorityV6:    100,
		BypassTableV4: 102,
		BypassTableV6: 103,
		BypassPrioV4:  50,
		BypassPrioV6:  50,
		ManageEngine:  false,
		StatePath:     filepath.Join(t.TempDir(), "state.json"),
	}
}

// systemFake emulates a host with a working WAN and optionally active engine.
func systemFake(active bool) *execx.Fake {
	return &execx.Fake{Handler: func(name string, args []string, _ []byte) (execx.Result, error) {
		joined := strings.TrimSpace(name + " " + strings.Join(args, " "))
		switch {
		case joined == "ip -j route show default":
			return execx.Result{Stdout: []byte(`[{"dst":"default","gateway":"192.168.1.1","dev":"wlan0","metric":600}]`)}, nil
		case joined == "ip -6 -j route show default":
			return execx.Result{Stdout: []byte(`[]`)}, nil
		case joined == "ip -j addr show dev wlan0":
			return execx.Result{Stdout: []byte(`[{"ifname":"wlan0","addr_info":[{"family":"inet","local":"192.168.1.100","prefixlen":24,"scope":"global"}]}]`)}, nil
		case joined == "ip rule show" || joined == "ip -6 rule show":
			return execx.Result{Stdout: []byte("100:\tfrom all fwmark 0x100 lookup 100\n")}, nil
		case strings.HasPrefix(joined, "ip route show table"):
			return execx.Result{Stdout: []byte("default via 192.168.1.1 dev wlan0\n")}, nil
		case joined == "ip -j link show type tun":
			return execx.Result{Stdout: []byte(`[{"ifname":"tun0"}]`)}, nil
		case strings.HasPrefix(joined, "nft"):
			return execx.Result{}, nil
		case strings.HasPrefix(joined, "systemctl is-active"):
			if active {
				return execx.Result{Stdout: []byte("active\n")}, nil
			}
			return execx.Result{Stdout: []byte("inactive\n"), Code: 3}, &execx.ExitError{Name: "systemctl", Code: 3, Stderr: []byte("inactive")}
		case strings.HasPrefix(joined, "systemctl start"):
			return execx.Result{}, nil
		case strings.HasPrefix(joined, "systemctl stop"):
			return execx.Result{}, nil
		}
		return execx.Result{}, nil
	}}
}

func TestValidatePlan(t *testing.T) {
	cases := map[string]func(*proto.Plan){
		"zero mark":      func(p *proto.Plan) { p.Mark = 0 },
		"reserved table": func(p *proto.Plan) { p.TableV4 = 254 },
		"bad priority":   func(p *proto.Plan) { p.PriorityV4 = 99999 },
		"bad cidr":       func(p *proto.Plan) { p.BypassCIDRs = []string{"nope"} },
		"relative state": func(p *proto.Plan) { p.StatePath = "state.json" },
		"bad iface":      func(p *proto.Plan) { p.WANInterface = "bad iface!" },
		"bad gateway":    func(p *proto.Plan) { p.WANGateway4 = "not-ip" },
	}
	for name, mutate := range cases {
		p := basePlan(t)
		mutate(&p)
		if err := validatePlan(p); err == nil {
			t.Errorf("case %q: expected error", name)
		}
	}
	if err := validatePlan(basePlan(t)); err != nil {
		t.Fatalf("valid plan rejected: %v", err)
	}
}

func TestApplyDryRunChangesNothing(t *testing.T) {
	p := basePlan(t)
	p.DryRun = true
	res := Execute(context.Background(), systemFake(false), p)
	if !res.OK {
		t.Fatalf("dry run failed: %+v", res)
	}
	if res.Changed {
		t.Fatal("dry run must not report changes")
	}
}

func TestApplySkipEngineSucceeds(t *testing.T) {
	p := basePlan(t)
	p.SkipEngine = true
	f := systemFake(false)
	res := Execute(context.Background(), f, p)
	if !res.OK {
		t.Fatalf("apply failed: %+v", res)
	}
	if _, err := filepath.Glob(p.StatePath); err != nil {
		t.Fatal(err)
	}
	if !hasCall(f.Calls, "nft", "-f", "-") {
		t.Fatal("expected nft apply call")
	}
}

func TestApplyRollsBackOnHealthFailure(t *testing.T) {
	p := basePlan(t)
	p.ManageEngine = true
	p.ServiceName = "sing-box"
	p.SkipEngine = false
	f := systemFake(false) // engine never becomes active
	res := Execute(context.Background(), f, p)
	if res.OK {
		t.Fatal("expected apply to fail")
	}
	if !strings.Contains(res.Error, "rolled back") {
		t.Fatalf("expected rollback message, got %q", res.Error)
	}
	if !hasCall(f.Calls, "nft", "destroy", "table", "inet", "prx") {
		t.Fatal("expected rollback to destroy the nft table")
	}
}

func TestRemoveIsIdempotent(t *testing.T) {
	p := basePlan(t)
	p.Action = proto.ActionRemove
	p.StatePath = filepath.Join(t.TempDir(), "state.json")
	f := &execx.Fake{Handler: func(name string, args []string, _ []byte) (execx.Result, error) {
		if strings.HasPrefix(strings.TrimSpace(name+" "+strings.Join(args, " ")), "nft") {
			return execx.Result{}, nil
		}
		return execx.Result{Stdout: []byte("")}, nil
	}}
	for i := 0; i < 2; i++ {
		res := Execute(context.Background(), f, p)
		if !res.OK {
			t.Fatalf("remove iteration %d failed: %+v", i, res)
		}
	}
}

func TestInspectReturnsChecks(t *testing.T) {
	p := basePlan(t)
	p.Action = proto.ActionInspect
	res := Execute(context.Background(), systemFake(true), p)
	if !res.OK {
		t.Fatalf("inspect failed: %+v", res)
	}
	if len(res.Checks) == 0 {
		t.Fatal("expected checks")
	}
}

func TestInitEngineRejectsBadJSON(t *testing.T) {
	p := proto.Plan{
		Action:           proto.ActionInitEngine,
		EngineConfig:     filepath.Join(t.TempDir(), "config.json"),
		EngineConfigBody: []byte("not json"),
	}
	res := Execute(context.Background(), systemFake(false), p)
	if res.OK {
		t.Fatal("expected rejection of invalid JSON")
	}
}

func hasCall(calls []execx.Call, name string, args ...string) bool {
	for _, c := range calls {
		if c.Name != name || len(c.Args) != len(args) {
			continue
		}
		match := true
		for i := range args {
			if c.Args[i] != args[i] {
				match = false
				break
			}
		}
		if match {
			return true
		}
	}
	return false
}

// TestApplyRollsBackPartialRoutingFailure verifies that a failure after the
// policy rule was added still cleans up the routing table. The rollback for a
// table must be registered before the table is touched.
func TestApplyRollsBackPartialRoutingFailure(t *testing.T) {
	p := basePlan(t)
	p.SkipEngine = true

	f := &execx.Fake{Handler: func(name string, args []string, _ []byte) (execx.Result, error) {
		joined := strings.TrimSpace(name + " " + strings.Join(args, " "))
		switch {
		case joined == "ip -j route show default":
			return execx.Result{Stdout: []byte(`[{"dst":"default","gateway":"192.168.1.1","dev":"wlan0","metric":600}]`)}, nil
		case joined == "ip -6 -j route show default":
			return execx.Result{Stdout: []byte(`[]`)}, nil
		case joined == "ip -j addr show dev wlan0":
			return execx.Result{Stdout: []byte(`[{"ifname":"wlan0","addr_info":[{"family":"inet","local":"192.168.1.100","prefixlen":24,"scope":"global"}]}]`)}, nil
		case joined == "ip rule show":
			// No rule present yet, so EnsureRule will add one.
			return execx.Result{Stdout: []byte("")}, nil
		case joined == "ip route show table 100":
			// RouteExists cannot read the table: force EnsureRoute to fail.
			return execx.Result{Code: 2, Stderr: []byte("RTNETLINK answers: Operation not permitted")}, &execx.ExitError{Name: "ip", Code: 2, Stderr: []byte("RTNETLINK answers: Operation not permitted")}
		case strings.HasPrefix(joined, "nft"):
			return execx.Result{}, nil
		}
		return execx.Result{}, nil
	}}

	res := Execute(context.Background(), f, p)
	if res.OK {
		t.Fatalf("expected apply to fail: %+v", res)
	}
	if !hasCall(f.Calls, "ip", "rule", "add", "fwmark", "0x100", "lookup", "100", "priority", "100") {
		t.Fatal("expected the policy rule to have been added")
	}
	if !hasCall(f.Calls, "ip", "route", "flush", "table", "100") {
		t.Fatal("partial failure leaked routing table 100 (rollback not registered early enough)")
	}
	if !hasCall(f.Calls, "nft", "destroy", "table", "inet", "prx") {
		t.Fatal("expected rollback to destroy the nft table")
	}
	if _, err := os.Stat(p.StatePath); !os.IsNotExist(err) {
		t.Fatalf("state file should not exist after rollback, stat err = %v", err)
	}
}

func TestApplyAcquiresStateLock(t *testing.T) {
	p := basePlan(t)
	p.SkipEngine = true
	if res := Execute(context.Background(), systemFake(false), p); !res.OK {
		t.Fatalf("apply failed: %+v", res)
	}
	lock := filepath.Join(filepath.Dir(p.StatePath), "prx.lock")
	if _, err := os.Stat(lock); err != nil {
		t.Fatalf("expected a lock file next to the state: %v", err)
	}
}

// TestApplyReportsIncompleteRollback verifies that a rollback step which fails
// is surfaced to the caller instead of being silently discarded.
func TestApplyReportsIncompleteRollback(t *testing.T) {
	p := basePlan(t)
	p.SkipEngine = true

	f := &execx.Fake{Handler: func(name string, args []string, _ []byte) (execx.Result, error) {
		joined := strings.TrimSpace(name + " " + strings.Join(args, " "))
		switch {
		case joined == "ip -j route show default":
			return execx.Result{Stdout: []byte(`[{"dst":"default","gateway":"192.168.1.1","dev":"wlan0","metric":600}]`)}, nil
		case joined == "ip -6 -j route show default":
			return execx.Result{Stdout: []byte(`[]`)}, nil
		case joined == "ip -j addr show dev wlan0":
			return execx.Result{Stdout: []byte(`[{"ifname":"wlan0","addr_info":[{"family":"inet","local":"192.168.1.100","prefixlen":24,"scope":"global"}]}]`)}, nil
		case joined == "ip rule show":
			return execx.Result{Stdout: []byte("")}, nil
		case joined == "ip route show table 100":
			return execx.Result{Code: 2, Stderr: []byte("RTNETLINK answers: Operation not permitted")}, &execx.ExitError{Name: "ip", Code: 2, Stderr: []byte("RTNETLINK answers: Operation not permitted")}
		case joined == "ip route flush table 100":
			// The rollback itself fails.
			return execx.Result{Code: 2, Stderr: []byte("flush refused")}, &execx.ExitError{Name: "ip", Code: 2, Stderr: []byte("flush refused")}
		case strings.HasPrefix(joined, "nft"):
			return execx.Result{}, nil
		}
		return execx.Result{}, nil
	}}

	res := Execute(context.Background(), f, p)
	if res.OK {
		t.Fatalf("expected apply to fail: %+v", res)
	}
	if !strings.Contains(res.Error, "rollback incomplete") {
		t.Fatalf("res.Error = %q, want it to flag an incomplete rollback", res.Error)
	}
	found := false
	for _, m := range res.Messages {
		if strings.Contains(m, "rollback warning") {
			found = true
		}
	}
	if !found {
		t.Fatalf("expected a rollback warning message, got %v", res.Messages)
	}
}