From 6a80fc24cf4e16eb6e0f8b0494348b1e0698969d Mon Sep 17 00:00:00 2001
From: Will Howard <will@influitive.com>
Date: Tue, 5 Dec 2017 14:10:40 -0500
Subject: [PATCH] Parse the normalized container.PortMappings presented by the
 Marathon 1.5.x API Fixes #3465

---
 discovery/marathon/marathon.go      | 14 +++++-
 discovery/marathon/marathon_test.go | 75 +++++++++++++++++++++++++++++
 2 files changed, 88 insertions(+), 1 deletion(-)

diff --git a/discovery/marathon/marathon.go b/discovery/marathon/marathon.go
index 03f85fa40..7a897cfe1 100644
--- a/discovery/marathon/marathon.go
+++ b/discovery/marathon/marathon.go
@@ -222,7 +222,8 @@ type DockerContainer struct {
 
 // Container describes the runtime an app in running in.
 type Container struct {
-	Docker DockerContainer `json:"docker"`
+	Docker       DockerContainer `json:"docker"`
+	PortMappings []PortMappings  `json:"portMappings"`
 }
 
 // PortDefinitions describes which load balancer port you should access to access the service.
@@ -344,12 +345,23 @@ func targetsForApp(app *App) []model.LabelSet {
 					target[model.LabelName(ln)] = model.LabelValue(lv)
 				}
 			}
+			// Prior to Marathon 1.5 the port mappings could be found at the path
+			// "container.docker.portMappings".  When support for Marathon 1.4
+			// is dropped then this section of code can be removed.
 			if i < len(app.Container.Docker.PortMappings) {
 				for ln, lv := range app.Container.Docker.PortMappings[i].Labels {
 					ln = portMappingLabelPrefix + strutil.SanitizeLabelName(ln)
 					target[model.LabelName(ln)] = model.LabelValue(lv)
 				}
 			}
+			// In Marathon 1.5.x the container.docker.portMappings object was moved
+			// to container.portMappings.
+			if i < len(app.Container.PortMappings) {
+				for ln, lv := range app.Container.PortMappings[i].Labels {
+					ln = portMappingLabelPrefix + strutil.SanitizeLabelName(ln)
+					target[model.LabelName(ln)] = model.LabelValue(lv)
+				}
+			}
 			targets = append(targets, target)
 		}
 	}
diff --git a/discovery/marathon/marathon_test.go b/discovery/marathon/marathon_test.go
index 31ce60471..5351bdb56 100644
--- a/discovery/marathon/marathon_test.go
+++ b/discovery/marathon/marathon_test.go
@@ -472,3 +472,78 @@ func TestMarathonSDSendGroupWithoutPortDefinitions(t *testing.T) {
 		t.Fatal("Did not get a target group.")
 	}
 }
+
+func marathonTestAppListWithContainerPortMappings(labels map[string]string, runningTasks int) *AppList {
+	var (
+		task = Task{
+			ID:    "test-task-1",
+			Host:  "mesos-slave1",
+			Ports: []uint32{31000, 32000},
+		}
+		docker = DockerContainer{
+			Image: "repo/image:tag",
+		}
+		container = Container{
+			Docker: docker,
+			PortMappings: []PortMappings{
+				{Labels: labels},
+				{Labels: make(map[string]string)},
+			},
+		}
+		app = App{
+			ID:           "test-service",
+			Tasks:        []Task{task},
+			RunningTasks: runningTasks,
+			Labels:       labels,
+			Container:    container,
+		}
+	)
+	return &AppList{
+		Apps: []App{app},
+	}
+}
+
+func TestMarathonSDSendGroupWithContainerPortMappings(t *testing.T) {
+	var (
+		ch     = make(chan []*config.TargetGroup, 1)
+		client = func(client *http.Client, url, token string) (*AppList, error) {
+			return marathonTestAppListWithContainerPortMappings(marathonValidLabel, 1), nil
+		}
+	)
+	if err := testUpdateServices(client, ch); err != nil {
+		t.Fatalf("Got error: %s", err)
+	}
+	select {
+	case tgs := <-ch:
+		tg := tgs[0]
+
+		if tg.Source != "test-service" {
+			t.Fatalf("Wrong target group name: %s", tg.Source)
+		}
+		if len(tg.Targets) != 2 {
+			t.Fatalf("Wrong number of targets: %v", tg.Targets)
+		}
+		tgt := tg.Targets[0]
+		if tgt[model.AddressLabel] != "mesos-slave1:31000" {
+			t.Fatalf("Wrong target address: %s", tgt[model.AddressLabel])
+		}
+		if tgt[model.LabelName(portMappingLabelPrefix+"prometheus")] != "yes" {
+			t.Fatalf("Wrong first portMappings label from the first port: %s", tgt[model.AddressLabel])
+		}
+		if tgt[model.LabelName(portDefinitionLabelPrefix+"prometheus")] != "" {
+			t.Fatalf("Wrong first portDefinitions label from the first port: %s", tgt[model.AddressLabel])
+		}
+		tgt = tg.Targets[1]
+		if tgt[model.AddressLabel] != "mesos-slave1:32000" {
+			t.Fatalf("Wrong target address: %s", tgt[model.AddressLabel])
+		}
+		if tgt[model.LabelName(portMappingLabelPrefix+"prometheus")] != "" {
+			t.Fatalf("Wrong portMappings label from the second port: %s", tgt[model.AddressLabel])
+		}
+		if tgt[model.LabelName(portDefinitionLabelPrefix+"prometheus")] != "" {
+			t.Fatalf("Wrong portDefinitions label from the second port: %s", tgt[model.AddressLabel])
+		}
+	default:
+		t.Fatal("Did not get a target group.")
+	}
+}