fix qemu and update docker with tests

hashicorp · Sep 4, 2019 · bd6bbc9 · bd6bbc9
1 parent 6190443
commit bd6bbc9
Show file tree

Hide file tree

Showing 5 changed files with 73 additions and 3 deletions.
diff --git a/client/taskenv/env.go b/client/taskenv/env.go
@@ -786,9 +786,14 @@ func buildUpstreamsEnv(envMap map[string]string, upstreams []structs.ConsulUpstr
 	}
 }
 
-func WithPortMapEnvs(envs map[string]string, ports map[string]int) map[string]string {
+// SetPortMapEnvs sets the PortMap related environment variables on the map
+func SetPortMapEnvs(envs map[string]string, ports map[string]int) map[string]string {
+	if envs == nil {
+		envs = map[string]string{}
+	}
+
 	for portLabel, port := range ports {
-		portEnv := PortPrefix + portLabel
+		portEnv := helper.CleanEnvVar(PortPrefix+portLabel, '_')
 		envs[portEnv] = strconv.Itoa(port)
 	}
 	return envs

diff --git a/client/taskenv/env_test.go b/client/taskenv/env_test.go
@@ -784,3 +784,23 @@ func TestEnvironment_Upstreams(t *testing.T) {
 	require.Equal(t, "127.0.0.1:1234", env["foo"])
 	require.Equal(t, "1234", env["bar"])
 }
+
+func TestEnvironment_SetPortMapEnvs(t *testing.T) {
+	envs := map[string]string{
+		"foo":            "bar",
+		"NOMAD_PORT_ssh": "2342",
+	}
+	ports := map[string]int{
+		"ssh":  22,
+		"http": 80,
+	}
+
+	envs = SetPortMapEnvs(envs, ports)
+
+	expected := map[string]string{
+		"foo":             "bar",
+		"NOMAD_PORT_ssh":  "22",
+		"NOMAD_PORT_http": "80",
+	}
+	require.Equal(t, expected, envs)
+}
diff --git a/drivers/docker/driver.go b/drivers/docker/driver.go
@@ -655,6 +655,9 @@ func (d *Driver) containerBinds(task *drivers.TaskConfig, driverConfig *TaskConf
 func (d *Driver) createContainerConfig(task *drivers.TaskConfig, driverConfig *TaskConfig,
 	imageID string) (docker.CreateContainerOptions, error) {
 
+	// ensure that PortMap variables are populated early on
+	task.Env = taskenv.SetPortMapEnvs(task.Env, driverConfig.PortMap)
+
 	logger := d.logger.With("task_name", task.Name)
 	var c docker.CreateContainerOptions
 	if task.Resources == nil {
@@ -967,7 +970,6 @@ func (d *Driver) createContainerConfig(task *drivers.TaskConfig, driverConfig *T
 		logger.Debug("applied labels on the container", "labels", config.Labels)
 	}
 
-	task.Env = taskenv.WithPortMapEnvs(task.Env, driverConfig.PortMap)
 	config.Env = task.EnvList()
 
 	containerName := fmt.Sprintf("%s-%s", strings.Replace(task.Name, "/", "_", -1), task.AllocID)

diff --git a/drivers/docker/driver_test.go b/drivers/docker/driver_test.go
@@ -1412,6 +1412,10 @@ func TestDockerDriver_PortsMapping(t *testing.T) {
 	container, err := client.InspectContainer(handle.containerID)
 	require.NoError(t, err)
 
+	// Verify that the port environment variables are set
+	require.Contains(t, container.Config.Env, "NOMAD_PORT_main=8080")
+	require.Contains(t, container.Config.Env, "NOMAD_PORT_REDIS=6379")
+
 	// Verify that the correct ports are EXPOSED
 	expectedExposedPorts := map[docker.Port]struct{}{
 		docker.Port("8080/tcp"): {},
@@ -1437,6 +1441,41 @@ func TestDockerDriver_PortsMapping(t *testing.T) {
 	require.Exactly(t, expectedPortBindings, container.HostConfig.PortBindings)
 }
 
+func TestDockerDriver_CreateContainerConfig_PortsMapping(t *testing.T) {
+	t.Parallel()
+
+	task, cfg, port := dockerTask(t)
+	res := port[0]
+	dyn := port[1]
+	cfg.PortMap = map[string]int{
+		"main":  8080,
+		"REDIS": 6379,
+	}
+	dh := dockerDriverHarness(t, nil)
+	driver := dh.Impl().(*Driver)
+
+	c, err := driver.createContainerConfig(task, cfg, "org/repo:0.1")
+	require.NoError(t, err)
+
+	require.Equal(t, "org/repo:0.1", c.Config.Image)
+	require.Contains(t, c.Config.Env, "NOMAD_PORT_main=8080")
+	require.Contains(t, c.Config.Env, "NOMAD_PORT_REDIS=6379")
+
+	// Verify that the correct ports are FORWARDED
+	hostIP := "127.0.0.1"
+	if runtime.GOOS == "windows" {
+		hostIP = ""
+	}
+	expectedPortBindings := map[docker.Port][]docker.PortBinding{
+		docker.Port("8080/tcp"): {{HostIP: hostIP, HostPort: fmt.Sprintf("%d", res)}},
+		docker.Port("8080/udp"): {{HostIP: hostIP, HostPort: fmt.Sprintf("%d", res)}},
+		docker.Port("6379/tcp"): {{HostIP: hostIP, HostPort: fmt.Sprintf("%d", dyn)}},
+		docker.Port("6379/udp"): {{HostIP: hostIP, HostPort: fmt.Sprintf("%d", dyn)}},
+	}
+	require.Exactly(t, expectedPortBindings, c.HostConfig.PortBindings)
+
+}
+
 func TestDockerDriver_CleanupContainer(t *testing.T) {
 	if !tu.IsCI() {
 		t.Parallel()

diff --git a/drivers/qemu/driver.go b/drivers/qemu/driver.go
@@ -14,6 +14,7 @@ import (
 
 	"github.com/coreos/go-semver/semver"
 	hclog "github.com/hashicorp/go-hclog"
+	"github.com/hashicorp/nomad/client/taskenv"
 	"github.com/hashicorp/nomad/drivers/shared/eventer"
 	"github.com/hashicorp/nomad/drivers/shared/executor"
 	"github.com/hashicorp/nomad/helper/pluginutils/hclutils"
@@ -304,6 +305,9 @@ func (d *Driver) StartTask(cfg *drivers.TaskConfig) (*drivers.TaskHandle, *drive
 		return nil, nil, fmt.Errorf("failed to decode driver config: %v", err)
 	}
 
+	// ensure that PortMap variables are populated early on
+	cfg.Env = taskenv.SetPortMapEnvs(cfg.Env, driverConfig.PortMap)
+
 	handle := drivers.NewTaskHandle(taskHandleVersion)
 	handle.Config = cfg