//go:build linux

package instanceexec

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"testing"
	"time"

	"golang.org/x/sys/unix"
)

// The init tests run the real runInit in a child process: this test
// binary, re-executed with ieHelperEnv naming a role. The "init" role
// runs runInit with the binary itself as the tenant, in the tenant role
// named by ieTenantEnv; every other role is a tenant.
const (
	ieHelperEnv = "CARLOS_INSTANCEEXEC_HELPER"
	ieTenantEnv = "CARLOS_INSTANCEEXEC_TENANT"
	ieScopeEnv  = "CARLOS_INSTANCEEXEC_SCOPE"
	ieArgEnv    = "CARLOS_INSTANCEEXEC_ARG"
)

func TestMain(m *testing.M) {
	if role := os.Getenv(ieHelperEnv); role != "" {
		os.Exit(runHelper(role))
	}
	os.Exit(m.Run())
}

// withRole is the environment with the helper role replaced, not
// appended: a Go child keeps the FIRST of duplicate keys (design §6.2).
func withRole(role string) []string {
	env := []string{ieHelperEnv + "=" + role}
	for _, kv := range os.Environ() {
		if !strings.HasPrefix(kv, ieHelperEnv+"=") {
			env = append(env, kv)
		}
	}
	return env
}

func runHelper(role string) int {
	self, err := os.Executable()
	if err != nil {
		fmt.Println("executable:", err)
		return 70
	}
	switch role {
	case "init":
		env := withRole(os.Getenv(ieTenantEnv))
		scope := func() error { return nil }
		if os.Getenv(ieScopeEnv) == "1" {
			scope = ScopeAbstractUnix
		}
		err := runInit(self, []string{self}, env, scope)
		fmt.Println("init:", err)
		return 71
	case "no-handler":
		// A tenant that never asks for SIGTERM, like most Go servers.
		fmt.Printf("ready pid=%d ppid=%d\n", os.Getpid(), os.Getppid())
		time.Sleep(30 * time.Second)
		return 0
	case "echo-signals":
		// Reports every forwarded signal; leaves on SIGTERM.
		got := make(chan os.Signal, 8)
		signal.Notify(got, forwarded[:]...)
		fmt.Println("ready")
		for s := range got {
			fmt.Println("got", s)
			if s == syscall.SIGTERM {
				return 0
			}
		}
		return 0
	case "exit":
		n, _ := strconv.Atoi(os.Getenv(ieArgEnv))
		return n
	case "kill-self":
		_ = syscall.Kill(os.Getpid(), syscall.SIGKILL)
		time.Sleep(30 * time.Second)
		return 0
	case "orphan":
		// Leave a grandchild behind: the intermediate starts it and exits
		// at once, so it is reparented to the nearest reaper.
		cmd := exec.Command(self)
		cmd.Env = withRole("intermediate")
		cmd.Stdout = os.Stdout
		if err := cmd.Run(); err != nil {
			fmt.Println("intermediate:", err)
			return 72
		}
		time.Sleep(30 * time.Second)
		return 0
	case "intermediate":
		cmd := exec.Command(self)
		cmd.Env = withRole("short-lived")
		if err := cmd.Start(); err != nil {
			fmt.Println("grandchild:", err)
			return 73
		}
		fmt.Printf("grandchild %d\n", cmd.Process.Pid)
		return 0
	case "short-lived":
		time.Sleep(1500 * time.Millisecond)
		return 0
	case "inherit":
		cwd, _ := os.Getwd()
		sid, _ := unix.Getsid(0)
		fmt.Printf("cwd=%s arg=%s sid-is-pid=%t\n", cwd, os.Getenv(ieArgEnv), sid == os.Getpid())
		f := os.NewFile(3, "fd3")
		if _, err := f.Write([]byte("fd3 reached\n")); err != nil {
			fmt.Println("fd3:", err)
		}
		return 0
	case "dial":
		c, err := net.Dial("unix", os.Getenv(ieArgEnv))
		switch {
		case err == nil:
			c.Close()
			fmt.Println("dial: connected")
		case errors.Is(unwrapErrno(err), syscall.EPERM):
			fmt.Println("dial: EPERM")
		default:
			fmt.Println("dial:", err)
		}
		return 0
	}
	fmt.Println("unknown helper role", role)
	return 74
}

// startInit starts the init helper with the given tenant role and extra
// environment, and returns it with a line reader on its stdout (which is
// the tenant's stdout too: the init passes fd 1 through).
func startInit(t *testing.T, tenant string, attr *syscall.SysProcAttr, extra ...string) (*exec.Cmd, *bufio.Reader) {
	t.Helper()
	self, err := os.Executable()
	if err != nil {
		t.Fatal(err)
	}
	cmd := exec.Command(self)
	cmd.Env = append(os.Environ(), append([]string{ieHelperEnv + "=init", ieTenantEnv + "=" + tenant}, extra...)...)
	cmd.Stderr = os.Stderr
	cmd.SysProcAttr = attr
	out, err := cmd.StdoutPipe()
	if err != nil {
		t.Fatal(err)
	}
	if err := cmd.Start(); err != nil {
		if attr != nil {
			t.Skipf("no unprivileged user and PID namespace here: %v", err)
		}
		t.Fatal(err)
	}
	t.Cleanup(func() {
		// Through the init first, so the tenant stops with it outside a
		// PID namespace too; the SIGKILL is for a test that failed early.
		_ = cmd.Process.Signal(syscall.SIGTERM)
		time.Sleep(100 * time.Millisecond)
		_ = cmd.Process.Kill()
		_ = cmd.Wait()
	})
	return cmd, bufio.NewReader(out)
}

// readLine reads one line within d, failing the test on a timeout.
func readLine(t *testing.T, r *bufio.Reader, d time.Duration) string {
	t.Helper()
	got := make(chan string, 1)
	go func() {
		line, err := r.ReadString('\n')
		if err != nil && line == "" {
			line = "read: " + err.Error()
		}
		got <- strings.TrimSpace(line)
	}()
	select {
	case line := <-got:
		return line
	case <-time.After(d):
		t.Fatalf("no line from the helper within %v", d)
		return ""
	}
}

// waitExit waits up to d for cmd to exit and returns its exit code, or -1
// with the signal if it died by one.
func waitExit(t *testing.T, cmd *exec.Cmd, d time.Duration) (int, syscall.Signal) {
	t.Helper()
	done := make(chan error, 1)
	go func() { done <- cmd.Wait() }()
	select {
	case <-done:
	case <-time.After(d):
		t.Fatalf("helper still running %v after the signal", d)
	}
	ws := cmd.ProcessState.Sys().(syscall.WaitStatus)
	if ws.Signaled() {
		return -1, ws.Signal()
	}
	return ws.ExitStatus(), 0
}

func TestInitExitStatusMapping(t *testing.T) {
	for _, c := range []struct {
		ws   syscall.WaitStatus
		want int
	}{
		{0, 0},
		{1 << 8, 1},
		{2 << 8, 2},
		{255 << 8, 255},
		{syscall.WaitStatus(syscall.SIGTERM), 143},
		{syscall.WaitStatus(syscall.SIGKILL), 137},
		{syscall.WaitStatus(syscall.SIGINT), 130},
		{syscall.WaitStatus(syscall.SIGSEGV) | 0x80, 139}, // core dumped
	} {
		if got := initExitCode(c.ws); got != c.want {
			t.Errorf("initExitCode(%#x) = %d, want %d", uint32(c.ws), got, c.want)
		}
	}
}

// End to end through the init, outside any namespace: the tenant's own
// exit code passes through, and a death by signal is 128+n.
func TestInitExitsWithTheTenantsStatus(t *testing.T) {
	cmd, _ := startInit(t, "exit", nil, ieArgEnv+"=7")
	if code, sig := waitExit(t, cmd, 10*time.Second); code != 7 {
		t.Errorf("init exit = %d (signal %v), want 7", code, sig)
	}
	cmd, _ = startInit(t, "kill-self", nil)
	if code, sig := waitExit(t, cmd, 10*time.Second); code != 128+int(syscall.SIGKILL) {
		t.Errorf("init exit = %d (signal %v), want 137", code, sig)
	}
}

// Every forwarded signal reaches the tenant.
func TestInitForwardsEachSignal(t *testing.T) {
	cmd, out := startInit(t, "echo-signals", nil)
	if line := readLine(t, out, 10*time.Second); line != "ready" {
		t.Fatalf("tenant said %q", line)
	}
	// SIGTERM last: the tenant leaves on it.
	for _, sig := range []syscall.Signal{syscall.SIGHUP, syscall.SIGINT, syscall.SIGQUIT, syscall.SIGUSR1, syscall.SIGUSR2, syscall.SIGTERM} {
		if err := cmd.Process.Signal(sig); err != nil {
			t.Fatal(err)
		}
		if line := readLine(t, out, 5*time.Second); line != "got "+sig.String() {
			t.Errorf("sent %v, tenant said %q", sig, line)
		}
	}
	if code, sig := waitExit(t, cmd, 5*time.Second); code != 0 {
		t.Errorf("init exit = %d (signal %v), want the tenant's 0", code, sig)
	}
}

// A tenant with no handlers acts on a forwarded signal by its default
// disposition, as the main process did before PrivatePIDs: it dies by
// the signal, and the init reports 128+n.
func TestInitStopsATenantWithNoHandler(t *testing.T) {
	for _, sig := range []syscall.Signal{syscall.SIGTERM, syscall.SIGINT, syscall.SIGHUP} {
		t.Run(sig.String(), func(t *testing.T) {
			cmd, out := startInit(t, "no-handler", nil)
			if line := readLine(t, out, 10*time.Second); !strings.HasPrefix(line, "ready ") {
				t.Fatalf("tenant said %q", line)
			}
			if err := cmd.Process.Signal(sig); err != nil {
				t.Fatal(err)
			}
			if code, got := waitExit(t, cmd, 5*time.Second); code != 128+int(sig) {
				t.Errorf("init exit = %d (signal %v), want %d", code, got, 128+int(sig))
			}
		})
	}
}

// The tenant gets the init's cwd, the launch env, a session of its own
// (as the unit's main process had) and any inherited fd beyond stdio.
func TestInitPassesTheTenantItsEnvironment(t *testing.T) {
	r, w, err := os.Pipe()
	if err != nil {
		t.Fatal(err)
	}
	defer r.Close()
	self, _ := os.Executable()
	cmd := exec.Command(self)
	cmd.Dir = t.TempDir()
	cmd.Env = append(os.Environ(), ieHelperEnv+"=init", ieTenantEnv+"=inherit", ieArgEnv+"=launch-value")
	cmd.ExtraFiles = []*os.File{w}
	var stdout strings.Builder
	cmd.Stdout = &stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		t.Fatalf("init: %v; output %q", err, stdout.String())
	}
	w.Close()
	want := fmt.Sprintf("cwd=%s arg=launch-value sid-is-pid=true", cmd.Dir)
	if got := strings.TrimSpace(stdout.String()); got != want {
		t.Errorf("tenant said %q, want %q", got, want)
	}
	b, _ := io.ReadAll(r)
	if string(b) != "fd3 reached\n" {
		t.Errorf("fd 3 carried %q", b)
	}
}

// A process orphaned inside the namespace is the init's to reap. The
// grandchild must be reparented to the init (it is a subreaper, and PID 1
// in production) and must not stay a zombie after it exits.
func TestInitReapsOrphans(t *testing.T) {
	cmd, out := startInit(t, "orphan", nil)
	line := readLine(t, out, 10*time.Second)
	var gpid int
	if _, err := fmt.Sscanf(line, "grandchild %d", &gpid); err != nil {
		t.Fatalf("tenant said %q", line)
	}
	initPid := cmd.Process.Pid
	adopted := false
	deadline := time.Now().Add(10 * time.Second)
	for time.Now().Before(deadline) {
		state, ppid, ok := procStat(gpid)
		if !ok {
			break // reaped
		}
		if ppid == initPid {
			adopted = true
		}
		if state == "Z" && time.Now().After(deadline.Add(-7*time.Second)) {
			t.Fatalf("grandchild %d is a zombie of %d and nobody reaps it", gpid, ppid)
		}
		time.Sleep(20 * time.Millisecond)
	}
	if _, _, ok := procStat(gpid); ok {
		t.Fatalf("grandchild %d still present", gpid)
	}
	if !adopted {
		t.Fatalf("grandchild %d was never the init's (%d) child", gpid, initPid)
	}
	// The tenant itself is still running: reaping an orphan is not the
	// tenant's exit.
	if err := cmd.Process.Signal(syscall.Signal(0)); err != nil {
		t.Fatalf("init gone after reaping an orphan: %v", err)
	}
}

// procStat returns the state and parent of pid from /proc/<pid>/stat.
func procStat(pid int) (state string, ppid int, ok bool) {
	b, err := os.ReadFile(fmt.Sprintf("/proc/%d/stat", pid))
	if err != nil {
		return "", 0, false
	}
	s := string(b)
	i := strings.LastIndexByte(s, ')')
	if i < 0 {
		return "", 0, false
	}
	f := strings.Fields(s[i+1:])
	if len(f) < 2 {
		return "", 0, false
	}
	ppid, _ = strconv.Atoi(f[1])
	return f[0], ppid, true
}

// The Landlock scope applied by the init holds in the tenant: the scope is
// on the thread that forks, and a domain is inherited across fork and
// exec. The unscoped run is the control, proving the dial can succeed.
func TestInitTenantKeepsTheAbstractSocketScope(t *testing.T) {
	abi, _, e := unix.Syscall(unix.SYS_LANDLOCK_CREATE_RULESET, 0, 0, unix.LANDLOCK_CREATE_RULESET_VERSION)
	if e != 0 || abi < landlockScopeABI {
		t.Skipf("no Landlock scope on this kernel (abi %d, %v)", abi, e)
	}
	name := "@carlos-init-test-" + randomName()
	ln, err := net.Listen("unix", name)
	if err != nil {
		t.Skipf("abstract listen: %v", err)
	}
	defer ln.Close()
	go func() {
		for {
			c, err := ln.Accept()
			if err != nil {
				return
			}
			c.Close()
		}
	}()
	for scope, want := range map[string]string{"0": "dial: connected", "1": "dial: EPERM"} {
		cmd, out := startInit(t, "dial", nil, ieArgEnv+"="+name, ieScopeEnv+"="+scope)
		if got := readLine(t, out, 10*time.Second); got != want {
			t.Errorf("scope=%s: tenant said %q, want %q", scope, got, want)
		}
		if code, sig := waitExit(t, cmd, 10*time.Second); code != 0 {
			t.Errorf("scope=%s: init exit = %d (signal %v)", scope, code, sig)
		}
	}
}

// The finding itself (#126, dev fleet): in a PID namespace, as under
// PrivatePIDs=yes, the first process is init and has no default signal
// actions. A Go tenant with no SIGTERM handler in that slot does not die
// by SIGTERM. Through the init it is the init's child and does, promptly. Needs an
// unprivileged user namespace; skipped where there is none.
func TestInitStopsATenantWithNoHandlerInAPIDNamespace(t *testing.T) {
	userPIDNS := func() *syscall.SysProcAttr {
		return &syscall.SysProcAttr{
			Cloneflags:  syscall.CLONE_NEWUSER | syscall.CLONE_NEWPID,
			UidMappings: []syscall.SysProcIDMap{{ContainerID: os.Getuid(), HostID: os.Getuid(), Size: 1}},
			GidMappings: []syscall.SysProcIDMap{{ContainerID: os.Getgid(), HostID: os.Getgid(), Size: 1}},
		}
	}

	t.Run("tenant as PID 1", func(t *testing.T) {
		self, _ := os.Executable()
		cmd := exec.Command(self)
		cmd.Env = append(os.Environ(), ieHelperEnv+"=no-handler")
		cmd.SysProcAttr = userPIDNS()
		cmd.Stderr = io.Discard // the runtime's own complaint, if any
		out, err := cmd.StdoutPipe()
		if err != nil {
			t.Fatal(err)
		}
		if err := cmd.Start(); err != nil {
			t.Skipf("no unprivileged user and PID namespace here: %v", err)
		}
		defer func() { _ = cmd.Process.Kill(); _ = cmd.Wait() }()
		if line := readLine(t, bufio.NewReader(out), 10*time.Second); !strings.HasPrefix(line, "ready pid=1 ") {
			t.Fatalf("tenant said %q, want it at PID 1", line)
		}
		_ = cmd.Process.Signal(syscall.SIGTERM)
		done := make(chan struct{})
		go func() { _ = cmd.Wait(); close(done) }()
		select {
		case <-done:
			ws := cmd.ProcessState.Sys().(syscall.WaitStatus)
			if ws.Signaled() && ws.Signal() == syscall.SIGTERM {
				t.Skip("this kernel delivered SIGTERM's default action to a namespace init")
			}
			t.Logf("as PID 1 the tenant did not die by SIGTERM: %v (the bug the init fixes)", cmd.ProcessState)
		case <-time.After(2 * time.Second):
			t.Logf("as PID 1 the tenant ignored SIGTERM (the bug the init fixes)")
		}
	})

	t.Run("through the init", func(t *testing.T) {
		cmd, out := startInit(t, "no-handler", userPIDNS())
		// Not PID 2 as such: the init's own threads take PIDs in the
		// namespace too. What matters is that the tenant is not PID 1.
		line := readLine(t, out, 10*time.Second)
		var pid, ppid int
		if _, err := fmt.Sscanf(line, "ready pid=%d ppid=%d", &pid, &ppid); err != nil || pid == 1 || ppid != 1 {
			t.Fatalf("tenant said %q, want a child of the init at PID 1", line)
		}
		start := time.Now()
		if err := cmd.Process.Signal(syscall.SIGTERM); err != nil {
			t.Fatal(err)
		}
		code, sig := waitExit(t, cmd, 5*time.Second)
		if code != 128+int(syscall.SIGTERM) {
			t.Errorf("init exit = %d (signal %v), want 143", code, sig)
		}
		t.Logf("stopped in %v", time.Since(start))
	})
}