helpers.go

/*
   Copyright The containerd Authors.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package shimtest

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"testing"
	"time"

	bootapi "github.com/containerd/containerd/api/runtime/bootstrap/v1"
	taskAPI "github.com/containerd/containerd/api/runtime/task/v3"
	"github.com/containerd/containerd/api/types"
	runcopt "github.com/containerd/containerd/api/types/runc/options"
	"github.com/containerd/containerd/v2/core/mount"
	"github.com/containerd/ttrpc"
	"github.com/containerd/typeurl/v2"
	"github.com/opencontainers/runtime-spec/specs-go"
	"google.golang.org/protobuf/encoding/protowire"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/anypb"
)

// uniqueTestNamespace returns a containerd namespace unique to this
// test invocation. The namespace is used to isolate each test's shim
// from others and acts as a key for zombie-shim attribution: because
// the shim is started with -namespace <ns>, the value is visible in
// /proc/<pid>/cmdline on Linux and can be reported when a leaked
// process is detected by registerShimLeakCheck.
//
// suite should be the short name of the calling suite (e.g. "run",
// "exec", "stress"). It is embedded in the namespace so that a leaked
// shim's cmdline immediately identifies both the suite and the
// specific container instance without needing a side-channel registry.
// The resulting format is "shimtest-<suite>-<random-hex>",
// e.g. "shimtest-stress-a1b2c3d4".
//
// For benchmarks the base shimtestNamespace constant is returned
// unchanged — benchmark iterations run in a tight loop and the
// per-iteration overhead of generating (and scanning for) unique
// namespaces is not worth the cost.
func uniqueTestNamespace(tb testing.TB, suite string) string {
	if _, isBench := tb.(*testing.B); isBench {
		return shimtestNamespace
	}
	if suite != "" {
		return shimtestNamespace + "-" + suite + "-" + randomSuffix()
	}
	return shimtestNamespace + "-" + randomSuffix()
}

// pidDiff returns the PIDs present in after but not in before.
func pidDiff(before, after map[int]struct{}) []int {
	var out []int
	for pid := range after {
		if _, ok := before[pid]; !ok {
			out = append(out, pid)
		}
	}
	return out
}

// registerShimLeakCheck snapshots the set of running shim processes
// at call time and registers a t.Cleanup that fires after all
// subtests (and their own cleanups) have completed. Any shim processes
// that are new relative to the snapshot are reported as leaks via
// t.Errorf, causing the suite-level test to fail.
//
// On Linux each leaked PID is annotated with shimCmdlineInfo so the
// -namespace and -id values (visible in /proc/<pid>/cmdline) appear
// in the error message. Because each test passes a namespace generated
// by uniqueTestNamespace and a container id generated by containerID,
// both of which embed a random suffix, the attribution string
// identifies the specific test invocation that created the zombie.
//
// Intended to be called once near the top of each suite's Run method.
// The single /proc scan at cleanup time is far cheaper than a per-test
// scan inside startShim, which would add a 500ms grace-period sleep
// for every iteration of a stress test.
func registerShimLeakCheck(t *testing.T, shimBin string) {
	t.Helper()
	before := shimPIDs(t, shimBin)
	t.Cleanup(func() {
		// Give the kernel a moment to reap any zombies left by the
		// last shutdown RPC before we scan.
		time.Sleep(500 * time.Millisecond)
		after := shimPIDs(t, shimBin)
		leaked := pidDiff(before, after)
		if len(leaked) == 0 {
			return
		}
		for _, pid := range leaked {
			info := shimCmdlineInfo(pid)
			if info != "" {
				t.Errorf("leaked shim process: PID %d (%s)", pid, info)
			} else {
				t.Errorf("leaked shim process: PID %d", pid)
			}
		}
	})
}

// containerID returns a safe container id derived from the test name,
// suitable for use as a path component. Both regular tests and
// benchmarks supply a name via tb.Name().
func containerID(tb testing.TB) string {
	tb.Helper()
	name := tb.Name()
	name = strings.NewReplacer("/", "-", " ", "-").Replace(name)
	if len(name) > 60 {
		name = name[:60]
	}
	return strings.ToLower(name) + "-" + randomSuffix()
}

// createOCISpec writes a minimal OCI spec config.json under
// bundleDir. Each opt is applied in order before the spec is
// written. When the process is rootless and the rootfs isn't being
// mounted by the shim, a user namespace + uid/gid mappings are added
// automatically.
func createOCISpec(tb testing.TB, bundleDir string, args []string, cfg Config, opts ...func(*specs.Spec)) {
	tb.Helper()

	spec := specs.Spec{
		Version: specs.Version,
		Root: &specs.Root{
			Path:     "rootfs",
			Readonly: false,
		},
		Process: &specs.Process{
			Args: args,
			Cwd:  "/",
			Env:  []string{"PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"},
		},
		Mounts: []specs.Mount{
			{Destination: "/proc", Type: "proc", Source: "proc"},
			{Destination: "/dev", Type: "tmpfs", Source: "tmpfs"},
		},
		Linux: &specs.Linux{
			Namespaces: []specs.LinuxNamespace{
				{Type: specs.MountNamespace},
			},
		},
	}

	if os.Getuid() != 0 && !cfg.FormatMounts {
		spec.Linux.Namespaces = append(spec.Linux.Namespaces,
			specs.LinuxNamespace{Type: specs.UserNamespace},
			specs.LinuxNamespace{Type: specs.PIDNamespace},
			specs.LinuxNamespace{Type: specs.NetworkNamespace},
		)
		spec.Linux.UIDMappings = []specs.LinuxIDMapping{
			{ContainerID: 0, HostID: uint32(os.Getuid()), Size: 1},
		}
		spec.Linux.GIDMappings = []specs.LinuxIDMapping{
			{ContainerID: 0, HostID: uint32(os.Getgid()), Size: 1},
		}
	}

	for _, opt := range opts {
		opt(&spec)
	}

	data, err := json.Marshal(spec)
	if err != nil {
		tb.Fatal("failed to marshal OCI spec:", err)
	}
	if err := os.WriteFile(filepath.Join(bundleDir, "config.json"), data, 0644); err != nil {
		tb.Fatal("failed to write config.json:", err)
	}
}

// withExtraMounts returns a CreateOCISpec opt that appends mounts to
// the spec.
func withExtraMounts(mounts ...specs.Mount) func(*specs.Spec) {
	return func(s *specs.Spec) {
		s.Mounts = append(s.Mounts, mounts...)
	}
}

// withMemoryLimit returns a CreateOCISpec opt that sets the memory
// limit (in bytes) on the spec, with swap clamped equal to the limit
// so the container cannot grow via swap before the OOM killer fires.
func withMemoryLimit(bytes int64) func(*specs.Spec) {
	return func(s *specs.Spec) {
		if s.Linux.Resources == nil {
			s.Linux.Resources = &specs.LinuxResources{}
		}
		s.Linux.Resources.Memory = &specs.LinuxMemory{
			Limit: &bytes,
			Swap:  &bytes,
		}
	}
}

// shimSetup resolves the shim binary, creates a bundle directory, and
// builds rootfs mounts from the embedded testbin. Returns the shim
// binary's absolute path, the bundle directory, and the rootfs
// mounts to pass to CreateTaskRequest.
func shimSetup(tb testing.TB, cfg Config) (shimBin, bundleDir string, rootfsMounts []*types.Mount) {
	tb.Helper()

	shimBin, err := exec.LookPath(cfg.ShimBinary)
	if err != nil {
		tb.Fatalf("shim binary %q not found in PATH: %v", cfg.ShimBinary, err)
	}

	// Ensure the shim binary's directory is in PATH so sibling
	// helpers (kernels, libraries) co-located with the shim resolve.
	shimDir := filepath.Dir(shimBin)
	if !strings.Contains(os.Getenv("PATH"), shimDir) {
		os.Setenv("PATH", shimDir+string(os.PathListSeparator)+os.Getenv("PATH"))
	}

	bundleDir = tb.TempDir()
	bundleDir, err = filepath.EvalSymlinks(bundleDir)
	if err != nil {
		tb.Fatal("failed to resolve bundle dir:", err)
	}

	rootfsDir := filepath.Join(bundleDir, "rootfs")
	if err := os.MkdirAll(rootfsDir, 0755); err != nil {
		tb.Fatal("failed to create rootfs dir:", err)
	}
	tb.Cleanup(func() { mount.Unmount(rootfsDir, 0) })

	rootfsMounts = buildEmbeddedRootfs(tb, bundleDir, cfg)

	// Always start an events recorder so the shim's event-publish connection
	// never blocks. On Linux, a missing events socket causes an immediate
	// "connection refused" which shims handle gracefully. On Windows, a missing
	// named-pipe server causes DialPipe to retry for several seconds (or
	// indefinitely), deadlocking the Create RPC. Binding the server here
	// ensures it exists before the shim needs it, regardless of which suite
	// method is running.
	startEventsRecorder(tb, bundleDir)

	return shimBin, bundleDir, rootfsMounts
}

// shortSocketPaths caches the events-socket path per bundleDir so that
// all callers within a single test get the same address. The concrete
// implementation of containerdSockPath (unix socket vs. named pipe) lives
// in connect_unix.go / connect_windows.go.
var shortSocketPaths sync.Map

// newCreateTaskRequest builds a CreateTaskRequest with runc Options.
// SystemdCgroup is forced off so root and rootless runs both use the
// cgroupfs manager — otherwise root on systemd hosts pays tens of ms
// of DBus round-trips per container, skewing benchmarks. Rootless
// (Linux non-root) paths additionally get IoUid/IoGid set and a writable
// Root. On Windows, where containers run inside a VM as root, none of
// the rootless tweaks apply (and os.Getuid returns -1, which would
// otherwise trigger the rootless branch and pass a Windows path as
// runc's --root, which then fails inside the Linux VM).
func newCreateTaskRequest(tb testing.TB, id, bundle, stdout, stderr string, rootfs []*types.Mount) *taskAPI.CreateTaskRequest {
	tb.Helper()
	req := &taskAPI.CreateTaskRequest{
		ID:     id,
		Bundle: bundle,
		Stdout: stdout,
		Stderr: stderr,
		Rootfs: rootfs,
	}
	opts := &runcopt.Options{SystemdCgroup: false}
	if runtime.GOOS != "windows" {
		uid := os.Getuid()
		gid := os.Getgid()
		if uid > 0 {
			runcRoot := filepath.Join(os.TempDir(), "shimtest-runc")
			os.MkdirAll(runcRoot, 0700)
			opts.IoUid = uint32(uid)
			opts.IoGid = uint32(gid)
			opts.Root = runcRoot
		}
	}
	any, err := typeurl.MarshalAnyToProto(opts)
	if err != nil {
		tb.Fatal("failed to marshal runc options:", err)
	}
	req.Options = &anypb.Any{TypeUrl: any.TypeUrl, Value: any.Value}
	return req
}

// bootstrapParams is the JSON / protobuf payload returned on stdout
// from `shim start`.
type bootstrapParams struct {
	Version  int    `json:"version"`
	Address  string `json:"address"`
	Protocol string `json:"protocol"`
}

// parseBootstrapResult tries to decode the shim's start response.
// Newer shims (containerd v2.3+) return protobuf; older ones return
// JSON.
func parseBootstrapResult(data []byte, params *bootstrapParams) error {
	if len(data) > 0 && data[0] == '{' {
		return json.Unmarshal(data, params)
	}
	b := data
	for len(b) > 0 {
		num, wtype, n := protowire.ConsumeTag(b)
		if n < 0 {
			return fmt.Errorf("invalid protobuf tag")
		}
		b = b[n:]
		switch wtype {
		case protowire.VarintType:
			v, n := protowire.ConsumeVarint(b)
			if n < 0 {
				return fmt.Errorf("invalid protobuf varint")
			}
			b = b[n:]
			if num == 1 {
				params.Version = int(v)
			}
		case protowire.BytesType:
			v, n := protowire.ConsumeBytes(b)
			if n < 0 {
				return fmt.Errorf("invalid protobuf bytes")
			}
			b = b[n:]
			switch num {
			case 2:
				params.Address = string(v)
			case 3:
				params.Protocol = string(v)
			}
		default:
			return fmt.Errorf("unexpected protobuf wire type %d for field %d", wtype, num)
		}
	}
	if params.Address == "" {
		return fmt.Errorf("no address in bootstrap result")
	}
	return nil
}

// startShim runs the shim binary's "start" subcommand and returns
// the bootstrap params. Registers cleanup that ensures the shim
// process exits before the test ends.
func startShim(tb testing.TB, shimBin, bundleDir, id, ns string, cfg Config) bootstrapParams {
	tb.Helper()

	socketDir, err := os.MkdirTemp(unixSafeDir(), "nb-")
	if err != nil {
		tb.Fatal("failed to create socket dir:", err)
	}
	tb.Cleanup(func() { os.RemoveAll(socketDir) })

	// setupLogPipe is platform-specific (io_unix.go / io_windows.go).
	// On Linux it creates bundleDir/log as a FIFO for the shim to write to.
	// On Windows it dials \\.\pipe\containerd-shim-<ns>-<id>-log, which the
	// shim creates as a server — mirroring containerd's shim_windows.go.
	logReader := setupLogPipe(tb, bundleDir, ns, id)

	// Buffer shim logs and only dump them on test failure.
	// Benchmarks always discard.
	_, isBench := tb.(*testing.B)
	var logBuf bytes.Buffer
	done := make(chan struct{})
	go func() {
		defer close(done)
		buf := make([]byte, 32768)
		for {
			n, err := logReader.Read(buf)
			if n > 0 && !isBench {
				logBuf.Write(buf[:n])
			}
			if err != nil {
				return
			}
		}
	}()
	if !isBench {
		tb.Cleanup(func() {
			if !tb.Failed() {
				return
			}
			if logBuf.Len() == 0 {
				return
			}
			tb.Logf("shim logs:\n%s", logBuf.String())
		})
	}
	tb.Cleanup(func() {
		logReader.Close()
		// Cap the wait: if the shim holds its log pipe open (e.g. nerdbox
		// doesn't exit cleanly after Shutdown), don't block indefinitely.
		select {
		case <-done:
		case <-time.After(shutdownTimeout):
		}
	})

	containerdAddr := containerdSockPath(tb, bundleDir)

	bootParams := &bootapi.BootstrapParams{
		InstanceID:            id,
		Namespace:             ns,
		ContainerdGrpcAddress: containerdAddr,
		SocketDir:             &socketDir,
	}
	if cfg.Debug {
		bootParams.LogLevel = bootapi.LogLevel_LOG_LEVEL_DEBUG
	}
	bootData, err := proto.Marshal(bootParams)
	if err != nil {
		tb.Fatal("marshal bootstrap params:", err)
	}

	shimArgs := []string{
		"-namespace", ns,
		"-id", id,
		"-address", containerdAddr,
	}
	if cfg.Debug {
		shimArgs = append(shimArgs, "-debug")
	}
	shimArgs = append(shimArgs, "start")
	cmd := exec.Command(shimBin, shimArgs...)
	cmd.Dir = bundleDir
	cmd.Stdin = bytes.NewReader(bootData)
	cmd.Env = append(os.Environ(),
		"GOMAXPROCS=2",
		"SHIM_SOCKET_DIR="+socketDir,
		"TTRPC_ADDRESS="+containerdAddr,
	)
	for k, v := range cfg.Env {
		cmd.Env = append(cmd.Env, k+"="+v)
	}

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	if err := cmd.Run(); err != nil {
		tb.Fatalf("shim start failed: %v\nstderr: %s", err, stderr.String())
	}

	var params bootstrapParams
	if err := parseBootstrapResult(stdout.Bytes(), &params); err != nil {
		tb.Fatalf("failed to parse bootstrap params: %v\nraw stdout: %s", err, stdout.String())
	}
	if params.Address == "" {
		tb.Fatal("shim returned empty address")
	}

	tb.Cleanup(func() {
		// Match containerd's cleanup behavior: invoke the shim binary's
		// `delete` subcommand for protocol-prescribed bundle/runc/mount
		// cleanup. The daemon itself exits via TTRPC Shutdown earlier
		// in shutdownShim; if it didn't, containerd has no protocol
		// method to kill it, so neither do we — leaks are surfaced by
		// the StressSuite leak detector instead of papered over with
		// SIGKILL on a potentially stale pid.
		deleteShim(tb, shimBin, bundleDir, id, ns, cfg)
	})

	return params
}

// deleteShim invokes the shim binary's `delete` subcommand to perform
// the protocol-prescribed bundle/runc/mount cleanup. This is what
// containerd does to clean up "dead shim" state (see
// core/runtime/v2/binary.go). Failures are logged but don't fail the
// test — best-effort cleanup.
func deleteShim(tb testing.TB, shimBin, bundleDir, id, ns string, cfg Config) {
	tb.Helper()
	containerdAddr := containerdSockPath(tb, bundleDir)
	args := []string{
		"-namespace", ns,
		"-id", id,
		"-bundle", bundleDir,
		"-address", containerdAddr,
	}
	if cfg.Debug {
		args = append(args, "-debug")
	}
	args = append(args, "delete")

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	cmd := exec.CommandContext(ctx, shimBin, args...)
	cmd.Dir = bundleDir

	var stderr bytes.Buffer
	cmd.Stderr = &stderr
	if err := cmd.Run(); err != nil {
		tb.Logf("shim delete failed: %v: %s", err, strings.TrimSpace(stderr.String()))
	}
}

// shimPidViaConnect dials the shim's TTRPC address and asks for its
// pid via the task service Connect RPC. Retries with a short backoff
// for up to retryFor since the server may take a few milliseconds to
// start accepting connections after the address is reported. Returns
// the last error if the deadline is reached without a successful
// response (some shims, e.g. nerdbox, return FailedPrecondition until
// a task exists). Used by RSS monitoring in StressSuite, which calls
// it after tc.Create when every conformant shim responds.
//
// dialShimConn is platform-specific (connect_unix.go / connect_windows.go).
func shimPidViaConnect(address, id string, retryFor time.Duration) (int, error) {
	deadline := time.Now().Add(retryFor)
	var lastErr error
	for {
		conn, err := dialShimConn(address, 500*time.Millisecond)
		if err != nil {
			lastErr = fmt.Errorf("dial: %w", err)
		} else {
			client := ttrpc.NewClient(conn)
			tc := taskAPI.NewTTRPCTaskClient(client)
			ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
			resp, callErr := tc.Connect(ctx, &taskAPI.ConnectRequest{ID: id})
			cancel()
			client.Close()
			if callErr != nil {
				lastErr = fmt.Errorf("connect RPC: %w", callErr)
			} else if resp.ShimPid == 0 {
				lastErr = fmt.Errorf("connect returned ShimPid=0")
			} else {
				return int(resp.ShimPid), nil
			}
		}
		if time.Now().After(deadline) {
			return 0, lastErr
		}
		time.Sleep(20 * time.Millisecond)
	}
}

// unixSafeDir returns a base directory for os.MkdirTemp that keeps the
// resulting unix socket paths within the 104-byte AF_UNIX pathname limit
// on macOS. The shim creates its socket at:
//
//	<socketDir>/<sha256hex>
//
// where socketDir is the result of os.MkdirTemp(base, "nb-") — roughly
// len(base) + 14 bytes ("/nb-" + up to 10 random digits). sha256hex is
// always 64 bytes, so the maximum safe length for base is:
//
//	104 (AF_UNIX limit) - 64 (sha256hex) - 14 (dir overhead) = 26 bytes
//
// If os.TempDir() exceeds that threshold, fall back to /tmp, which is
// always short enough. On Windows /tmp does not exist, so os.TempDir()
// is always used regardless of length.
func unixSafeDir() string {
	if runtime.GOOS == "windows" {
		return os.TempDir()
	}
	const (
		afUnixLimit   = 104                                       // macOS AF_UNIX pathname length limit
		socketNameLen = 64                                        // sha256 hex digest
		dirOverhead   = 14                                        // "/nb-" prefix + up to 10 random digits
		maxBaseLen    = afUnixLimit - socketNameLen - dirOverhead // 26
	)
	if len(os.TempDir()) > maxBaseLen {
		return "/tmp"
	}
	return os.TempDir()
}

// shutdownTask calls Shutdown on tc with a bounded shutdownTimeout derived
// from ctx. Shims that exit without sending a ttrpc response will surface as
// a context-deadline error rather than hanging the test indefinitely.
func shutdownTask(ctx context.Context, tc taskAPI.TTRPCTaskService, id string) {
	shutCtx, cancel := context.WithTimeout(ctx, shutdownTimeout)
	defer cancel()
	tc.Shutdown(shutCtx, &taskAPI.ShutdownRequest{ID: id})
}

// parseIntBytes parses a decimal integer from a byte slice.
func parseIntBytes(b []byte) (int, error) {
	s := strings.TrimSpace(string(b))
	var n int
	for _, c := range s {
		if c < '0' || c > '9' {
			return 0, nil
		}
		n = n*10 + int(c-'0')
	}
	return n, nil
}