docker_test.go

package shim

import (
	"context"
	"encoding/hex"
	"math/rand"
	"os"
	"runtime"
	"sync"
	"testing"
	"time"

	"github.com/docker/docker/api/types/mount"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// TestDocker_SSHServer pulls ubuntu image (without sshd), installs openssh-server and exits
// Basically, it indirectly tests a shell script generated by getSSHShellCommands
func TestDocker_SSHServer(t *testing.T) {
	if testing.Short() || (os.Getenv("CI") == "true" && runtime.GOOS == "darwin") {
		t.Skip()
	}
	t.Parallel()

	params := &dockerParametersMock{
		commands:         []string{"/usr/sbin/sshd -V 2>&1 | grep OpenSSH"},
		sshShellCommands: true,
		runnerDir:        t.TempDir(),
	}

	timeout := 180 // seconds
	ctx, cancel := context.WithTimeout(t.Context(), time.Duration(timeout)*time.Second)
	defer cancel()

	dockerRunner, err := NewDockerRunner(ctx, params)
	require.NoError(t, err)

	taskConfig := createTaskConfig(t)
	defer dockerRunner.Remove(t.Context(), taskConfig.ID)

	assert.NoError(t, dockerRunner.Submit(ctx, taskConfig))
	assert.NoError(t, dockerRunner.Run(ctx, taskConfig.ID))
}

func TestDocker_ShmNoexecByDefault(t *testing.T) {
	if testing.Short() || (os.Getenv("CI") == "true" && runtime.GOOS == "darwin") {
		t.Skip()
	}
	t.Parallel()

	params := &dockerParametersMock{
		commands:  []string{"mount | grep '/dev/shm .*size=65536k' | grep noexec"},
		runnerDir: t.TempDir(),
	}

	timeout := 180 // seconds
	ctx, cancel := context.WithTimeout(t.Context(), time.Duration(timeout)*time.Second)
	defer cancel()

	dockerRunner, err := NewDockerRunner(ctx, params)
	require.NoError(t, err)

	taskConfig := createTaskConfig(t)
	defer dockerRunner.Remove(t.Context(), taskConfig.ID)

	assert.NoError(t, dockerRunner.Submit(ctx, taskConfig))
	assert.NoError(t, dockerRunner.Run(ctx, taskConfig.ID))
}

func TestDocker_ShmExecIfSizeSpecified(t *testing.T) {
	if testing.Short() || (os.Getenv("CI") == "true" && runtime.GOOS == "darwin") {
		t.Skip()
	}
	t.Parallel()

	params := &dockerParametersMock{
		commands:  []string{"mount | grep '/dev/shm .*size=1024k' | grep -v noexec"},
		runnerDir: t.TempDir(),
	}

	timeout := 180 // seconds
	ctx, cancel := context.WithTimeout(t.Context(), time.Duration(timeout)*time.Second)
	defer cancel()

	dockerRunner, err := NewDockerRunner(ctx, params)
	require.NoError(t, err)

	taskConfig := createTaskConfig(t)
	taskConfig.ShmSize = 1024 * 1024
	defer dockerRunner.Remove(t.Context(), taskConfig.ID)

	assert.NoError(t, dockerRunner.Submit(ctx, taskConfig))
	assert.NoError(t, dockerRunner.Run(ctx, taskConfig.ID))
}

/* Mocks */

type dockerParametersMock struct {
	commands         []string
	sshShellCommands bool
	runnerDir        string
}

func (c *dockerParametersMock) DockerPrivileged() bool {
	return false
}

func (c *dockerParametersMock) DockerPJRTDevice() string {
	return ""
}

func (c *dockerParametersMock) DockerShellCommands(authorizedKeys []string, runnerHttpAddress string) []string {
	commands := make([]string, 0)
	if c.sshShellCommands {
		commands = append(commands, getSSHShellCommands()...)
	}
	commands = append(commands, c.commands...)
	return commands
}

func (c *dockerParametersMock) DockerPorts() []int {
	return []int{}
}

func (c *dockerParametersMock) RunnerHTTPPort() int {
	return 10999
}

func (c *dockerParametersMock) DockerMounts(string) ([]mount.Mount, error) {
	return nil, nil
}

func (c *dockerParametersMock) MakeRunnerDir(string) (string, error) {
	return c.runnerDir, nil
}

/* Utilities */

var (
	randSrc = rand.New(rand.NewSource(time.Now().UnixNano()))
	randMu  = sync.Mutex{}
)

func generateID(t *testing.T) string {
	const idLen = 16
	b := make([]byte, idLen/2)
	randMu.Lock()
	defer randMu.Unlock()
	_, err := randSrc.Read(b)
	require.Nil(t, err)
	return hex.EncodeToString(b)[:idLen]
}

func createTaskConfig(t *testing.T) TaskConfig {
	return TaskConfig{
		ID:        generateID(t),
		Name:      t.Name(),
		ImageName: "ubuntu",
	}
}