//go:build linux

package cgroup

import (
	"context"
	"errors"
	"os"
	"path/filepath"
	"reflect"
	"strconv"
	"testing"
	"time"
)

func cleanupFixture(t *testing.T) (*Manager, string, string) {
	t.Helper()
	host, root := "cleanup.example", t.TempDir()
	dir := filepath.Join(root, ExecName(host))
	if err := os.Mkdir(dir, 0700); err != nil {
		t.Fatal(err)
	}
	for name, body := range map[string]string{"cgroup.events": "populated 1\n", "cgroup.procs": "123\n456\n", "cgroup.kill": ""} {
		if err := os.WriteFile(filepath.Join(dir, name), []byte(body), 0600); err != nil {
			t.Fatal(err)
		}
	}
	return &Manager{root: root, logf: t.Logf}, host, dir
}

func fixtureCleanup(t *testing.T, m *Manager, host string) *Cleanup {
	t.Helper()
	c, err := m.openCleanup(host, func(*os.File) error { return nil })
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = c.Close() })
	return c
}

func TestCleanupReopensAndDrains(t *testing.T) {
	m, host, dir := cleanupFixture(t)
	c := fixtureCleanup(t, m, host)
	if _, exists := reflect.TypeOf(c).MethodByName("Start"); exists {
		t.Fatal("cleanup grants spawn")
	}
	if populated, err := c.Populated(); err != nil || !populated {
		t.Fatal(populated, err)
	}
	if members, err := c.Members(); err != nil || !reflect.DeepEqual(members, []int{123, 456}) {
		t.Fatal(members, err)
	}
	ctx, cancel := context.WithTimeout(t.Context(), 30*time.Millisecond)
	defer cancel()
	if err := c.KillAndEmpty(ctx); !errors.Is(err, context.DeadlineExceeded) {
		t.Fatal(err)
	}
	if b, err := os.ReadFile(filepath.Join(dir, "cgroup.kill")); err != nil || string(b) != "1" {
		t.Fatalf("kill: %q %v", b, err)
	}
	if err := os.WriteFile(filepath.Join(dir, "cgroup.events"), []byte("populated 0\n"), 0600); err != nil {
		t.Fatal(err)
	}
	ctx2, cancel2 := context.WithTimeout(t.Context(), time.Second)
	defer cancel2()
	// Population alone is insufficient: the members still prevent success.
	ctx3, cancel3 := context.WithTimeout(t.Context(), 30*time.Millisecond)
	defer cancel3()
	if err := c.KillAndEmpty(ctx3); !errors.Is(err, context.DeadlineExceeded) {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(dir, "cgroup.procs"), nil, 0600); err != nil {
		t.Fatal(err)
	}
	if err := c.KillAndEmpty(ctx2); err != nil {
		t.Fatal(err)
	}
	if err := c.Close(); err != nil {
		t.Fatal(err)
	}
	if _, err := c.Populated(); !errors.Is(err, ErrStrictState) {
		t.Fatal(err)
	}
	if _, err := os.Stat(dir); err != nil {
		t.Fatal("close removed group", err)
	}
}

func TestCleanupRefusals(t *testing.T) {
	var absent *Manager
	if _, err := absent.OpenCleanup("cleanup.example"); !errors.Is(err, ErrNoDelegation) {
		t.Fatal(err)
	}
	if _, err := (&Manager{}).OpenCleanup("cleanup.example"); !errors.Is(err, ErrNoDelegation) {
		t.Fatal(err)
	}
	m, host, dir := cleanupFixture(t)
	if _, err := m.OpenCleanup(host); !errors.Is(err, ErrStrictUnsupported) {
		t.Fatal("ordinary filesystem admitted", err)
	}
	if _, err := m.openCleanup("../outside", func(*os.File) error { return nil }); !errors.Is(err, ErrStrictRefused) {
		t.Fatal(err)
	}
	if _, err := m.openCleanup("missing.example", func(*os.File) error { return nil }); !errors.Is(err, ErrNoCgroup) {
		t.Fatal(err)
	}
	if _, err := os.Stat(filepath.Join(m.root, ExecName("missing.example"))); !errors.Is(err, os.ErrNotExist) {
		t.Fatal("created missing group", err)
	}
	c := fixtureCleanup(t, m, host)
	for _, body := range []string{"", "populated 2\n", "populated 0\npopulated 0\n", "broken\n"} {
		if err := os.WriteFile(filepath.Join(dir, "cgroup.events"), []byte(body), 0600); err != nil {
			t.Fatal(err)
		}
		if _, err := c.Populated(); err == nil {
			t.Fatal("malformed events accepted")
		}
	}
	for _, body := range []string{"0\n", "-1\n", "oops\n", "9999999999999999999999999\n"} {
		if err := os.WriteFile(filepath.Join(dir, "cgroup.procs"), []byte(body), 0600); err != nil {
			t.Fatal(err)
		}
		if _, err := c.Members(); err == nil {
			t.Fatal("malformed members accepted")
		}
	}
	if err := c.KillAndEmpty(context.Background()); !errors.Is(err, ErrStrictDeadline) {
		t.Fatal(err)
	}
	if err := c.KillAndEmpty(nil); !errors.Is(err, ErrStrictDeadline) {
		t.Fatal(err)
	}
	if err := os.Rename(dir, dir+".old"); err != nil {
		t.Fatal(err)
	}
	if err := os.Mkdir(dir, 0700); err != nil {
		t.Fatal(err)
	}
	if _, err := c.Populated(); err == nil {
		t.Fatal("replaced identity accepted")
	}
	if _, err := c.Members(); err == nil {
		t.Fatal("replaced membership accepted")
	}
	if err := c.Close(); err != nil {
		t.Fatal(err)
	}
}

func TestCleanupDeadlineWhileLocked(t *testing.T) {
	m, host, _ := cleanupFixture(t)
	c := fixtureCleanup(t, m, host)
	c.state.pin.mu.Lock()
	ctx, cancel := context.WithTimeout(t.Context(), 30*time.Millisecond)
	defer cancel()
	done := make(chan error, 1)
	go func() { done <- c.KillAndEmpty(ctx) }()
	select {
	case err := <-done:
		c.state.pin.mu.Unlock()
		if !errors.Is(err, context.DeadlineExceeded) {
			t.Fatal(err)
		}
	case <-time.After(time.Second):
		c.state.pin.mu.Unlock()
		<-done
		t.Fatal("cleanup exceeded its deadline waiting for a handle lock")
	}
}

func TestCleanupRejectsMissingControlsAndSymlinks(t *testing.T) {
	for _, name := range []string{"cgroup.events", "cgroup.procs", "cgroup.kill"} {
		t.Run(name, func(t *testing.T) {
			m, host, dir := cleanupFixture(t)
			if err := os.Remove(filepath.Join(dir, name)); err != nil {
				t.Fatal(err)
			}
			if c, err := m.openCleanup(host, func(*os.File) error { return nil }); err == nil || c != nil || errors.Is(err, ErrNoCgroup) {
				t.Fatalf("missing control: %v %v", c, err)
			}
			if err := os.Symlink(filepath.Join(dir, "elsewhere"), filepath.Join(dir, name)); err != nil {
				t.Fatal(err)
			}
			if c, err := m.openCleanup(host, func(*os.File) error { return nil }); err == nil || c != nil {
				t.Fatalf("symlink control: %v %v", c, err)
			}
		})
	}
}

// A kernel before 5.14 has no cgroup.kill. A cleanup opened there names
// that (ErrNoKill), so a caller can fall back to the leader-only proof
// instead of reading it as an unproven kill on every stop. A cgroup.kill
// that is there but a symlink is not that case.
func TestCleanupNamesAMissingKillFile(t *testing.T) {
	m, host, dir := cleanupFixture(t)
	if err := os.Remove(filepath.Join(dir, "cgroup.kill")); err != nil {
		t.Fatal(err)
	}
	if c, err := m.openCleanup(host, func(*os.File) error { return nil }); c != nil || !errors.Is(err, ErrNoKill) {
		t.Fatalf("missing cgroup.kill: %v %v, want ErrNoKill", c, err)
	}
	if err := os.Symlink(filepath.Join(dir, "elsewhere"), filepath.Join(dir, "cgroup.kill")); err != nil {
		t.Fatal(err)
	}
	if c, err := m.openCleanup(host, func(*os.File) error { return nil }); c != nil || err == nil || errors.Is(err, ErrNoKill) {
		t.Fatalf("symlinked cgroup.kill: %v %v, want a refusal that is not ErrNoKill", c, err)
	}
}

func mkCgroupFixture(t *testing.T, dir, procs, events string) {
	t.Helper()
	if err := os.MkdirAll(dir, 0o700); err != nil {
		t.Fatal(err)
	}
	for name, body := range map[string]string{"cgroup.events": events, "cgroup.procs": procs, "cgroup.kill": ""} {
		if err := os.WriteFile(filepath.Join(dir, name), []byte(body), 0o600); err != nil {
			t.Fatal(err)
		}
	}
}

func fixtureCleanupAt(t *testing.T, root, host, child string) *Cleanup {
	t.Helper()
	c, err := (&Manager{root: root}).openCleanupAt(host, child, func(*os.File) error { return nil })
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = c.Close() })
	return c
}

// Stop evidence is taken on the parent (spec §7.3). A generation parent's
// own cgroup.procs is always empty; its generations' are not.
func TestCleanupMembersReachEveryGeneration(t *testing.T) {
	host, root := "gens.example", t.TempDir()
	parent := filepath.Join(root, ExecName(host))
	mkCgroupFixture(t, parent, "", "populated 1\n")
	mkCgroupFixture(t, filepath.Join(parent, "g1"), "11\n", "populated 1\n")
	mkCgroupFixture(t, filepath.Join(parent, "g2"), "22\n", "populated 1\n")
	// Review focus 3: a flat tenant runs as the agent's uid and can make
	// cgroups of its own under its generation.
	mkCgroupFixture(t, filepath.Join(parent, "g2", "tenant"), "33\n", "populated 1\n")
	c := fixtureCleanupAt(t, root, host, "")
	members, err := c.Members()
	if err != nil || !reflect.DeepEqual(members, []int{11, 22, 33}) {
		t.Fatalf("members = %v, %v; want every generation's", members, err)
	}
}

func TestCleanupMembersRefuseATreeTooDeep(t *testing.T) {
	host, root := "deep.example", t.TempDir()
	dir := filepath.Join(root, ExecName(host))
	mkCgroupFixture(t, dir, "", "populated 1\n")
	for i := range maxCleanupDepth + 1 {
		dir = filepath.Join(dir, "d"+strconv.Itoa(i))
		mkCgroupFixture(t, dir, "", "populated 0\n")
	}
	c := fixtureCleanupAt(t, root, host, "")
	if _, err := c.Members(); !errors.Is(err, ErrStrictRefused) {
		t.Fatalf("a tree deeper than %d: %v, want ErrStrictRefused", maxCleanupDepth, err)
	}
}

// The group budget is a width bound too, and it is enforced while the
// directory is read: a tree one group too wide is refused before any child
// is opened. Every child here is unreadable (mode 0), so a walk that
// opened one before counting them all would fail with EACCES instead.
func TestCleanupMembersRefuseATreeTooWide(t *testing.T) {
	if os.Geteuid() == 0 {
		t.Skip("root opens a mode-0 directory; the probe needs an ordinary user")
	}
	host, root := "wide.example", t.TempDir()
	dir := filepath.Join(root, ExecName(host))
	mkCgroupFixture(t, dir, "", "populated 1\n")
	for i := range maxCleanupGroups + 1 {
		if err := os.Mkdir(filepath.Join(dir, "w"+strconv.Itoa(i)), 0); err != nil {
			t.Fatal(err)
		}
	}
	c := fixtureCleanupAt(t, root, host, "")
	if _, err := c.Members(); !errors.Is(err, ErrStrictRefused) {
		t.Fatalf("%d groups at one level: %v, want ErrStrictRefused before any child is opened", maxCleanupGroups+1, err)
	}
}

func TestCleanupMembersAcceptATreeExactlyAtTheBound(t *testing.T) {
	host, root := "bound.example", t.TempDir()
	dir := filepath.Join(root, ExecName(host))
	mkCgroupFixture(t, dir, "", "populated 1\n")
	for i := range maxCleanupGroups {
		if err := os.Mkdir(filepath.Join(dir, "w"+strconv.Itoa(i)), 0o700); err != nil {
			t.Fatal(err)
		}
	}
	mkCgroupFixture(t, filepath.Join(dir, "w7"), "77\n", "populated 1\n")
	c := fixtureCleanupAt(t, root, host, "")
	if members, err := c.Members(); err != nil || !reflect.DeepEqual(members, []int{77}) {
		t.Fatalf("%d groups: %v, %v; want [77]", maxCleanupGroups, members, err)
	}
}

// KillAndEmpty on a generation acts only on that generation's child.
func TestCleanupChildKillsOnlyItsGeneration(t *testing.T) {
	host, root := "kill.example", t.TempDir()
	parent := filepath.Join(root, ExecName(host))
	mkCgroupFixture(t, parent, "", "populated 1\n")
	mkCgroupFixture(t, filepath.Join(parent, "g1"), "11\n", "populated 1\n")
	mkCgroupFixture(t, filepath.Join(parent, "g2"), "", "populated 0\n")
	c := fixtureCleanupAt(t, root, host, "g2")
	ctx, cancel := context.WithTimeout(t.Context(), time.Second)
	defer cancel()
	if err := c.KillAndEmpty(ctx); err != nil {
		t.Fatal(err)
	}
	for dir, want := range map[string]string{filepath.Join(parent, "g2"): "1", filepath.Join(parent, "g1"): "", parent: ""} {
		if got := readFixture(t, filepath.Join(dir, "cgroup.kill")); got != want {
			t.Errorf("%s/cgroup.kill = %q, want %q", dir, got, want)
		}
	}
}

func TestOpenCleanupChildRefusesForeignNamesAndMissingChildren(t *testing.T) {
	host, root := "names.example", t.TempDir()
	mkCgroupFixture(t, filepath.Join(root, ExecName(host)), "", "populated 0\n")
	m := &Manager{root: root}
	for _, child := range []string{"child", "../x", "g0"} {
		if _, err := m.openCleanupAt(host, child, func(*os.File) error { return nil }); !errors.Is(err, ErrStrictRefused) {
			t.Errorf("child %q: %v, want ErrStrictRefused", child, err)
		}
	}
	if _, err := m.openCleanupAt(host, "g9", func(*os.File) error { return nil }); !errors.Is(err, ErrNoCgroup) {
		t.Errorf("missing g9: %v, want ErrNoCgroup", err)
	}
}

// Review focus 4, Ruling 6: a process being placed under the host (here a
// generation into g1) is a member the kernel cannot see yet. While it is in
// flight the host's cleanup reports populated, refuses to list members,
// and kills nothing; once it resolves, Populated reads cgroup.events again.
func TestCleanupCountsPlacementInFlight(t *testing.T) {
	m, host, dir := cleanupFixture(t)
	writeFixture(t, filepath.Join(dir, "cgroup.events"), "populated 0\n")
	writeFixture(t, filepath.Join(dir, "cgroup.procs"), "")
	c := fixtureCleanup(t, m, host)
	done := m.beginPlacement(filepath.Join(dir, "g1"))
	defer done()
	if populated, err := c.Populated(); err != nil || !populated {
		t.Fatalf("Populated = %v, %v with a placement in flight: want true", populated, err)
	}
	if members, err := c.Members(); !errors.Is(err, ErrPlacementInFlight) {
		t.Fatalf("Members = %v, %v with a placement in flight: want ErrPlacementInFlight", members, err)
	}
	ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond)
	defer cancel()
	if err := c.KillAndEmpty(ctx); err == nil {
		t.Fatal("KillAndEmpty proved the host empty with a placement in flight")
	}
	if got := readFixture(t, filepath.Join(dir, "cgroup.kill")); got != "" {
		t.Fatalf("cgroup.kill = %q: the kill was written before the placement resolved", got)
	}
	done()
	if populated, err := c.Populated(); err != nil || populated {
		t.Fatalf("Populated = %v, %v once the placement resolved: want cgroup.events' answer, false", populated, err)
	}
	if members, err := c.Members(); err != nil || len(members) != 0 {
		t.Fatalf("Members = %v, %v once the placement resolved: want none", members, err)
	}
}

// Ruling 6: a placement that registers after KillAndEmpty's wait, here
// right after the kill is written, keeps the proof from finishing: every
// emptiness check after the kill counts it. The fake cgroupfs says empty
// throughout, so only the registry stands between the proof and a false
// empty.
func TestKillAndEmptySeesAPlacementRegisteredMidProof(t *testing.T) {
	m, host, dir := cleanupFixture(t)
	writeFixture(t, filepath.Join(dir, "cgroup.events"), "populated 0\n")
	writeFixture(t, filepath.Join(dir, "cgroup.procs"), "")
	c := fixtureCleanup(t, m, host)
	var done func()
	orig := cleanupKillWrite
	cleanupKillWrite = func(f *os.File) error {
		err := orig(f)
		done = m.beginPlacement(filepath.Join(dir, "g1"))
		return err
	}
	t.Cleanup(func() {
		cleanupKillWrite = orig
		if done != nil {
			done()
		}
	})
	ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond)
	defer cancel()
	if err := c.KillAndEmpty(ctx); err == nil {
		t.Fatal("KillAndEmpty proved the host empty while a placement registered mid-proof was in flight")
	}
	if done == nil {
		t.Fatal("the kill was never written: the test proves nothing")
	}
	// The control: with nothing in flight the same proof finishes.
	done()
	ctx2, cancel2 := context.WithTimeout(t.Context(), time.Second)
	defer cancel2()
	cleanupKillWrite = orig
	if err := c.KillAndEmpty(ctx2); err != nil {
		t.Fatalf("KillAndEmpty with nothing in flight = %v, want the proof", err)
	}
}

// Whole-branch review, finding 1, first schedule: a placement registers
// while the kernel is being read (the hook runs right after the read, before
// the registry is asked again). The fake cgroupfs says empty, so only the
// registry can stop a false empty: Populated must answer populated and
// Members must refuse, even though nothing was in flight when either began.
func TestCleanupCountsAPlacementRegisteredDuringTheRead(t *testing.T) {
	m, host, dir := cleanupFixture(t)
	writeFixture(t, filepath.Join(dir, "cgroup.events"), "populated 0\n")
	writeFixture(t, filepath.Join(dir, "cgroup.procs"), "")
	c := fixtureCleanup(t, m, host)
	var dones []func()
	t.Cleanup(func() {
		for _, d := range dones {
			d()
		}
	})
	reads := 0
	setPlacementReadHook(t, func(string) {
		reads++
		dones = append(dones, m.beginPlacement(filepath.Join(dir, "g1")))
	})
	if populated, err := c.Populated(); err != nil || !populated {
		t.Fatalf("Populated = %v, %v with a placement registered during the read: want true", populated, err)
	}
	for _, d := range dones {
		d()
	}
	if members, err := c.Members(); !errors.Is(err, ErrPlacementInFlight) {
		t.Fatalf("Members = %v, %v with a placement registered during the read: want ErrPlacementInFlight", members, err)
	}
	if reads != 2 {
		t.Fatalf("the read hook ran %d times, want 2: the test proves nothing", reads)
	}
	// The control: the same reads with nothing registering answer empty.
	for _, d := range dones {
		d()
	}
	placementReadHook = nil
	if populated, err := c.Populated(); err != nil || populated {
		t.Fatalf("Populated = %v, %v with nothing in flight: want false", populated, err)
	}
	if members, err := c.Members(); err != nil || len(members) != 0 {
		t.Fatalf("Members = %v, %v with nothing in flight: want none", members, err)
	}
}

// Finding 1, second schedule: a placement in flight when the kernel is read
// lands and ends its registration between the kernel read and the
// registry's. Asked in that order, the registry would read zero beside a
// kernel view taken before the process arrived. Both answers must count it.
func TestCleanupCountsAPlacementThatEndsBetweenTheReads(t *testing.T) {
	m, host, dir := cleanupFixture(t)
	writeFixture(t, filepath.Join(dir, "cgroup.events"), "populated 0\n")
	writeFixture(t, filepath.Join(dir, "cgroup.procs"), "")
	c := fixtureCleanup(t, m, host)
	var done func()
	setPlacementReadHook(t, func(string) {
		if done != nil {
			done()
		}
	})
	done = m.beginPlacement(filepath.Join(dir, "g1"))
	if populated, err := c.Populated(); err != nil || !populated {
		t.Fatalf("Populated = %v, %v with a placement that ended after the kernel read: want true", populated, err)
	}
	done = m.beginPlacement(filepath.Join(dir, "g1"))
	if members, err := c.Members(); !errors.Is(err, ErrPlacementInFlight) {
		t.Fatalf("Members = %v, %v with a placement that ended after the kernel read: want ErrPlacementInFlight", members, err)
	}
	if m.placementsUnder(dir) != 0 {
		t.Fatal("the hook did not end the placements: the test proves nothing")
	}
}

// Finding 1 through KillAndEmpty: a placement registered during each
// emptiness read after the kill keeps the proof from finishing, so it ends
// unproven at ctx's end and never answers empty.
func TestKillAndEmptyCountsAPlacementRegisteredDuringARead(t *testing.T) {
	m, host, dir := cleanupFixture(t)
	writeFixture(t, filepath.Join(dir, "cgroup.events"), "populated 0\n")
	writeFixture(t, filepath.Join(dir, "cgroup.procs"), "")
	c := fixtureCleanup(t, m, host)
	var dones []func()
	t.Cleanup(func() {
		for _, d := range dones {
			d()
		}
	})
	setPlacementReadHook(t, func(string) {
		// Each read sees one registration begin; the previous one ends, so
		// a check that asked the registry only before its read would see
		// nothing in flight.
		for _, d := range dones {
			d()
		}
		dones = append(dones, m.beginPlacement(filepath.Join(dir, "g1")))
	})
	ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond)
	defer cancel()
	if err := c.KillAndEmpty(ctx); !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("KillAndEmpty = %v, want unproven at ctx's end (context.DeadlineExceeded)", err)
	}
	if len(dones) == 0 {
		t.Fatal("no emptiness read ran: the test proves nothing")
	}
}