From 02f8038747574f7b552284af43c48588449c329a Mon Sep 17 00:00:00 2001
From: lxc <lxcypp@163.com>
Date: Tue, 27 Nov 2018 10:37:05 +0800
Subject: [PATCH] Add container address to DNS A record. Such as
 nginx.marathon.container.mesos

---
 records/generator.go   | 5 +++++
 records/state/state.go | 9 +++++++++
 2 files changed, 14 insertions(+)

diff --git a/records/generator.go b/records/generator.go
index 8700b363..89d07af8 100644
--- a/records/generator.go
+++ b/records/generator.go
@@ -394,6 +394,7 @@ type context struct {
 	slaveID  string
 	taskIPs  []net.IP
 	slaveIPs []string
+	containerIP string
 }
 
 func (rg *RecordGenerator) taskRecord(task state.Task, f state.Framework, domain string, spec labels.Func, ipSources []string, enumFW *EnumerableFramework) {
@@ -409,6 +410,7 @@ func (rg *RecordGenerator) taskRecord(task state.Task, f state.Framework, domain
 		slaveIDTail(task.SlaveID),
 		task.IPs(ipSources...),
 		task.SlaveIPs,
+		task.ContainerIPAddress(),
 	}
 
 	// use DiscoveryInfo name if defined instead of task name
@@ -454,6 +456,9 @@ func (rg *RecordGenerator) taskContextRecord(ctx context, task state.Task, f sta
 			rg.insertTaskRR(canonical+".slave"+tail, sIPStr, A, enumTask)
 		}
 	}
+	if ctx.containerIP != "" {
+		rg.insertTaskRR(arec+".container"+tail, ctx.containerIP, A, enumTask)
+	}
 
 	// recordName generates records for ctx.taskName, given some generation chain
 	recordName := func(gen chain) { gen("_" + ctx.taskName) }
diff --git a/records/state/state.go b/records/state/state.go
index c1907e84..ee018d34 100644
--- a/records/state/state.go
+++ b/records/state/state.go
@@ -128,6 +128,15 @@ func (t *Task) IPs(srcs ...string) (ips []net.IP) {
 	return ips
 }
 
+func (t *Task) ContainerIPAddress() string{
+	for _, status := range(t.Statuses) {
+		if status.State == "TASK_RUNNING" {
+			return status.ContainerStatus.NetworkInfos[0].IPAddresses[0].IPAddress
+		}
+	}
+	return ""
+}
+
 // sources maps the string representation of IP sources to their functions.
 var sources = map[string]func(*Task) []string{
 	"host":    hostIPs,