cep/docker/container.go

package docker

import (
	"archive/tar"
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net/netip"
	"strings"
	"time"

	fw "git.buran.team/main/fairwind"

	"github.com/docker/docker/pkg/stdcopy"
	mobycontainer "github.com/moby/moby/api/types/container"
	mobynetwork "github.com/moby/moby/api/types/network"
	moby "github.com/moby/moby/client"
)

var ErrContainerExecReturnedNonZero = errors.New("container exec returned non-zero")
var ErrContainerNotExited = errors.New("container not exited")

type ExecutionResult struct {
	Code   int
	Stdout []byte
	Stderr []byte
	Error  error
}

type ContainerHook func(ctx context.Context) error

type ContainerSettingsPermissions struct {
	Privileged   bool
	Capabilities []string
}

type ContainerSettingsNetwork struct {
	IP string
}

type ContainerResources struct {
	CPU    int64
	Memory int64
	// TODO: disk iops, network iops
}

type Variable struct {
	Key   string
	Value string
}

type ContainerSettingsHost struct {
	Name string
	IP   string
}

type ContainerSettings struct {
	Name        string
	Command     string
	Variables   []Variable
	Hosts       []ContainerSettingsHost
	Binds       []string
	Permissions ContainerSettingsPermissions
	Network     ContainerSettingsNetwork
	Resources   ContainerResources
}

type Container struct {
	ctx      context.Context
	log      *fw.Log
	client   *Docker
	network  *Network
	image    *Image
	settings ContainerSettings
}

func NewContainer(ctx context.Context, log *fw.Log, client *Docker, network *Network, image *Image, settings ContainerSettings) *Container {
	return &Container{
		ctx:      ctx,
		log:      log,
		client:   client,
		network:  network,
		image:    image,
		settings: settings,
	}
}

func (this *Container) Image() *Image {
	return this.image
}

func (this *Container) Start() error {
	err := this.image.Pull()
	if err != nil {
		return fmt.Errorf("can't start container: %w", err)
	}

	binds := []string{}
	for _, bind := range this.settings.Binds {
		binds = append(binds, fmt.Sprintf("%s:%s", bind, bind))
	}

	command := []string{}
	if len(this.settings.Command) > 0 {
		command = strings.Split(this.settings.Command, " ")
	} else {
		command = nil
	}

	env := []string{}
	if this.settings.Variables != nil {
		for _, environment := range this.settings.Variables {
			env = append(
				env,
				fmt.Sprintf("%s=%s", environment.Key, environment.Value),
			)
		}
	}

	hosts := []string{}
	if this.settings.Hosts != nil {
		for _, host := range this.settings.Hosts {
			hosts = append(
				hosts,
				fmt.Sprintf("%s:%s", host.Name, host.IP),
			)
		}
	}

	resources := mobycontainer.Resources{}
	if this.settings.Resources.CPU > 0 {
		resources.NanoCPUs = this.settings.Resources.CPU
	}

	if this.settings.Resources.Memory > 0 {
		resources.Memory = this.settings.Resources.Memory
	}

	// Create container
	_, err = this.client.Docker.ContainerCreate(
		this.ctx,
		moby.ContainerCreateOptions{
			Name: this.settings.Name,
			Config: &mobycontainer.Config{
				Image: this.image.Name(),
				Cmd:   command,
				Env:   env,
			},
			HostConfig: &mobycontainer.HostConfig{
				Binds: binds,
				RestartPolicy: mobycontainer.RestartPolicy{
					Name: mobycontainer.RestartPolicyDisabled,
				},
				Privileged:  this.settings.Permissions.Privileged,
				CapAdd:      this.settings.Permissions.Capabilities,
				NetworkMode: "default",
				Resources:   resources,
				ExtraHosts:  hosts,
			},
			NetworkingConfig: &mobynetwork.NetworkingConfig{
				EndpointsConfig: map[string]*mobynetwork.EndpointSettings{
					this.network.Name(): {
						IPAMConfig: &mobynetwork.EndpointIPAMConfig{
							IPv4Address: netip.MustParseAddr(this.settings.Network.IP),
						},
					},
				},
			},
		},
	)
	if err != nil {
		return fmt.Errorf("can't start container: %w", err)
	}

	// Start
	_, err = this.client.Docker.ContainerStart(
		this.ctx,
		this.settings.Name,
		moby.ContainerStartOptions{},
	)
	if err != nil {
		return fmt.Errorf("can't start container: %w", err)
	}

	return nil
}

func (this *Container) Stop() error {
	_, err := this.client.Docker.ContainerRemove(
		this.ctx,
		this.settings.Name,
		moby.ContainerRemoveOptions{
			Force: true,
		},
	)
	if err != nil {
		return err
	}

	return nil
}

func (this *Container) Status(shouldExit bool) ExecutionResult {
	// Get status
	result0, err := this.client.Docker.ContainerInspect(
		this.ctx,
		this.settings.Name,
		moby.ContainerInspectOptions{},
	)
	if err != nil {
		return ExecutionResult{
			Error: err,
		}
	}

	if shouldExit {
		if result0.Container.State.Status != "exited" {
			return ExecutionResult{
				Error: ErrContainerNotExited,
			}
		}
	}

	// Get logs
	result1, err := this.client.Docker.ContainerLogs(
		this.ctx,
		this.settings.Name,
		moby.ContainerLogsOptions{
			ShowStdout: true,
			ShowStderr: true,
		},
	)
	if err != nil {
		return ExecutionResult{
			Error: err,
		}
	}

	var stdout bytes.Buffer
	var stderr bytes.Buffer
	_, err = stdcopy.StdCopy(&stdout, &stderr, result1)
	if err != nil {
		return ExecutionResult{
			Error: err,
		}
	}

	// Done
	return ExecutionResult{
		Code:   result0.Container.State.ExitCode,
		Stdout: stdout.Bytes(),
		Stderr: stderr.Bytes(),
	}
}

func (this *Container) Execute(timeout int, command string) ExecutionResult {
	// Context
	ctx, cancel := context.WithTimeout(
		this.ctx,
		time.Duration(timeout)*time.Millisecond,
	)
	defer cancel()

	// Create
	resultCreate, err := this.client.Docker.ExecCreate(
		ctx,
		this.settings.Name,
		moby.ExecCreateOptions{
			AttachStdin:  false,
			AttachStdout: true,
			AttachStderr: true,
			Cmd:          strings.Split(command, " "),
		},
	)
	if err != nil {
		return ExecutionResult{
			Error: fmt.Errorf("can't exec command in container: %w", err),
		}
	}

	resultAttach, err := this.client.Docker.ExecAttach(
		ctx,
		resultCreate.ID,
		moby.ExecAttachOptions{},
	)
	if err != nil {
		return ExecutionResult{
			Error: fmt.Errorf("can't exec command in container: %w", err),
		}
	}
	defer resultAttach.Close()

	// Start
	_, err = this.client.Docker.ExecStart(
		ctx,
		resultCreate.ID,
		moby.ExecStartOptions{},
	)
	if err != nil {
		return ExecutionResult{
			Error: fmt.Errorf("can't exec command in container: %w", err),
		}
	}

	// Wait
	exitCode := 0
	for {
		resultInspect, err := this.client.Docker.ExecInspect(
			ctx,
			resultCreate.ID,
			moby.ExecInspectOptions{},
		)
		if err != nil {
			return ExecutionResult{
				Error: fmt.Errorf("can't exec command in container: %w", err),
			}
		}

		if resultInspect.Running {
			time.Sleep(10 * time.Millisecond)
			continue
		}

		exitCode = resultInspect.ExitCode
		break
	}

	var stdout bytes.Buffer
	var stderr bytes.Buffer
	_, err = stdcopy.StdCopy(&stdout, &stderr, resultAttach.Reader)
	if err != nil {
		return ExecutionResult{
			Error: fmt.Errorf("can't exec command in container: %w", err),
		}
	}

	// Done
	return ExecutionResult{
		Code:   exitCode,
		Stdout: stdout.Bytes(),
		Stderr: stderr.Bytes(),
	}
}

func (this *Container) Read(timeout int, path string) ([]byte, error) {
	// Context
	ctx, cancel := context.WithTimeout(
		this.ctx,
		time.Duration(timeout)*time.Millisecond,
	)
	defer cancel()

	// Read
	result, err := this.client.Docker.CopyFromContainer(
		ctx,
		this.settings.Name,
		moby.CopyFromContainerOptions{
			SourcePath: path,
		},
	)
	if err != nil {
		return nil, fmt.Errorf("can't copy from container: %w", err)
	}
	defer result.Content.Close()

	// Tar
	buffer := []byte{}
	reader := tar.NewReader(result.Content)

	_, err = reader.Next()
	if err == io.EOF {
		return nil, fmt.Errorf("can't copy from container: %w", err)
	}

	if err != nil {
		return nil, fmt.Errorf("can't copy from container: %w", err)
	}

	buffer, err = io.ReadAll(reader)
	if err != nil {
		return nil, fmt.Errorf("can't copy from container: %w", err)
	}

	// Done
	return buffer, nil
}