src/pkg/containers/containers.go

package containers

import (
	"context"
	"fmt"
	"log"
	"os"
	"syscall"

	containerd "github.com/containerd/containerd/v2/client"
	"github.com/containerd/containerd/v2/pkg/cio"
	"github.com/containerd/containerd/v2/pkg/namespaces"
	"github.com/containerd/containerd/v2/pkg/oci"
	"github.com/opencontainers/runtime-spec/specs-go"
)

// GetOrDownloadImage gets an image from the containerd client if the image is
// already available locally, and downloads and unpacks it if it is not.
func GetOrDownloadImage(ctx context.Context, client *containerd.Client, imageName string) (containerd.Image, error) {
	image, err := client.GetImage(ctx, imageName)
	if err != nil {
		log.Printf("error getting local image: %v\n", err)
	}
	if err != nil || image == nil {
		image, err = client.Pull(ctx, imageName, containerd.WithPullUnpack)
		if err != nil {
			return nil, fmt.Errorf("error pulling image: %v", err)
		} else {
			log.Println("finished pulling image")
		}
	}

	return image, nil
}

// A ContainerConfig is a simplified set of configurations which can be used to
// build an OCI container.
type ContainerConfig struct {
	// Name is the name that will be given to the container.
	Name       string
	// Command is the command that will serve as the container entrypoint. If empty,
	// the default command for the container will be used.
	Command    string
	// Args is the set of arguments which will be passed to the container command.
	Args       []string
	// Cwd is the working directory where the command will be run.
	Cwd        string
	// Env is a map from variable names to values which will be passed to the container
	// as environment variables.
	Env        map[string]string
	// DirMounts is the set of directories which will be mapped from the host to the container.
	// The directories will be created automatically if they don't exist.
	DirMounts     []specs.Mount
	// FileMounts is the set of files that will be mapped from the host to the container.
	// It is an error if the expected file doesn't exist in the host.
	FileMounts    []specs.Mount
	// Priviliged determines whether or not the container will be run in privileged mode.
	Privileged bool
	// NetHost determines whether or not the container will use the host's networking.
	NetHost    bool
	// Opts is any additional options which will be used to create the container.
	Opts       []oci.SpecOpts
}

// Run downloads an image if it's not available locally, then runs the container
// command specified by config and waits for its completion.
// Returns the exit status of the container, or 1 on other failure.
func Run(ctx context.Context, imageName string, config ContainerConfig, containerNamespace, containerdSocket string) (uint32, error) {
	client, err := containerd.New(containerdSocket)
	if err != nil {
		log.Fatalf("error connecting to containerd socket: %v\n", err)
	}
	defer client.Close()

	ctx = namespaces.WithNamespace(ctx, containerNamespace)

	image, err := GetOrDownloadImage(ctx, client, imageName)
	if err != nil {
		return 1, err
	}

	container, err := NewContainer(ctx, client, image, config)
	if err != nil {
		return 1, err
	}
	defer func() {
		if err := DeleteContainer(ctx, client, container); err != nil {
			log.Printf("could not delete container: %v\n", err)
		}
	}()

	return RunContainerTask(ctx, client, container)
}

// NewContainer builds and returns a new container from the provided image and config.
func NewContainer(
	ctx context.Context,
	client *containerd.Client,
	image containerd.Image,
	config ContainerConfig,
) (containerd.Container, error) {
	processArgs := append([]string{config.Command}, config.Args...)

	var envVarsArray []string
	for name, value := range config.Env {
		envVarsArray = append(envVarsArray, fmt.Sprintf("%s=%s", name, value))
	}

	for _, mount := range config.DirMounts {
		if err := os.MkdirAll(mount.Source, 0777); err != nil {
			return nil, fmt.Errorf("could not find or create mount directory %s: %v", mount.Source, err)
		}
	}

	// The order matters here! The image config has to come first.
	opts := []oci.SpecOpts{
		oci.WithImageConfig(image),
		oci.WithEnv(envVarsArray),
		oci.WithProcessArgs(processArgs...),
		oci.WithProcessCwd("/"),
		oci.WithMounts(config.DirMounts),
		oci.WithMounts(config.FileMounts),
	}

	if config.Privileged {
		opts = append(opts, oci.WithPrivileged)
	}

	if config.NetHost {
		hostname, err := os.Hostname()
		if err != nil {
			return nil, fmt.Errorf("could not get hostname: %v\n", err)
		}

		hostNetworkOpts := []oci.SpecOpts{
			oci.WithHostNamespace(specs.NetworkNamespace),
			oci.WithHostHostsFile,
			oci.WithHostResolvconf,
			oci.WithEnv([]string{fmt.Sprintf("HOSTNAME=%s", hostname)}),
		}
		opts = append(opts, hostNetworkOpts...)
	}

	opts = append(opts, config.Opts...)

	spec := containerd.WithNewSpec(opts...)

	container, err := client.NewContainer(
		ctx,
		config.Name,
		containerd.WithNewSnapshot(fmt.Sprintf("%s-rootfs", config.Name), image),
		spec,
	)
	if err != nil {
		return nil, fmt.Errorf("error creating container: %v", err)
	}

	return container, nil
}

// RunContainerTask runs a container's command and waits for its completion.
// Returns the exit status of the container, or 1 on other failure.
func RunContainerTask(ctx context.Context, client *containerd.Client, container containerd.Container) (uint32, error) {
	cioCreator := cio.NewCreator(cio.WithStdio)
	task, err := container.NewTask(ctx, cioCreator)
	if err != nil {
		return 1, fmt.Errorf("error creating task: %v", err)
	}
	defer func() {
		if err := DeleteTask(ctx, task); err != nil {
			log.Printf("could not delete container task: %v\n", err)
		}
	}()

	exitStatusC, err := task.Wait(ctx)
	if err != nil {
		return 1, fmt.Errorf("error waiting for task completion: %v", err)
	}

	err = task.Start(ctx)
	if err != nil {
		return 1, fmt.Errorf("error starting task: %v", err)
	}

	status := <-exitStatusC
	code, _, err := status.Result()
	if err != nil {
		return code, err
	}

	return code, nil
}

// DeleteContainer deletes a container and its associated snapshot.
func DeleteContainer(ctx context.Context, client *containerd.Client, container containerd.Container) error {
	if err := container.Delete(ctx, containerd.WithSnapshotCleanup); err != nil {
		return fmt.Errorf("failed to delete container: %v", err)
	}

	return nil
}

// DeleteTask deletes a task.
// If the task is not already stopped, this will kill it by calling KillTask.
func DeleteTask(ctx context.Context, task containerd.Task) error {
	taskStatus, err := task.Status(ctx)
	if err != nil {
		log.Printf("failed to get task status: %v\n", err)
	}
	if taskStatus.Status != containerd.Stopped {
		if err := KillTask(ctx, task); err != nil {
			return fmt.Errorf("failed to kill task: %v", err)
		}
	}

	if _, err = task.Delete(ctx); err != nil {
		log.Printf("error deleting task: %v\n", err)
	}

	return nil
}

// KillTask attempts to kill a task by sending SIGTERM.
// If the SIGTERM does not kill the task, a SIGKILL will be sent as well.
func KillTask(ctx context.Context, task containerd.Task) error {
	if err := task.Kill(ctx, syscall.SIGTERM); err != nil {
		return fmt.Errorf("could not kill task: %v", err)
	}

	taskStatus, err := task.Status(ctx)
	if err != nil || taskStatus.Status != containerd.Stopped {
		log.Printf("task could not be killed with SIGTERM, sending SIGKILL: %v\n", err)
		if err := task.Kill(ctx, syscall.SIGKILL); err != nil {
			return fmt.Errorf("could not kill task: %v", err)
		}
	}

	return nil
}