agent/agentexec/cli_linux_test.go

Rama main · réplica de lectura del repositorio original.

//go:build linux
// +build linux

package agentexec_test

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"slices"
	"strconv"
	"strings"
	"syscall"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"golang.org/x/sys/unix"
	"golang.org/x/xerrors"

	"github.com/coder/coder/v2/agent/agentexec"
	"github.com/coder/coder/v2/testutil"
)

//nolint:paralleltest // This test is sensitive to environment variables
func TestCLI(t *testing.T) {
	t.Run("OK", func(t *testing.T) {
		ctx := testutil.Context(t, testutil.WaitMedium)
		cmd, path := cmd(ctx, t, 123, 12)
		err := cmd.Start()
		require.NoError(t, err)
		go cmd.Wait()

		waitForSentinel(ctx, t, cmd, path)
		requireOOMScore(t, cmd.Process.Pid, 123)
		requireNiceScore(t, cmd.Process.Pid, 12)
	})

	t.Run("FiltersEnv", func(t *testing.T) {
		ctx := testutil.Context(t, testutil.WaitMedium)
		cmd, path := cmd(ctx, t, 123, 12)
		cmd.Env = append(cmd.Env, fmt.Sprintf("%s=true", agentexec.EnvProcPrioMgmt))
		cmd.Env = append(cmd.Env, fmt.Sprintf("%s=123", agentexec.EnvProcOOMScore))
		cmd.Env = append(cmd.Env, fmt.Sprintf("%s=12", agentexec.EnvProcNiceScore))
		// Ensure unrelated environment variables are preserved.
		cmd.Env = append(cmd.Env, "CODER_TEST_ME_AGENTEXEC=true")
		err := cmd.Start()
		require.NoError(t, err)
		go cmd.Wait()
		waitForSentinel(ctx, t, cmd, path)

		env := procEnv(t, cmd.Process.Pid)
		hasExecEnvs := slices.ContainsFunc(
			env,
			func(e string) bool {
				return strings.HasPrefix(e, agentexec.EnvProcPrioMgmt) ||
					strings.HasPrefix(e, agentexec.EnvProcOOMScore) ||
					strings.HasPrefix(e, agentexec.EnvProcNiceScore)
			})
		require.False(t, hasExecEnvs, "expected environment variables to be filtered")
		userEnv := slices.Contains(env, "CODER_TEST_ME_AGENTEXEC=true")
		require.True(t, userEnv, "expected user environment variables to be preserved")
	})

	t.Run("Defaults", func(t *testing.T) {
		ctx := testutil.Context(t, testutil.WaitMedium)
		cmd, path := cmd(ctx, t, 0, 0)
		err := cmd.Start()
		require.NoError(t, err)
		go cmd.Wait()

		waitForSentinel(ctx, t, cmd, path)

		expectedNice := expectedNiceScore(t)
		expectedOOM := expectedOOMScore(t)
		requireOOMScore(t, cmd.Process.Pid, expectedOOM)
		requireNiceScore(t, cmd.Process.Pid, expectedNice)
	})

	t.Run("Capabilities", func(t *testing.T) {
		testdir := filepath.Dir(TestBin)
		capDir := filepath.Join(testdir, "caps")
		err := os.Mkdir(capDir, 0o755)
		require.NoError(t, err)
		bin := buildBinary(capDir)
		// Try to set capabilities on the binary. This should work fine in CI but
		// it's possible some developers may be working in an environment where they don't have the necessary permissions.
		err = setCaps(t, bin, "cap_net_admin")
		if os.Getenv("CI") != "" {
			require.NoError(t, err)
		} else if err != nil {
			t.Skipf("unable to set capabilities for test: %v", err)
		}
		ctx := testutil.Context(t, testutil.WaitMedium)
		cmd, path := binCmd(ctx, t, bin, 123, 12)
		err = cmd.Start()
		require.NoError(t, err)
		go cmd.Wait()

		waitForSentinel(ctx, t, cmd, path)
		// This is what we're really testing, a binary with added capabilities requires setting dumpable.
		requireOOMScore(t, cmd.Process.Pid, 123)
		requireNiceScore(t, cmd.Process.Pid, 12)
	})
}

func requireNiceScore(t *testing.T, pid int, score int) {
	t.Helper()

	nice, err := unix.Getpriority(unix.PRIO_PROCESS, pid)
	require.NoError(t, err)
	// See https://linux.die.net/man/2/setpriority#Notes
	require.Equal(t, score, 20-nice)
}

func requireOOMScore(t *testing.T, pid int, expected int) {
	t.Helper()

	actual, err := os.ReadFile(fmt.Sprintf("/proc/%d/oom_score_adj", pid))
	require.NoError(t, err)
	score := strings.TrimSpace(string(actual))
	require.Equal(t, strconv.Itoa(expected), score)
}

func waitForSentinel(ctx context.Context, t *testing.T, cmd *exec.Cmd, path string) {
	t.Helper()

	ticker := time.NewTicker(testutil.IntervalFast)
	defer ticker.Stop()

	// RequireEventually doesn't work well with require.NoError or similar require functions.
	for {
		err := cmd.Process.Signal(syscall.Signal(0))
		require.NoError(t, err)

		_, err = os.Stat(path)
		if err == nil {
			return
		}

		select {
		case <-ticker.C:
		case <-ctx.Done():
			require.NoError(t, ctx.Err())
		}
	}
}

func binCmd(ctx context.Context, t *testing.T, bin string, oom, nice int) (*exec.Cmd, string) {
	var (
		args = execArgs(oom, nice)
		dir  = t.TempDir()
		file = filepath.Join(dir, "sentinel")
	)

	args = append(args, "sh", "-c", fmt.Sprintf("touch %s && sleep 10m", file))
	//nolint:gosec
	cmd := exec.CommandContext(ctx, bin, args...)

	// We set this so we can also easily kill the sleep process the shell spawns.
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Setpgid: true,
	}

	cmd.Env = os.Environ()
	var buf bytes.Buffer
	cmd.Stdout = &buf
	cmd.Stderr = &buf
	t.Cleanup(func() {
		// Print output of a command if the test fails.
		if t.Failed() {
			t.Logf("cmd %q output: %s", cmd.Args, buf.String())
		}
		if cmd.Process != nil {
			// We use -cmd.Process.Pid to kill the whole process group.
			_ = syscall.Kill(-cmd.Process.Pid, syscall.SIGINT)
		}
	})
	return cmd, file
}

func cmd(ctx context.Context, t *testing.T, oom, nice int) (*exec.Cmd, string) {
	return binCmd(ctx, t, TestBin, oom, nice)
}

func expectedOOMScore(t *testing.T) int {
	t.Helper()

	score, err := os.ReadFile(fmt.Sprintf("/proc/%d/oom_score_adj", os.Getpid()))
	require.NoError(t, err)

	scoreInt, err := strconv.Atoi(strings.TrimSpace(string(score)))
	require.NoError(t, err)

	if scoreInt < 0 {
		return 0
	}
	if scoreInt >= 998 {
		return 1000
	}
	return 998
}

// procEnv returns the environment variables for a given process.
func procEnv(t *testing.T, pid int) []string {
	t.Helper()

	env, err := os.ReadFile(fmt.Sprintf("/proc/%d/environ", pid))
	require.NoError(t, err)
	return strings.Split(string(env), "\x00")
}

func expectedNiceScore(t *testing.T) int {
	t.Helper()

	score, err := unix.Getpriority(unix.PRIO_PROCESS, os.Getpid())
	require.NoError(t, err)

	// Priority is niceness + 20.
	score = 20 - score
	score += 5
	if score > 19 {
		return 19
	}
	return score
}

func execArgs(oom int, nice int) []string {
	execArgs := []string{"agent-exec"}
	if oom != 0 {
		execArgs = append(execArgs, fmt.Sprintf("--coder-oom=%d", oom))
	}
	if nice != 0 {
		execArgs = append(execArgs, fmt.Sprintf("--coder-nice=%d", nice))
	}
	execArgs = append(execArgs, "--")
	return execArgs
}

func setCaps(t *testing.T, bin string, caps ...string) error {
	t.Helper()

	setcap := fmt.Sprintf("sudo -n setcap %s=ep %s", strings.Join(caps, ", "), bin)
	out, err := exec.CommandContext(context.Background(), "sh", "-c", setcap).CombinedOutput()
	if err != nil {
		return xerrors.Errorf("setcap %q (%s): %w", setcap, out, err)
	}
	return nil
}