//go:build linux

package cgroup

import (
	"context"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

// The arm-time placement check (zdd readiness 3, Task 3). The ungated tests
// run anywhere; the kernel ones need CARLOS_SHIM_KERNEL=1 inside a unit with
// Delegate=yes, as shim_kernel_linux_test.go's do:
//
//	systemd-run --user --wait --pipe --collect -p Delegate=yes \
//	  -E TMPDIR=/tmp/k/zddr3 -E CARLOS_SHIM_KERNEL=1 -E CARLOS_GEN_KERNEL=1 \
//	  cgroup.test -test.run 'VerifyPlacement' -test.v -test.count=1

func assertNoPlacementCheckDir(t *testing.T, root string) {
	t.Helper()
	if _, err := os.Lstat(filepath.Join(root, placementCheckName)); !errors.Is(err, os.ErrNotExist) {
		t.Fatalf("%s left behind under %s (%v)", placementCheckName, root, err)
	}
}

func TestVerifyPlacementNilManager(t *testing.T) {
	var m *Manager
	if err := m.VerifyPlacement(context.Background()); !errors.Is(err, ErrNoDelegation) {
		t.Fatalf("nil manager: %v, want ErrNoDelegation", err)
	}
}

func TestVerifyPlacementWithoutTheShim(t *testing.T) {
	prev := shimHooked.Load()
	shimHooked.Store(false)
	t.Cleanup(func() { shimHooked.Store(prev) })
	root := t.TempDir()
	m := &Manager{root: root, logf: t.Logf}
	err := m.VerifyPlacement(context.Background())
	if err == nil || !strings.Contains(err.Error(), "shim") {
		t.Fatalf("no shim: %v, want an error naming the shim", err)
	}
	assertNoPlacementCheckDir(t, root)
}

// A manager over a directory that is no cgroup answers an error and leaves
// nothing behind, a leftover from an earlier run included.
func TestVerifyPlacementOverATempRootFails(t *testing.T) {
	root := t.TempDir()
	if err := os.Mkdir(filepath.Join(root, placementCheckName), 0o755); err != nil {
		t.Fatal(err)
	}
	m := &Manager{root: root, logf: t.Logf}
	err := m.VerifyPlacement(context.Background())
	if err == nil {
		t.Fatal("a temp directory passed the placement check")
	}
	t.Logf("temp root: %v", err)
	assertNoPlacementCheckDir(t, root)
	if m.placementsUnder(root) != 0 {
		t.Fatal("a placement is still registered")
	}
}

// TestVerifyPlacementKernelChild is the probe that never exits: it sleeps
// when it finds itself in the placement check's cgroup, and otherwise (a
// normal test run) returns at once.
func TestVerifyPlacementKernelChild(t *testing.T) {
	raw, err := os.ReadFile("/proc/self/cgroup")
	if err != nil {
		return
	}
	rel, ok := ParseSelfCgroup(string(raw))
	if !ok || filepath.Base(rel) != placementCheckName {
		return
	}
	time.Sleep(time.Minute)
}

// awaitProbe judges a probe that has already exited by its own exit, never
// by which of its two channels the select saw first: with ctx already done
// and the probe a zombie, an exit 0 passes and an exit 1 fails as the
// probe's failure, every time. One still running is killed, reaped, and
// fails on ctx's error. Ungated: plain processes, no cgroup.
func TestAwaitProbeJudgesAnExitedProbeByItsExit(t *testing.T) {
	done, cancel := context.WithCancel(context.Background())
	cancel()
	for i := 0; i < 50; i++ {
		ok := exitedProbe(t, "/bin/true")
		if err := awaitProbe(done, ok); err != nil {
			t.Fatalf("round %d: a probe that exited 0 before ctx ended: %v, want nil", i, err)
		}
		bad := exitedProbe(t, "/bin/false")
		err := awaitProbe(done, bad)
		if err == nil || errors.Is(err, context.Canceled) || !strings.Contains(err.Error(), "the probe failed: exit status 1") {
			t.Fatalf("round %d: a probe that exited 1 before ctx ended: %v, want its own exit", i, err)
		}
	}
	hung := exec.Command("/bin/sleep", "60")
	if err := hung.Start(); err != nil {
		t.Fatal(err)
	}
	err := awaitProbe(done, hung)
	if !errors.Is(err, context.Canceled) || !strings.Contains(err.Error(), "did not exit") {
		t.Fatalf("a running probe with ctx done: %v, want ctx's error", err)
	}
	if hung.ProcessState == nil {
		t.Fatal("the killed probe was not reaped")
	}
}

// exitedProbe starts path and returns once the process has exited and is
// a zombie, not yet reaped.
func exitedProbe(t *testing.T, path string) *exec.Cmd {
	t.Helper()
	cmd := exec.Command(path)
	if err := cmd.Start(); err != nil {
		t.Fatal(err)
	}
	stat := fmt.Sprintf("/proc/%d/stat", cmd.Process.Pid)
	for deadline := time.Now().Add(5 * time.Second); time.Now().Before(deadline); time.Sleep(time.Millisecond) {
		raw, err := os.ReadFile(stat)
		if err != nil {
			t.Fatal(err)
		}
		// The state follows the parenthesised command name.
		if i := strings.LastIndexByte(string(raw), ')'); i >= 0 && strings.HasPrefix(string(raw[i:]), ") Z") {
			return cmd
		}
	}
	t.Fatalf("%s did not exit", path)
	return nil
}

func TestVerifyPlacementKernel(t *testing.T) {
	m := shimKernelManager(t)
	if err := m.VerifyPlacement(context.Background()); err != nil {
		t.Fatalf("VerifyPlacement: %v", err)
	}
	assertNoPlacementCheckDir(t, m.root)
	if m.placementsUnder(m.root) != 0 {
		t.Fatal("a placement is still registered")
	}
}

// A leftover from an earlier run (empty: the unit's restart killed its
// members) is replaced, and the check still passes.
func TestVerifyPlacementKernelReplacesALeftover(t *testing.T) {
	m := shimKernelManager(t)
	if err := os.Mkdir(filepath.Join(m.root, placementCheckName), 0o755); err != nil {
		t.Fatal(err)
	}
	if err := m.VerifyPlacement(context.Background()); err != nil {
		t.Fatalf("VerifyPlacement over a leftover: %v", err)
	}
	assertNoPlacementCheckDir(t, m.root)
}

// A shim that cannot place (here, one that is not the carlos binary at all)
// fails the check, and the scratch cgroup is removed.
func TestVerifyPlacementKernelBrokenShim(t *testing.T) {
	m := shimKernelManager(t)
	prev := shimExe
	shimExe = "/bin/true"
	t.Cleanup(func() { shimExe = prev })
	err := m.VerifyPlacement(context.Background())
	if err == nil {
		t.Fatal("a shim that cannot place passed the check")
	}
	t.Logf("broken shim: %v", err)
	assertNoPlacementCheckDir(t, m.root)
}

// A probe whose output is not its cgroup, or that exits non-zero, fails
// the check: placement is proven only by what the placed process saw.
func TestVerifyPlacementKernelChecksTheProbe(t *testing.T) {
	m := shimKernelManager(t)
	for name, args := range map[string][]string{
		// The test binary with no tests: exits 0, prints "PASS" and a warning.
		"wrong output": {"-test.run=^$"},
		// An unknown flag: exits 2.
		"non-zero exit": {"-test.run=^$", "-test.no-such-flag"},
	} {
		t.Run(name, func(t *testing.T) {
			prev := placementProbeArgs
			placementProbeArgs = args
			t.Cleanup(func() { placementProbeArgs = prev })
			err := m.VerifyPlacement(context.Background())
			if err == nil {
				t.Fatal("passed with a probe that did not report its cgroup")
			}
			t.Logf("%s: %v", name, err)
			assertNoPlacementCheckDir(t, m.root)
		})
	}
}

// A probe that does not exit is killed when ctx ends, and reaped before the
// scratch cgroup is removed: the removal succeeds.
func TestVerifyPlacementKernelKillsAProbeThatHangs(t *testing.T) {
	m := shimKernelManager(t)
	prev := placementProbeArgs
	placementProbeArgs = []string{"-test.run=^TestVerifyPlacementKernelChild$", "-test.count=1"}
	t.Cleanup(func() { placementProbeArgs = prev })
	ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
	defer cancel()
	start := time.Now()
	err := m.VerifyPlacement(ctx)
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("hung probe: %v, want the context's deadline", err)
	}
	if d := time.Since(start); d > 10*time.Second {
		t.Fatalf("took %v to give up on a hung probe", d)
	}
	t.Logf("hung probe: %v", err)
	assertNoPlacementCheckDir(t, m.root)
}

// A shim that stalls, before placement or after the release, cannot hold
// the arm-time check past its budget: the decline, the reap and the removal
// of the scratch cgroup all land within placementCheckBudget plus a few
// milliseconds, under 5 s, with the caller's ctx carrying no deadline at
// all, as placementCheck's does not. A sooner ctx deadline caps the shim's
// waits the same way.
func TestVerifyPlacementKernelStalledShimDeclinesWithinBudget(t *testing.T) {
	const hard = 5 * time.Second
	for _, c := range []struct {
		name string
		hold int
		ctx  time.Duration // 0: context.Background()
		want time.Duration // the wait that must end it
	}{
		{"before placement", 1, 0, placementCheckBudget},
		{"after release", 2, 0, placementCheckBudget},
		{"before placement, sooner ctx", 1, 300 * time.Millisecond, 300 * time.Millisecond},
	} {
		t.Run(c.name, func(t *testing.T) {
			m := shimKernelManager(t)
			setShimVar(t, &shimTestHold, c.hold)
			ctx := context.Background()
			if c.ctx > 0 {
				var cancel context.CancelFunc
				ctx, cancel = context.WithTimeout(ctx, c.ctx)
				defer cancel()
			}
			start := time.Now()
			err := m.VerifyPlacement(ctx)
			took := time.Since(start)
			t.Logf("hold %d: took %v: %v", c.hold, took, err)
			if err == nil || !isPlacementError(err) {
				t.Fatalf("a stalled shim: %v, want a placement error", err)
			}
			if took < c.want-50*time.Millisecond {
				t.Fatalf("declined after %v, before the %v wait ended: not the stall the test is about", took, c.want)
			}
			if took > c.want+500*time.Millisecond || took > hard {
				t.Fatalf("declined after %v, want within %v plus 500ms and under %v", took, c.want, hard)
			}
			assertNoPlacementCheckDir(t, m.root)
			assertNothingInFlight(t, m)
		})
	}
}