//go:build linux

package instanceexec

import (
	"fmt"
	"os"
	"os/signal"
	"runtime"
	"syscall"

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

// The trampoline stays behind as the namespace's init (#126, found on the
// dev fleet). The templates run with PrivatePIDs=yes, so the unit's main
// process is PID 1 of its own PID namespace, and the kernel gives a
// namespace init no default signal actions: a signal it has no handler
// for is dropped, whoever sends it (only SIGKILL from outside gets
// through). A tenant exec'd into that slot that never asked for SIGTERM,
// which is most Go servers, cannot be stopped by it: Go's runtime tries
// to die by the signal, the kernel drops it, and the runtime exits 2
// instead; a tenant with no handler at all waits out TimeoutStopSec for
// the SIGKILL. Every hibernate, restart, move and deploy then left the
// unit failed. A flat tenant was never PID 1.
//
// So the trampoline forks the tenant as its child, which is not PID 1 and
// keeps its default actions exactly as on the flat layout, and does the
// two things PID 1 owes the namespace:
//
//   - Forward. KillMode=mixed sends the stop's SIGTERM to the main
//     process, which is this one; it and the other signals in forwarded
//     go on to the tenant, which acts on them as it would have as the
//     main process. SIGKILL needs no forwarding: after TimeoutStopSec
//     systemd sends it to the whole cgroup, and SIGKILL to a namespace
//     init from outside always kills it.
//   - Reap. A process whose parent dies is reparented to the namespace's
//     init, and stays a zombie until init waits for it. The tenant's own
//     children that outlive their parent land here.
//
// When the tenant exits, the init exits with its status: the tenant's
// exit code, or 128+n for a death by signal n (the shell's convention; a
// namespace init cannot die by a signal it sends itself). The run and
// sidecar templates count 129, 130, 141 and 143 as success, the signals
// systemd itself counts as a clean death of a main process. The kernel
// then kills whatever is left in the namespace, as it did when the tenant
// was PID 1.
//
// The init is small on purpose and allocates nothing per signal: one
// buffered channel, one WaitStatus, a kill or a wait4 per signal.

// forwarded are the signals the init passes on to the tenant.
var forwarded = [...]os.Signal{syscall.SIGTERM, syscall.SIGINT, syscall.SIGHUP, syscall.SIGQUIT, syscall.SIGUSR1, syscall.SIGUSR2}

// scopedInit is the trampoline's launch: scope, start the tenant under
// the init, and exit with its status. It returns only if the tenant could
// not be started; a scope that cannot be applied refuses the start (the
// probe has already failed readiness on such a box, so no nested wake
// gets here).
func scopedInit(argv0 string, argv, env []string) error {
	return runInit(argv0, argv, env, ScopeAbstractUnix)
}

func runInit(argv0 string, argv, env []string, scope func() error) error {
	return runInitChecked(argv0, argv, env, scope, nil)
}

func runInitChecked(argv0 string, argv, env []string, scope, check func() error) error {
	// Before anything can start the tenant: a stop that arrives early is
	// queued and forwarded, not dropped by PID 1's missing defaults.
	sigs := make(chan os.Signal, 32)
	signal.Notify(sigs, append(forwarded[:], syscall.SIGCHLD)...)

	// Landlock restricts a thread, and a fork copies the forking thread's
	// domain, which execve keeps: scoping this locked thread and forking
	// from it scopes the tenant. The thread stays locked until exit.
	runtime.LockOSThread()
	if err := scope(); err != nil {
		return fmt.Errorf("not starting without the abstract socket scope: %w", err)
	}
	// PID 1 in production already reaps orphans; this makes the init
	// their reaper wherever it runs.
	if err := unix.Prctl(unix.PR_SET_CHILD_SUBREAPER, 1, 0, 0, 0); err != nil {
		return fmt.Errorf("subreaper: %w", err)
	}
	// Same argv, env, cwd and stdio as the exec this replaces. Fds beyond
	// stdio pass as they did there: anything not close-on-exec stays
	// open across fork and exec. Setsid gives the tenant the session of
	// its own that systemd gave the main process.
	if check != nil {
		if err := check(); err != nil {
			return err
		}
	}
	tenant, err := syscall.ForkExec(argv0, argv, &syscall.ProcAttr{
		Env:   env,
		Files: []uintptr{0, 1, 2},
		Sys:   &syscall.SysProcAttr{Setsid: true},
	})
	if err != nil {
		return err
	}
	os.Exit(superviseTenant(tenant, sigs))
	return nil
}

// superviseTenant forwards signals to the tenant and reaps every child
// until the tenant itself is reaped, and returns the init's exit code.
func superviseTenant(tenant int, sigs <-chan os.Signal) int {
	var ws syscall.WaitStatus
	for s := range sigs {
		sig := s.(syscall.Signal)
		if sig != syscall.SIGCHLD {
			_ = syscall.Kill(tenant, sig)
			continue
		}
		// One SIGCHLD can stand for several exits: drain them all.
		for {
			pid, err := syscall.Wait4(-1, &ws, syscall.WNOHANG, nil)
			if err == syscall.EINTR {
				continue
			}
			if err != nil || pid <= 0 {
				break
			}
			if pid == tenant {
				return initExitCode(ws)
			}
		}
	}
	return 1 // unreachable: signal.Notify never closes the channel
}

// initExitCode maps the tenant's wait status to the init's exit code.
func initExitCode(ws syscall.WaitStatus) int {
	if ws.Signaled() {
		return 128 + int(ws.Signal())
	}
	return ws.ExitStatus()
}