From dff556a761299a83e16b5c98665d07bf052dbccf Mon Sep 17 00:00:00 2001
From: Mahmood Ali <mahmood@hashicorp.com>
Date: Wed, 4 Sep 2019 09:33:35 -0400
Subject: [PATCH] fix qemu and update docker with tests

---
 client/taskenv/env.go         |  9 +++++++--
 drivers/docker/driver.go      |  4 +++-
 drivers/docker/driver_test.go | 32 ++++++++++++++++++++++++++++++++
 drivers/qemu/driver.go        |  4 ++++
 4 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/client/taskenv/env.go b/client/taskenv/env.go
index a90991852c78..96236d9ddb70 100644
--- a/client/taskenv/env.go
+++ b/client/taskenv/env.go
@@ -730,9 +730,14 @@ func buildPortEnv(envMap map[string]string, p structs.Port, ip string, driverNet
 	}
 }
 
-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/drivers/docker/driver.go b/drivers/docker/driver.go
index 21f84286910f..ac83f376d865 100644
--- 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
index 50e6d9778a52..d30252ba66fb 100644
--- a/drivers/docker/driver_test.go
+++ b/drivers/docker/driver_test.go
@@ -1437,6 +1437,38 @@ 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"
+	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
index a4fef4b39dfe..8cbe7d9f8f97 100644
--- 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