//go:build linux

package host

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

// The generation census that carlos-system-recover runs, on the two
// refusals the harness box hit: a non-root service still running the
// /opt/carlos/carlos an apply renamed over ("(deleted)"), and short-lived
// tasks the service lookup refused with a bare "permission denied". These
// drive the production service lookup (generationServiceEnv.reader) over a
// fake /proc, replacing only systemctl, the unit files' ownership check,
// the signed-binary proof and the owner uid a non-root test cannot have.

type deletedExeFixture struct {
	generationProcFixture
	opt, image string
	env        generationServiceEnv
	current    string // digest the env proves for <opt>/carlos
	prev       string // and for <opt>/carlos.prev; "" is no carlos.prev
	execStart  string
	mainPID    string
	unit       string
	tt         *testing.T
	// deletedName is the exe link's base name; "" is "carlos (deleted)".
	deletedName string
	// imageFile is the open descriptor of the task's current image.
	imageFile *os.File
}

func (f *deletedExeFixture) t() *testing.T { return f.tt }

const deletedExeBody = "carlos router build, now renamed over"

func deletedExe(t *testing.T) *deletedExeFixture {
	t.Helper()
	g := generationProc(t)
	f := &deletedExeFixture{generationProcFixture: g, opt: filepath.Join(t.TempDir(), "opt", "carlos"), mainPID: "41", unit: "carlos-relay.service"}
	if err := os.MkdirAll(f.opt, 0o755); err != nil {
		t.Fatal(err)
	}
	f.execStart = "{ path=" + f.opt + "/carlos ; argv[]=" + f.opt + "/carlos relay ; ignore_errors=no ; start_time=[n/a] ; stop_time=[n/a] ; pid=41 ; code=(null) ; status=0/0 }"
	generationWrite(t, f.task+"/status", "Tgid:\t41\nUid:\t1001\t1001\t1001\t1001\nKthread:\t0\n")
	group := "/system.slice/carlos-relay.service"
	generationWrite(t, f.task+"/cgroup", "0::"+group+"\n")
	generationWrite(t, f.proc+"/41/cgroup", "0::"+group+"\n")
	// The image the task runs: open, then unlinked, as the apply's rename
	// leaves the edge's. The exe link reads "<opt>/carlos (deleted)" and
	// opens that image, as /proc/<pid>/exe does.
	f.image = filepath.Join(t.TempDir(), "image")
	f.writeImage(t, deletedExeBody, 0o755, true)
	sum := sha256.Sum256([]byte(deletedExeBody))
	f.current = "not-this-build"
	f.prev = hex.EncodeToString(sum[:])
	f.env = generationServiceEnv{proc: f.proc, opt: f.opt, owner: uint32(os.Getuid()),
		show: func(_ context.Context, unit string) (map[string]string, error) {
			if unit != f.unit {
				return nil, fmt.Errorf("systemctl show %s: no such unit", unit)
			}
			return map[string]string{"ControlGroup": group, "MainPID": f.mainPID, "ExecStart": f.execStart, "FragmentPath": "/etc/systemd/system/" + unit}, nil
		},
		trust: func(string) error { return nil },
		binary: func(_ context.Context, opt, path string) (rootFloorBinary, error) {
			switch path {
			case filepath.Join(opt, "carlos"):
				return rootFloorBinary{Digest: f.current}, nil
			case filepath.Join(opt, "carlos.prev"):
				if f.prev == "" {
					return rootFloorBinary{}, os.ErrNotExist
				}
				return rootFloorBinary{Digest: f.prev}, nil
			}
			return rootFloorBinary{}, os.ErrPermission
		},
	}
	return f
}

// writeImage replaces the task's image with body at mode, unlinked or not.
func (f *deletedExeFixture) writeImage(t *testing.T, body string, mode os.FileMode, unlink bool) {
	t.Helper()
	path := f.image + fmt.Sprint(time.Now().UnixNano())
	if err := os.WriteFile(path, []byte(body), 0o600); err != nil {
		t.Fatal(err)
	}
	if err := os.Chmod(path, mode); err != nil {
		t.Fatal(err)
	}
	img, err := os.Open(path)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { img.Close() })
	f.imageFile = img
	if unlink {
		if err := os.Remove(path); err != nil {
			t.Fatal(err)
		}
	}
	name := f.deletedName
	if name == "" {
		name = "carlos (deleted)"
	}
	deleted := filepath.Join(f.opt, name)
	os.Remove(deleted)
	generationLink(t, fmt.Sprintf("/proc/%d/fd/%d", os.Getpid(), img.Fd()), deleted)
	os.Remove(f.task + "/exe")
	generationLink(t, deleted, f.task+"/exe")
}

// kernelCensus is one census pass through the production lookup, as
// productionGenerationKernelCensus wires it: classify and recheck.
func (f *deletedExeFixture) kernelCensus(ctx context.Context) error {
	service, recheck := f.env.reader(ctx)
	_, err := (generationKernel{proc: f.proc, service: service, recheck: recheck}).absent(ctx, f.reservation)
	return err
}

func (f *deletedExeFixture) census(t *testing.T) error {
	t.Helper()
	return f.kernelCensus(t.Context())
}

func TestRecoverCensusAcceptsADeletedExeOfTheCurrentOrPreviousBuild(t *testing.T) {
	for _, which := range []string{"previous", "current", "current-no-prev"} {
		t.Run(which, func(t *testing.T) {
			f := deletedExe(t)
			switch which {
			case "current":
				f.current, f.prev = f.prev, "other-build"
			case "current-no-prev":
				f.current, f.prev = f.prev, ""
			}
			if err := f.census(t); err != nil {
				t.Fatalf("census refused a deleted exe of the %s signed build: %v", which, err)
			}
		})
	}
}

func TestRecoverCensusAcceptsTheEdgesDeletedExeInSupervise(t *testing.T) {
	// The harness's b2 task: carlos-edge.service's main process, in its
	// supervise/ group, still running the pre-apply image, with and
	// without the guard drop-in (whose ExecStart is the guard and whose
	// selected program is <opt>/carlos).
	for _, guarded := range []bool{false, true} {
		t.Run(map[bool]string{false: "unguarded", true: "guarded"}[guarded], func(t *testing.T) {
			f := deletedExe(t)
			f.unit = "carlos-edge.service"
			if guarded {
				f.execStart = "{ path=/opt/carlos/system/guard/carlos ; argv[]=/opt/carlos/system/guard/carlos system guard-exec edge ; ignore_errors=no ; start_time=[n/a] ; stop_time=[n/a] ; pid=41 ; code=(null) ; status=0/0 }"
				binary := f.env.binary
				f.env.binary = func(ctx context.Context, opt, path string) (rootFloorBinary, error) {
					if path == "/opt/carlos/system/guard/carlos" {
						return rootFloorBinary{Digest: "guard"}, nil
					}
					b, err := binary(ctx, opt, path)
					b.File.Inode = 7 // the live <opt>/carlos, not the deleted image
					return b, err
				}
			}
			base := "/system.slice/carlos-edge.service"
			generationWrite(t, f.task+"/cgroup", "0::"+base+"/supervise\n")
			generationWrite(t, f.proc+"/41/cgroup", "0::"+base+"/supervise\n")
			show := f.env.show
			f.env.show = func(ctx context.Context, unit string) (map[string]string, error) {
				props, err := show(ctx, unit)
				if err == nil {
					props["ControlGroup"] = base
				}
				return props, err
			}
			if err := f.census(t); err != nil {
				t.Fatalf("census refused the edge's deleted signed image: %v", err)
			}
			// The same edge running some other deleted image still refuses.
			f.writeImage(t, "an image no identity signs", 0o755, true)
			if err := f.census(t); err == nil {
				t.Fatal("census accepted an unsigned deleted image in the edge")
			}
		})
	}
}

func TestRecoverCensusRefusesOtherDeletedExes(t *testing.T) {
	for _, attack := range []string{"unsigned", "no-proof", "linked", "group-writable", "not-router-binary", "other-program", "deleted-helper", "other-path"} {
		t.Run(attack, func(t *testing.T) {
			f := deletedExe(t)
			switch attack {
			case "unsigned":
				f.writeImage(t, "some other build", 0o755, true)
			case "no-proof":
				f.env.binary = func(context.Context, string, string) (rootFloorBinary, error) {
					return rootFloorBinary{}, errRootFloorHistory
				}
			case "linked":
				f.writeImage(t, deletedExeBody, 0o755, false)
			case "group-writable":
				f.writeImage(t, deletedExeBody, 0o775, true)
			case "not-router-binary":
				// The unit runs another program: a deleted <opt>/carlos is
				// not what it selects.
				f.execStart = strings.ReplaceAll(f.execStart, f.opt+"/carlos", f.opt+"/helper")
			case "other-program":
				f.execStart = strings.ReplaceAll(f.execStart, f.opt+"/carlos", "/usr/bin/carlos")
			case "other-path":
				// The unit selects <opt>/carlos, but the task's image was
				// somewhere else when it was deleted.
				f.deletedName = "carlos-copy (deleted)"
				f.writeImage(t, deletedExeBody, 0o755, true)
			case "deleted-helper":
				// The unit's own program, deleted, with a signed router
				// build's bytes: only <opt>/carlos is ever accepted.
				f.execStart = strings.ReplaceAll(f.execStart, f.opt+"/carlos", f.opt+"/helper")
				f.deletedName = "helper (deleted)"
				f.writeImage(t, deletedExeBody, 0o755, true)
			}
			err := f.census(t)
			if err == nil {
				t.Fatalf("%s: census accepted the deleted exe", attack)
			}
			if !errors.Is(err, os.ErrPermission) && attack != "no-proof" {
				t.Fatalf("%s: refusal %v is not a permission refusal", attack, err)
			}
			if !strings.Contains(err.Error(), "deleted exe check refused") {
				t.Fatalf("%s: refusal %q does not name the deleted exe check", attack, err)
			}
		})
	}
}

func TestRecoverCensusRefusesADeletedExeNotOwnedByRoot(t *testing.T) {
	if os.Getuid() == 0 {
		t.Skip("the fixture's image is root-owned when run as root")
	}
	f := deletedExe(t)
	f.env.owner = 0 // production: the image must be root's
	if err := f.census(t); err == nil || !strings.Contains(err.Error(), "not an unlinked root-owned file") {
		t.Fatalf("census = %v, want a refusal of a non-root-owned image", err)
	}
}

func TestRecoverCensusStillChecksReferencesOfAnAcceptedDeletedExe(t *testing.T) {
	for _, ref := range []string{"fd", "mapped", "cwd"} {
		t.Run(ref, func(t *testing.T) {
			f := deletedExe(t)
			b := f.reservation.Binding
			switch ref {
			case "fd":
				generationLink(t, b.Path, f.task+"/fd/3")
			case "mapped":
				generationWrite(t, f.task+"/maps", fmt.Sprintf("1000-2000 r-xp 0000 %x:%x %d %s\n", generationMajor(b.Executable.Device), generationMinor(b.Executable.Device), b.Executable.Inode, b.Path))
			case "cwd":
				generationRelink(t, b.Tree, f.task+"/cwd")
			}
			if err := f.census(t); err == nil {
				t.Fatalf("a deleted-exe task referencing a retired generation (%s) was accepted", ref)
			}
		})
	}
}

func TestRecoverCensusRefusalsNameCheckUnitCgroupAndExe(t *testing.T) {
	for _, tc := range []struct {
		name, check, unit, group string
		setup                    func(*deletedExeFixture)
	}{
		{"activating", "main pid", "carlos-relay.service", "/system.slice/carlos-relay.service", func(f *deletedExeFixture) { f.mainPID = "0" }},
		{"session", "unit", "session-3.scope", "/user.slice/user-1000.slice/session-3.scope", func(f *deletedExeFixture) {
			generationWrite(f.t(), f.task+"/cgroup", "0::/user.slice/user-1000.slice/session-3.scope\n")
		}},
		{"live-exe", "exe", "carlos-relay.service", "/system.slice/carlos-relay.service", func(f *deletedExeFixture) {
			// A live (not deleted) exe the lookup cannot walk to as a
			// root-owned file under the task's root.
			generationRelink(f.t(), f.unrelated, f.task+"/exe")
			generationRelink(f.t(), "/", f.task+"/root")
		}},
	} {
		t.Run(tc.name, func(t *testing.T) {
			f := deletedExe(t)
			f.tt = t
			tc.setup(f)
			exe, _ := os.Readlink(f.task + "/exe")
			err := f.census(t)
			if err == nil {
				t.Fatal("census accepted")
			}
			for _, want := range []string{"/41/task/41: service: " + tc.check, "service: " + tc.check + " check refused", "unit " + tc.unit, "cgroup " + tc.group, "exe " + exe} {
				if !strings.Contains(err.Error(), want) {
					t.Fatalf("refusal %q does not name %q", err, want)
				}
			}
		})
	}
}

func TestRecoverCensusRetriesATransientRefusalWithinTheWindow(t *testing.T) {
	clock := time.Unix(0, 0)
	var slept []time.Duration
	now := func() time.Time { return clock }
	sleep := func(_ context.Context, d time.Duration) error {
		if clock.Sub(time.Unix(0, 0)) > time.Minute {
			return errors.New("unbounded retry") // ends a run with no window
		}
		slept = append(slept, d)
		clock = clock.Add(d)
		return nil
	}

	// A unit still activating, through the real census pass: it clears on
	// the fourth pass.
	f := deletedExe(t)
	passes := 0
	inspect := func(ctx context.Context) (string, string, error) {
		passes++
		f.mainPID = "0"
		if passes == 4 {
			f.mainPID = "41"
		}
		err := f.kernelCensus(ctx)
		return "boot", "digest", err
	}
	if _, _, err := generationCensusRetryFor(t.Context(), inspect, 20*time.Second, now, sleep); err != nil {
		t.Fatalf("a transient refusal did not clear within the window: %v", err)
	}
	if passes != 4 || fmt.Sprint(slept) != "[250ms 500ms 1s]" {
		t.Fatalf("passes=%d slept=%v, want 4 passes after 250ms, 500ms, 1s", passes, slept)
	}

	// The same refusal, persistent: refused once the next wait would end
	// past the window, with the last pass's refusal.
	clock, slept, passes = time.Unix(0, 0), nil, 0
	persistent := func(ctx context.Context) (string, string, error) {
		passes++
		f.mainPID = "0"
		err := f.kernelCensus(ctx)
		return "", "", err
	}
	_, _, err := generationCensusRetryFor(t.Context(), persistent, 20*time.Second, now, sleep)
	if err == nil || !strings.Contains(err.Error(), "main pid check refused") {
		t.Fatalf("persistent refusal = %v", err)
	}
	var total time.Duration
	for _, d := range slept {
		total += d
	}
	if total > 20*time.Second || total < 15*time.Second || passes != len(slept)+1 {
		t.Fatalf("persistent refusal: %d passes, slept %v (%s), want about 20 s then refuse", passes, slept, total)
	}
}

func TestRecoverCensusRetriesOnlyTaskLevelTransients(t *testing.T) {
	sleep := func(context.Context, time.Duration) error { return nil }
	for _, tc := range []struct {
		name  string
		err   error
		retry bool
	}{
		{"service", &generationTaskRefusal{task: "/proc/9/task/9", err: fmt.Errorf("service: %w", &generationServiceRefusal{check: "unit", err: os.ErrPermission})}, true},
		{"exited", &generationTaskRefusal{task: "/proc/9/task/9", err: os.ErrNotExist}, true},
		{"replaced", &generationTaskRefusal{task: "/proc/9/task/9", err: errGenerationTaskReplaced}, true},
		{"reference", &generationTaskRefusal{task: "/proc/9/task/9", err: fmt.Errorf("fd: %w", os.ErrPermission)}, false},
		{"mounts", &generationTaskRefusal{task: "/proc/9/task/9", err: fmt.Errorf("mounts: %w", os.ErrPermission)}, false},
		{"structural", os.ErrPermission, false},
	} {
		t.Run(tc.name, func(t *testing.T) {
			passes := 0
			_, _, err := generationCensusRetryFor(t.Context(), func(context.Context) (string, string, error) {
				passes++
				if passes == 2 {
					return "boot", "digest", nil
				}
				return "", "", tc.err
			}, 20*time.Second, time.Now, sleep)
			if tc.retry && (err != nil || passes != 2) {
				t.Fatalf("transient %s: passes=%d err=%v", tc.name, passes, err)
			}
			if !tc.retry && (err == nil || passes != 1) {
				t.Fatalf("%s retried: passes=%d err=%v", tc.name, passes, err)
			}
		})
	}
	// A cancelled wait refuses with the refusal it was waiting out.
	ctx, cancel := context.WithCancel(t.Context())
	refusal := &generationTaskRefusal{task: "/proc/9/task/9", err: os.ErrNotExist}
	_, _, err := generationCensusRetryFor(ctx, func(context.Context) (string, string, error) { return "", "", refusal }, 20*time.Second, time.Now, func(context.Context, time.Duration) error { cancel(); return context.Canceled })
	if !errors.Is(err, refusal) {
		t.Fatalf("cancelled wait = %v", err)
	}
}

func TestProductionGenerationCensusRefusalsAreTyped(t *testing.T) {
	// Every refusal of the production census is a generationCensusRefusal,
	// which is what lets a no-journal recovery defer it. As a non-root test
	// the census refuses at its first check.
	if os.Geteuid() == 0 {
		t.Skip("the census does not refuse at once as root")
	}
	_, _, err := productionGenerationKernelCensus(t.Context())
	var refusal *generationCensusRefusal
	if !errors.As(err, &refusal) || !onlyCensusRefusal(err) || !errors.Is(err, os.ErrPermission) {
		t.Fatalf("census refusal = %#v, want a typed census refusal", err)
	}
}

// execRace runs one census pass through the production lookup with a hook
// around each of its rechecks: before(final) runs before the recheck and
// after(final) after it. The hooks stand in for the task exec'ing (the
// same PID, the same start time) while the census inspects it.
func (f *deletedExeFixture) execRace(t *testing.T, after func(final bool, call int), before func(final bool, call int)) error {
	t.Helper()
	classify, recheck := f.env.reader(t.Context())
	calls := 0
	hooked := func(path string, final bool) error {
		calls++
		if before != nil {
			before(final, calls)
		}
		err := recheck(path, final)
		if after != nil {
			after(final, calls)
		}
		return err
	}
	_, err := (generationKernel{proc: f.proc, service: classify, recheck: hooked}).absent(t.Context(), f.reservation)
	return err
}

func TestRecoverCensusRefusesATaskThatExecsAfterItsImageWasProven(t *testing.T) {
	helper := func(f *deletedExeFixture) { generationRelink(t, f.unrelated, f.task+"/exe") }
	for _, tc := range []struct {
		name string
		race func(f *deletedExeFixture) (after, before func(bool, int))
	}{
		// The task execs an unsigned helper between the hash and the first
		// use of the trust.
		{"after-classification", func(f *deletedExeFixture) (func(bool, int), func(bool, int)) {
			return nil, func(_ bool, call int) {
				if call == 1 {
					helper(f)
				}
			}
		}},
		// It execs another image that reads the same, "<opt>/carlos
		// (deleted)" with signed bytes, but is not the inode hashed.
		{"same-name-other-inode", func(f *deletedExeFixture) (func(bool, int), func(bool, int)) {
			return nil, func(_ bool, call int) {
				if call == 1 {
					f.writeImage(t, deletedExeBody, 0o755, true)
				}
			}
		}},
		// It execs the helper after the first check and is back on the
		// proven image by the end: the check at the mounts use sees it.
		{"between-uses", func(f *deletedExeFixture) (func(bool, int), func(bool, int)) {
			target, err := os.Readlink(f.task + "/exe")
			if err != nil {
				t.Fatal(err)
			}
			return func(final bool, call int) {
					if !final && call == 1 {
						helper(f)
					}
				}, func(final bool, _ int) {
					if final {
						generationRelink(t, target, f.task+"/exe")
					}
				}
		}},
		// The proven image itself is made group-writable mid-inspection.
		{"image-made-writable", func(f *deletedExeFixture) (func(bool, int), func(bool, int)) {
			return nil, func(_ bool, call int) {
				if call == 1 {
					if err := f.imageFile.Chmod(0o775); err != nil {
						t.Fatal(err)
					}
				}
			}
		}},
		// It execs the helper after the last use, before the end.
		{"before-the-end", func(f *deletedExeFixture) (func(bool, int), func(bool, int)) {
			return func(final bool, call int) {
				if !final && call == 2 {
					helper(f)
				}
			}, nil
		}},
	} {
		t.Run(tc.name, func(t *testing.T) {
			f := deletedExe(t)
			if err := f.execRace(t, nil, nil); err != nil {
				t.Fatalf("census without a race refused: %v", err)
			}
			after, before := tc.race(f)
			err := f.execRace(t, after, before)
			if err == nil {
				t.Fatal("census accepted a task that exec'd away from its proven image")
			}
			if !strings.Contains(err.Error(), "exe binding check refused") {
				t.Fatalf("refusal %q does not name the exe binding check", err)
			}
			var task *generationTaskRefusal
			if !errors.As(err, &task) || !task.transient() {
				t.Fatalf("refusal %v is not a transient task refusal a later pass can clear", err)
			}
		})
	}
}

func TestRecoverCensusNeedsABindingForATrustedTask(t *testing.T) {
	// A kernel wired with a lookup but no recheck refuses every task the
	// lookup trusts.
	f := deletedExe(t)
	classify, _ := f.env.reader(t.Context())
	if _, err := (generationKernel{proc: f.proc, service: classify}).absent(t.Context(), f.reservation); err == nil || !strings.Contains(err.Error(), "no executable binding") {
		t.Fatalf("census without a recheck = %v", err)
	}
}

func TestRecoverCensusDeadlineCountsInspectionTime(t *testing.T) {
	clock := time.Unix(0, 0)
	now := func() time.Time { return clock }
	sleep := func(_ context.Context, d time.Duration) error { clock = clock.Add(d); return nil }
	transient := &generationTaskRefusal{task: "/proc/9/task/9", err: os.ErrNotExist}

	// A slow pass that succeeds past the deadline refuses.
	_, _, err := generationCensusRetryFor(t.Context(), func(context.Context) (string, string, error) {
		clock = clock.Add(25 * time.Second)
		return "boot", "digest", nil
	}, 20*time.Second, now, sleep)
	if !errors.Is(err, errGenerationCensusLate) {
		t.Fatalf("a pass that ended past the window = %v, want a late refusal", err)
	}

	// Slow refused passes: inspection time and waits share the deadline,
	// no pass starts after it, and the census refuses inside it.
	clock = time.Unix(0, 0)
	var starts []time.Duration
	_, _, err = generationCensusRetryFor(t.Context(), func(context.Context) (string, string, error) {
		starts = append(starts, clock.Sub(time.Unix(0, 0)))
		clock = clock.Add(6 * time.Second)
		return "", "", transient
	}, 20*time.Second, now, sleep)
	if err == nil {
		t.Fatal("persistent refusal accepted")
	}
	for _, s := range starts {
		if s >= 20*time.Second {
			t.Fatalf("a pass started at %s, past the window (%v)", s, starts)
		}
	}
	if end := clock.Sub(time.Unix(0, 0)); end > 26*time.Second {
		t.Fatalf("census ran to %s with 6 s passes; it must stop at the first pass that ends past 20 s", end)
	}
	if len(starts) < 2 {
		t.Fatalf("passes started at %v: the slow passes were not retried inside the window", starts)
	}

	// A wait that overruns (a timer firing late) never starts a pass past
	// the deadline.
	clock, starts = time.Unix(0, 0), nil
	oversleep := func(_ context.Context, d time.Duration) error { clock = clock.Add(d + 3*time.Second); return nil }
	_, _, err = generationCensusRetryFor(t.Context(), func(context.Context) (string, string, error) {
		starts = append(starts, clock.Sub(time.Unix(0, 0)))
		return "", "", transient
	}, 20*time.Second, now, oversleep)
	if !errors.Is(err, errGenerationCensusLate) {
		t.Fatalf("overrun wait = %v, want a late refusal", err)
	}
	for _, s := range starts {
		if s >= 20*time.Second {
			t.Fatalf("a pass started at %s, past the window (%v)", s, starts)
		}
	}
}

func TestRecoverCensusDeadlineReachesTheInspection(t *testing.T) {
	// On the real clock: the pass's context ends at the deadline, and a pass
	// that only returns when it does is refused, promptly.
	start := time.Now()
	var deadline time.Time
	var has bool
	_, _, err := generationCensusRetryFor(t.Context(), func(ctx context.Context) (string, string, error) {
		deadline, has = ctx.Deadline()
		select {
		case <-ctx.Done():
		case <-time.After(2 * time.Second):
		}
		return "boot", "digest", nil
	}, 50*time.Millisecond, time.Now, generationSleep)
	if !errors.Is(err, errGenerationCensusLate) {
		t.Fatalf("a pass blocked to the deadline = %v, want a late refusal", err)
	}
	if !has || deadline.Sub(start) > 50*time.Millisecond+10*time.Millisecond {
		t.Fatalf("pass context deadline = %v (set %t), want the census's", deadline.Sub(start), has)
	}
	if took := time.Since(start); took > 2*time.Second {
		t.Fatalf("census took %s past a 50 ms window", took)
	}
	// The lookup's own blocking reads get that context: systemctl show and
	// the image hash.
	f := deletedExe(t)
	var showDeadline bool
	show := f.env.show
	f.env.show = func(ctx context.Context, unit string) (map[string]string, error) {
		_, showDeadline = ctx.Deadline()
		return show(ctx, unit)
	}
	if _, _, err := generationCensusRetryFor(t.Context(), func(ctx context.Context) (string, string, error) {
		return "boot", "digest", f.kernelCensus(ctx)
	}, 20*time.Second, time.Now, generationSleep); err != nil {
		t.Fatal(err)
	}
	if !showDeadline {
		t.Fatal("systemctl show ran without the census deadline")
	}
	// The deadline passes after the lookup has started on the task: the
	// image hash stops.
	ctx, cancel := context.WithCancel(t.Context())
	f.env.show = func(ctx context.Context, unit string) (map[string]string, error) {
		cancel()
		return show(ctx, unit)
	}
	if err := f.kernelCensus(ctx); !errors.Is(err, context.Canceled) || !strings.Contains(err.Error(), "deleted exe check refused") {
		t.Fatalf("image hash under an ended context = %v", err)
	}
}

// guardedEdge makes the fixture's unit the guarded carlos-edge.service:
// ExecStart is the guard, the selected program <opt>/carlos, and the task
// is the edge's main process in its supervise/ group.
func (f *deletedExeFixture) guardedEdge(t *testing.T) {
	t.Helper()
	f.unit = "carlos-edge.service"
	f.execStart = "{ path=/opt/carlos/system/guard/carlos ; argv[]=/opt/carlos/system/guard/carlos system guard-exec edge ; ignore_errors=no ; start_time=[n/a] ; stop_time=[n/a] ; pid=41 ; code=(null) ; status=0/0 }"
	binary := f.env.binary
	f.env.binary = func(ctx context.Context, opt, path string) (rootFloorBinary, error) {
		if path == "/opt/carlos/system/guard/carlos" {
			return rootFloorBinary{Digest: "guard"}, nil
		}
		b, err := binary(ctx, opt, path)
		b.File.Inode = 7
		return b, err
	}
	base := "/system.slice/carlos-edge.service"
	generationWrite(t, f.task+"/cgroup", "0::"+base+"/supervise\n")
	generationWrite(t, f.proc+"/41/cgroup", "0::"+base+"/supervise\n")
	show := f.env.show
	f.env.show = func(ctx context.Context, unit string) (map[string]string, error) {
		props, err := show(ctx, unit)
		if err == nil {
			props["ControlGroup"] = base
		}
		return props, err
	}
}

func TestRecoverCensusDeletedExeOfARootTask(t *testing.T) {
	// The deleted-exe acceptance also covers a root task of a unit with a
	// selected image (the guarded edge), and only that among root tasks.
	// The task runs in a private mount namespace, which the strict root
	// scan refuses: only a trusted classification lets it pass.
	root := func(t *testing.T, f *deletedExeFixture) {
		generationWrite(t, f.task+"/status", "Tgid:\t41\nUid:\t0\t0\t0\t0\nKthread:\t0\n")
		generationRelink(t, "private-service", f.task+"/ns/mnt")
	}
	t.Run("selected-image-signed", func(t *testing.T) {
		f := deletedExe(t)
		f.guardedEdge(t)
		root(t, f)
		if err := f.census(t); err != nil {
			t.Fatalf("census refused a root guarded-edge task on its signed deleted image: %v", err)
		}
	})
	t.Run("selected-image-unsigned", func(t *testing.T) {
		f := deletedExe(t)
		f.guardedEdge(t)
		root(t, f)
		f.writeImage(t, "an image no identity signs", 0o755, true)
		if err := f.census(t); err == nil {
			t.Fatal("census accepted a root guarded-edge task on an unsigned deleted image")
		}
	})
	t.Run("no-selected-image", func(t *testing.T) {
		// A root task of a unit with no selected image keeps main's path
		// walk, which a deleted exe fails, so the strict root scan applies.
		f := deletedExe(t)
		root(t, f)
		if err := f.census(t); err == nil {
			t.Fatal("census trusted a root task of an unselected unit on its deleted exe")
		}
	})
	t.Run("selected-image-exec-race", func(t *testing.T) {
		f := deletedExe(t)
		f.guardedEdge(t)
		root(t, f)
		err := f.execRace(t, nil, func(_ bool, call int) {
			if call == 1 {
				generationRelink(t, f.unrelated, f.task+"/exe")
			}
		})
		if err == nil || !strings.Contains(err.Error(), "exe binding check refused") {
			t.Fatalf("root guarded-edge task that exec'd away = %v", err)
		}
	})
}

// slowImage returns one byte per read and ends ctx after the first.
type slowImage struct {
	cancel func()
	reads  int
}

func (s *slowImage) Read(p []byte) (int, error) {
	s.reads++
	if s.reads == 1 {
		s.cancel()
	}
	if len(p) == 0 {
		return 0, nil
	}
	p[0] = 'x'
	return 1, nil
}

func TestSignedBinaryProofStopsWhenTheCensusDeadlineEnds(t *testing.T) {
	// The hash: a slow image whose read outlives the deadline.
	ctx, cancel := context.WithCancel(t.Context())
	image := &slowImage{cancel: cancel}
	if _, err := rootFloorDigest(ctx, image); !errors.Is(err, context.Canceled) || image.reads > 2 {
		t.Fatalf("digest = %v after %d reads, want it stopped at the deadline", err, image.reads)
	}
	// The identity scan: once the deadline has passed it reads no identity.
	reads := 0
	ctx, cancel = context.WithCancel(t.Context())
	cancel()
	if _, err := rootFloorSignedIdentity(ctx, nil, rootFloorBinary{Digest: "d"}, []string{"a.json"}, func(string) ([]byte, error) { reads++; return nil, nil }); !errors.Is(err, context.Canceled) || reads != 0 {
		t.Fatalf("identity scan after the deadline = %v after %d reads", err, reads)
	}
}

func TestRecoverCensusDeadlineReachesTheSignedBinaryProof(t *testing.T) {
	// The lookup hands the census's deadline context to the signed-binary
	// proof, for the deleted image and for the guarded edge's guard and
	// selected images: a proof that only returns when that context ends
	// is cut off at the deadline.
	for _, guarded := range []bool{false, true} {
		t.Run(map[bool]string{false: "deleted-image", true: "guarded-edge"}[guarded], func(t *testing.T) {
			f := deletedExe(t)
			if guarded {
				f.guardedEdge(t)
			}
			proofs := 0
			f.env.binary = func(ctx context.Context, opt, path string) (rootFloorBinary, error) {
				proofs++
				if _, ok := ctx.Deadline(); !ok {
					t.Error("signed-binary proof ran without the census deadline")
				}
				select {
				case <-ctx.Done():
					return rootFloorBinary{}, ctx.Err()
				case <-time.After(5 * time.Second):
					return rootFloorBinary{}, errors.New("proof outlived the deadline")
				}
			}
			start := time.Now()
			_, _, err := generationCensusRetryFor(t.Context(), func(ctx context.Context) (string, string, error) {
				return "boot", "digest", f.kernelCensus(ctx)
			}, 100*time.Millisecond, time.Now, generationSleep)
			if !errors.Is(err, errGenerationCensusLate) {
				t.Fatalf("census = %v, want a late refusal", err)
			}
			if took := time.Since(start); took > 2*time.Second {
				t.Fatalf("census took %s past a 100 ms window: the proof ignored its context", took)
			}
			if proofs == 0 {
				t.Fatal("the proof never ran")
			}
		})
	}
}