Allow custom ports for services and checks

Fixes #3380

Adds address_mode to checks (but no auto) and allows services and checks
to set literal port numbers when using address_mode=driver.

This allows SDNs, overlays, etc to advertise internal and host addresses
as well as do checks against either.

command/agent/consul/client.go

@@ -569,23 +569,16 @@ func (c *ServiceClient) serviceRegs(ops *operations, allocID string, service *st
 		checkIDs:  make(map[string]struct{}, len(service.Checks)),
 	}
 
-	// Determine the address to advertise
+	// Service address modes default to auto
 	addrMode := service.AddressMode
-	if addrMode == structs.AddressModeAuto {
-		if net.Advertise() {
-			addrMode = structs.AddressModeDriver
-		} else {
-			// No driver network or shouldn't default to driver's network
-			addrMode = structs.AddressModeHost
-		}
+	if addrMode == "" {
+		addrMode = structs.AddressModeAuto
 	}
-	ip, port := task.Resources.Networks.Port(service.PortLabel)
-	if addrMode == structs.AddressModeDriver {
-		if net == nil {
-			return nil, fmt.Errorf("service %s cannot use driver's IP because driver didn't set one", service.Name)
-		}
-		ip = net.IP
-		port = net.PortMap[service.PortLabel]
+
+	// Determine the address to advertise based on the mode
+	ip, port, err := getAddress(addrMode, service.PortLabel, task.Resources.Networks, net)
+	if err != nil {
+		return nil, fmt.Errorf("unable to get address for service %q: %v", service.Name, err)
 	}
 
 	// Build the Consul Service registration request
@@ -641,13 +634,24 @@ func (c *ServiceClient) checkRegs(ops *operations, allocID, serviceID string, se
 
 		}
 
-		// Checks should always use the host ip:port
+		// Default to the service's port but allow check to override
 		portLabel := check.PortLabel
 		if portLabel == "" {
 			// Default to the service's port label
 			portLabel = service.PortLabel
 		}
-		ip, port := task.Resources.Networks.Port(portLabel)
+
+		// Checks address mode defaults to host for pre-#3380 backward compat
+		addrMode := check.AddressMode
+		if addrMode == "" {
+			addrMode = structs.AddressModeHost
+		}
+
+		ip, port, err := getAddress(addrMode, portLabel, task.Resources.Networks, net)
+		if err != nil {
+			return nil, fmt.Errorf("unable to get address for check %q: %v", check.Name, err)
+		}
+
 		checkReg, err := createCheckReg(serviceID, checkID, check, ip, port)
 		if err != nil {
 			return nil, fmt.Errorf("failed to add check %q: %v", check.Name, err)
@@ -709,8 +713,8 @@ func (c *ServiceClient) RegisterTask(allocID string, task *structs.Task, restart
 func (c *ServiceClient) UpdateTask(allocID string, existing, newTask *structs.Task, restarter TaskRestarter, exec driver.ScriptExecutor, net *cstructs.DriverNetwork) error {
 	ops := &operations{}
 
-	t := new(TaskRegistration)
-	t.Services = make(map[string]*ServiceRegistration, len(newTask.Services))
+	taskReg := new(TaskRegistration)
+	taskReg.Services = make(map[string]*ServiceRegistration, len(newTask.Services))
 
 	existingIDs := make(map[string]*structs.Service, len(existing.Services))
 	for _, s := range existing.Services {
@@ -745,7 +749,7 @@ func (c *ServiceClient) UpdateTask(allocID string, existing, newTask *structs.Ta
 			serviceID: existingID,
 			checkIDs:  make(map[string]struct{}, len(newSvc.Checks)),
 		}
-		t.Services[existingID] = sreg
+		taskReg.Services[existingID] = sreg
 
 		// PortLabel and AddressMode aren't included in the ID, so we
 		// have to compare manually.
@@ -755,7 +759,7 @@ func (c *ServiceClient) UpdateTask(allocID string, existing, newTask *structs.Ta
 			delete(newIDs, existingID)
 		}
 
-		// Check to see what checks were updated
+		// See what checks were updated
 		existingChecks := make(map[string]*structs.ServiceCheck, len(existingSvc.Checks))
 		for _, check := range existingSvc.Checks {
 			existingChecks[makeCheckID(existingID, check)] = check
@@ -806,11 +810,11 @@ func (c *ServiceClient) UpdateTask(allocID string, existing, newTask *structs.Ta
 			return err
 		}
 
-		t.Services[sreg.serviceID] = sreg
+		taskReg.Services[sreg.serviceID] = sreg
 	}
 
 	// Add the task to the allocation's registration
-	c.addTaskRegistration(allocID, newTask.Name, t)
+	c.addTaskRegistration(allocID, newTask.Name, taskReg)
 
 	c.commit(ops)
 
@@ -1079,3 +1083,44 @@ func isNomadService(id string) bool {
 	const prefix = nomadServicePrefix + "-executor"
 	return strings.HasPrefix(id, prefix)
 }
+
+// getAddress returns the ip and port to use for a service or check. An error
+// is returned if an ip and port cannot be determined.
+func getAddress(addrMode, portLabel string, networks structs.Networks, driverNet *cstructs.DriverNetwork) (string, int, error) {
+	switch addrMode {
+	case structs.AddressModeAuto:
+		if driverNet.Advertise() {
+			addrMode = structs.AddressModeDriver
+		} else {
+			addrMode = structs.AddressModeHost
+		}
+		return getAddress(addrMode, portLabel, networks, driverNet)
+	case structs.AddressModeHost:
+		// Default path: use host ip:port
+		ip, port := networks.Port(portLabel)
+		return ip, port, nil
+
+	case structs.AddressModeDriver:
+		// Require a driver network if driver address mode is used
+		if driverNet == nil {
+			return "", 0, fmt.Errorf(`cannot use address_mode="driver": no driver network exists`)
+		}
+
+		// If the port is a label, use the driver's port (not the host's)
+		if port, ok := driverNet.PortMap[portLabel]; ok {
+			return driverNet.IP, port, nil
+		}
+
+		// If port isn't a label, try to parse it as a literal port number
+		port, err := strconv.Atoi(portLabel)
+		if err != nil {
+			return "", 0, fmt.Errorf("invalid port %q: %v", portLabel, err)
+		}
+
+		return driverNet.IP, port, nil
+
+	default:
+		// Shouldn't happen due to validation, but enforce invariants
+		return "", 0, fmt.Errorf("invalid address mode %q", addrMode)
+	}
+}

command/agent/consul/unit_test.go

@@ -16,6 +16,7 @@ import (
 	cstructs "github.com/hashicorp/nomad/client/structs"
 	"github.com/hashicorp/nomad/nomad/structs"
 	"github.com/kr/pretty"
+	"github.com/stretchr/testify/assert"
 )
 
 const (
@@ -1440,3 +1441,141 @@ func TestCreateCheckReg(t *testing.T) {
 		t.Fatalf("diff:\n%s\n", strings.Join(diff, "\n"))
 	}
 }
+
+// TestGetAddress asserts Nomad uses the correct ip and port for services and
+// checks depending on port labels, driver networks, and address mode.
+func TestGetAddress(t *testing.T) {
+	const HostIP = "127.0.0.1"
+
+	cases := []struct {
+		Name string
+
+		// Parameters
+		Mode      string
+		PortLabel string
+		Host      map[string]int // will be converted to structs.Networks
+		Driver    *cstructs.DriverNetwork
+
+		// Results
+		IP          string
+		Port        int
+		ErrContains string
+	}{
+		{
+			Name:      "ExampleService",
+			Mode:      structs.AddressModeAuto,
+			PortLabel: "db",
+			Host:      map[string]int{"db": 12435},
+			Driver: &cstructs.DriverNetwork{
+				PortMap: map[string]int{"db": 6379},
+				IP:      "10.1.2.3",
+			},
+			IP:   HostIP,
+			Port: 12435,
+		},
+		{
+			Name:      "Host",
+			Mode:      structs.AddressModeHost,
+			PortLabel: "db",
+			Host:      map[string]int{"db": 12345},
+			Driver: &cstructs.DriverNetwork{
+				PortMap: map[string]int{"db": 6379},
+				IP:      "10.1.2.3",
+			},
+			IP:   HostIP,
+			Port: 12345,
+		},
+		{
+			Name:      "Driver",
+			Mode:      structs.AddressModeDriver,
+			PortLabel: "db",
+			Host:      map[string]int{"db": 12345},
+			Driver: &cstructs.DriverNetwork{
+				PortMap: map[string]int{"db": 6379},
+				IP:      "10.1.2.3",
+			},
+			IP:   "10.1.2.3",
+			Port: 6379,
+		},
+		{
+			Name:      "AutoDriver",
+			Mode:      structs.AddressModeAuto,
+			PortLabel: "db",
+			Host:      map[string]int{"db": 12345},
+			Driver: &cstructs.DriverNetwork{
+				PortMap:       map[string]int{"db": 6379},
+				IP:            "10.1.2.3",
+				AutoAdvertise: true,
+			},
+			IP:   "10.1.2.3",
+			Port: 6379,
+		},
+		{
+			Name:      "DriverCustomPort",
+			Mode:      structs.AddressModeDriver,
+			PortLabel: "7890",
+			Host:      map[string]int{"db": 12345},
+			Driver: &cstructs.DriverNetwork{
+				PortMap: map[string]int{"db": 6379},
+				IP:      "10.1.2.3",
+			},
+			IP:   "10.1.2.3",
+			Port: 7890,
+		},
+		{
+			Name:        "DriverWithoutNetwork",
+			Mode:        structs.AddressModeDriver,
+			PortLabel:   "db",
+			Host:        map[string]int{"db": 12345},
+			Driver:      nil,
+			ErrContains: "no driver network exists",
+		},
+		{
+			Name:      "DriverBadPort",
+			Mode:      structs.AddressModeDriver,
+			PortLabel: "bad-port-label",
+			Host:      map[string]int{"db": 12345},
+			Driver: &cstructs.DriverNetwork{
+				PortMap: map[string]int{"db": 6379},
+				IP:      "10.1.2.3",
+			},
+			ErrContains: "invalid port",
+		},
+		{
+			Name:        "InvalidMode",
+			Mode:        "invalid-mode",
+			ErrContains: "invalid address mode",
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.Name, func(t *testing.T) {
+			// convert host port map into a structs.Networks
+			networks := []*structs.NetworkResource{
+				{
+					IP:            HostIP,
+					ReservedPorts: make([]structs.Port, len(tc.Host)),
+				},
+			}
+
+			i := 0
+			for label, port := range tc.Host {
+				networks[0].ReservedPorts[i].Label = label
+				networks[0].ReservedPorts[i].Value = port
+				i++
+			}
+
+			// Run getAddress
+			ip, port, err := getAddress(tc.Mode, tc.PortLabel, networks, tc.Driver)
+
+			// Assert the results
+			assert.Equal(t, tc.IP, ip, "IP mismatch")
+			assert.Equal(t, tc.Port, port, "Port mismatch")
+			if tc.ErrContains == "" {
+				assert.Nil(t, err)
+			} else {
+				assert.Contains(t, err.Error(), tc.ErrContains)
+			}
+		})
+	}
+}

command/agent/job_endpoint.go

@@ -710,6 +710,7 @@ func ApiTaskToStructsTask(apiTask *api.Task, structsTask *structs.Task) {
 						Path:          check.Path,
 						Protocol:      check.Protocol,
 						PortLabel:     check.PortLabel,
+						AddressMode:   check.AddressMode,
 						Interval:      check.Interval,
 						Timeout:       check.Timeout,
 						InitialStatus: check.InitialStatus,

command/agent/job_endpoint_test.go

@@ -1222,6 +1222,7 @@ func TestJobs_ApiJobToStructsJob(t *testing.T) {
 										Path:          "/check",
 										Protocol:      "http",
 										PortLabel:     "foo",
+										AddressMode:   "driver",
 										Interval:      4 * time.Second,
 										Timeout:       2 * time.Second,
 										InitialStatus: "ok",
@@ -1418,6 +1419,7 @@ func TestJobs_ApiJobToStructsJob(t *testing.T) {
 										Path:          "/check",
 										Protocol:      "http",
 										PortLabel:     "foo",
+										AddressMode:   "driver",
 										Interval:      4 * time.Second,
 										Timeout:       2 * time.Second,
 										InitialStatus: "ok",