From 588b0c219c1f56f5bb3dc37d790e39b06cbf02e2 Mon Sep 17 00:00:00 2001
From: Nick Ethier <nethier@hashicorp.com>
Date: Fri, 19 Jun 2020 10:51:32 -0400
Subject: [PATCH] taskenv: populate NOMAD_IP|PORT|ADDR env from allocated ports

---
 client/taskenv/env.go      | 33 +++++++++++++++++++++++++++++++++
 client/taskenv/env_test.go | 16 ++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/client/taskenv/env.go b/client/taskenv/env.go
index f41609021..f92d5c0f4 100644
--- a/client/taskenv/env.go
+++ b/client/taskenv/env.go
@@ -70,15 +70,21 @@ const (
 	// The ip:port are always the host's.
 	AddrPrefix = "NOMAD_ADDR_"
 
+	HostAddrPrefix = "NOMAD_HOST_ADDR_"
+
 	// IpPrefix is the prefix for passing the host IP of a port allocation
 	// to a task.
 	IpPrefix = "NOMAD_IP_"
 
+	HostIpPrefix = "NOMAD_HOST_IP_"
+
 	// PortPrefix is the prefix for passing the port allocation to a task.
 	// It will be the task's port if a port map is specified. Task's should
 	// bind to this port.
 	PortPrefix = "NOMAD_PORT_"
 
+	AllocPortPrefix = "NOMAD_ALLOC_PORT_"
+
 	// HostPortPrefix is the prefix for passing the host port when a port
 	// map is specified.
 	HostPortPrefix = "NOMAD_HOST_PORT_"
@@ -620,6 +626,7 @@ func (b *Builder) setAlloc(alloc *structs.Allocation) *Builder {
 			}
 		}
 
+		// COMPAT(1.0): remove in 1.0 when AllocatedPorts can be used exclusivly
 		// Add ports from other tasks
 		for taskName, resources := range alloc.AllocatedResources.Tasks {
 			// Add ports from other tasks
@@ -637,6 +644,7 @@ func (b *Builder) setAlloc(alloc *structs.Allocation) *Builder {
 			}
 		}
 
+		// COMPAT(1.0): remove in 1.0 when AllocatedPorts can be used exclusivly
 		// Add ports from group networks
 		//TODO Expose IPs but possibly only via variable interpolation
 		for _, nw := range alloc.AllocatedResources.Shared.Networks {
@@ -647,6 +655,11 @@ func (b *Builder) setAlloc(alloc *structs.Allocation) *Builder {
 				addGroupPort(b.otherPorts, p)
 			}
 		}
+
+		// Add any allocated host ports
+		if alloc.AllocatedResources.Shared.Ports != nil {
+			addPorts(b.otherPorts, alloc.AllocatedResources.Shared.Ports)
+		}
 	}
 
 	upstreams := []structs.ConsulUpstream{}
@@ -857,3 +870,23 @@ func addGroupPort(m map[string]string, port structs.Port) {
 
 	m[HostPortPrefix+port.Label] = strconv.Itoa(port.Value)
 }
+
+func addPorts(m map[string]string, ports structs.AllocatedPorts) {
+	for _, p := range ports {
+		m[AddrPrefix+p.Label] = fmt.Sprintf("%s:%d", p.HostIP, p.Value)
+		m[HostAddrPrefix+p.Label] = fmt.Sprintf("%s:%d", p.HostIP, p.Value)
+		m[IpPrefix+p.Label] = p.HostIP
+		m[HostIpPrefix+p.Label] = p.HostIP
+		if p.To > 0 {
+			val := strconv.Itoa(p.To)
+			m[PortPrefix+p.Label] = val
+			m[AllocPortPrefix+p.Label] = val
+		} else {
+			val := strconv.Itoa(p.Value)
+			m[PortPrefix+p.Label] = val
+			m[AllocPortPrefix+p.Label] = val
+		}
+
+		m[HostPortPrefix+p.Label] = strconv.Itoa(p.Value)
+	}
+}
diff --git a/client/taskenv/env_test.go b/client/taskenv/env_test.go
index 3e5316907..08c96eb54 100644
--- a/client/taskenv/env_test.go
+++ b/client/taskenv/env_test.go
@@ -360,6 +360,15 @@ func TestEnvironment_AllValues(t *testing.T) {
 		},
 	}
 
+	a.AllocatedResources.Shared.Ports = structs.AllocatedPorts{
+		{
+			Label:  "admin",
+			Value:  32000,
+			To:     9000,
+			HostIP: "127.0.0.1",
+		},
+	}
+
 	sharedNet := a.AllocatedResources.Shared.Networks[0]
 
 	// Add group network port with only a host port.
@@ -463,6 +472,13 @@ func TestEnvironment_AllValues(t *testing.T) {
 		"NOMAD_HOST_PORT_hostonly":                  "9998",
 		"NOMAD_PORT_static":                         "97",
 		"NOMAD_HOST_PORT_static":                    "9997",
+		"NOMAD_ADDR_admin":                          "127.0.0.1:32000",
+		"NOMAD_HOST_ADDR_admin":                     "127.0.0.1:32000",
+		"NOMAD_IP_admin":                            "127.0.0.1",
+		"NOMAD_HOST_IP_admin":                       "127.0.0.1",
+		"NOMAD_PORT_admin":                          "9000",
+		"NOMAD_ALLOC_PORT_admin":                    "9000",
+		"NOMAD_HOST_PORT_admin":                     "32000",
 
 		// 0.9 style env map
 		`env["taskEnvKey"]`:        "taskEnvVal",