From 66a5387742ccb5bbf07b6988940342fdeb1bdb8d Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Sun, 8 Jun 2014 16:01:06 -0700
Subject: [PATCH] agent: First pass supporting DNS TTLs

---
 command/agent/dns.go | 35 +++++++++++++++++++++++------------
 1 file changed, 23 insertions(+), 12 deletions(-)

diff --git a/command/agent/dns.go b/command/agent/dns.go
index 58aa4a2379..697f314fcb 100644
--- a/command/agent/dns.go
+++ b/command/agent/dns.go
@@ -334,14 +334,15 @@ RPC:
 	}
 
 	// Add the node record
-	records := d.formatNodeRecord(&out.NodeServices.Node, req.Question[0].Name, qType)
+	records := d.formatNodeRecord(&out.NodeServices.Node, req.Question[0].Name,
+		qType, d.config.NodeTTL)
 	if records != nil {
 		resp.Answer = append(resp.Answer, records...)
 	}
 }
 
 // formatNodeRecord takes a Node and returns an A, AAAA, or CNAME record
-func (d *DNSServer) formatNodeRecord(node *structs.Node, qName string, qType uint16) (records []dns.RR) {
+func (d *DNSServer) formatNodeRecord(node *structs.Node, qName string, qType uint16, ttl time.Duration) (records []dns.RR) {
 	// Parse the IP
 	ip := net.ParseIP(node.Address)
 	var ipv4 net.IP
@@ -355,7 +356,7 @@ func (d *DNSServer) formatNodeRecord(node *structs.Node, qName string, qType uin
 				Name:   qName,
 				Rrtype: dns.TypeA,
 				Class:  dns.ClassINET,
-				Ttl:    0,
+				Ttl:    uint32(ttl / time.Second),
 			},
 			A: ip,
 		}}
@@ -366,7 +367,7 @@ func (d *DNSServer) formatNodeRecord(node *structs.Node, qName string, qType uin
 				Name:   qName,
 				Rrtype: dns.TypeAAAA,
 				Class:  dns.ClassINET,
-				Ttl:    0,
+				Ttl:    uint32(ttl / time.Second),
 			},
 			AAAA: ip,
 		}}
@@ -379,7 +380,7 @@ func (d *DNSServer) formatNodeRecord(node *structs.Node, qName string, qType uin
 				Name:   qName,
 				Rrtype: dns.TypeCNAME,
 				Class:  dns.ClassINET,
-				Ttl:    0,
+				Ttl:    uint32(ttl / time.Second),
 			},
 			Target: dns.Fqdn(node.Address),
 		}
@@ -436,6 +437,16 @@ RPC:
 		return
 	}
 
+	// Determine the TTL
+	var ttl time.Duration
+	if d.config.ServiceTTL != nil {
+		var ok bool
+		ttl, ok = d.config.ServiceTTL[service]
+		if !ok {
+			ttl = d.config.ServiceTTL["*"]
+		}
+	}
+
 	// Filter out any service nodes due to health checks
 	out.Nodes = d.filterServiceNodes(out.Nodes)
 
@@ -449,10 +460,10 @@ RPC:
 
 	// Add various responses depending on the request
 	qType := req.Question[0].Qtype
-	d.serviceNodeRecords(out.Nodes, req, resp)
+	d.serviceNodeRecords(out.Nodes, req, resp, ttl)
 
 	if qType == dns.TypeSRV {
-		d.serviceSRVRecords(datacenter, out.Nodes, req, resp)
+		d.serviceSRVRecords(datacenter, out.Nodes, req, resp, ttl)
 	}
 }
 
@@ -484,7 +495,7 @@ func shuffleServiceNodes(nodes structs.CheckServiceNodes) {
 }
 
 // serviceNodeRecords is used to add the node records for a service lookup
-func (d *DNSServer) serviceNodeRecords(nodes structs.CheckServiceNodes, req, resp *dns.Msg) {
+func (d *DNSServer) serviceNodeRecords(nodes structs.CheckServiceNodes, req, resp *dns.Msg, ttl time.Duration) {
 	qName := req.Question[0].Name
 	qType := req.Question[0].Qtype
 	handled := make(map[string]struct{})
@@ -498,7 +509,7 @@ func (d *DNSServer) serviceNodeRecords(nodes structs.CheckServiceNodes, req, res
 		handled[addr] = struct{}{}
 
 		// Add the node record
-		records := d.formatNodeRecord(&node.Node, qName, qType)
+		records := d.formatNodeRecord(&node.Node, qName, qType, ttl)
 		if records != nil {
 			resp.Answer = append(resp.Answer, records...)
 		}
@@ -506,7 +517,7 @@ func (d *DNSServer) serviceNodeRecords(nodes structs.CheckServiceNodes, req, res
 }
 
 // serviceARecords is used to add the SRV records for a service lookup
-func (d *DNSServer) serviceSRVRecords(dc string, nodes structs.CheckServiceNodes, req, resp *dns.Msg) {
+func (d *DNSServer) serviceSRVRecords(dc string, nodes structs.CheckServiceNodes, req, resp *dns.Msg, ttl time.Duration) {
 	handled := make(map[string]struct{})
 	for _, node := range nodes {
 		// Avoid duplicate entries, possible if a node has
@@ -523,7 +534,7 @@ func (d *DNSServer) serviceSRVRecords(dc string, nodes structs.CheckServiceNodes
 				Name:   req.Question[0].Name,
 				Rrtype: dns.TypeSRV,
 				Class:  dns.ClassINET,
-				Ttl:    0,
+				Ttl:    uint32(ttl / time.Second),
 			},
 			Priority: 1,
 			Weight:   1,
@@ -533,7 +544,7 @@ func (d *DNSServer) serviceSRVRecords(dc string, nodes structs.CheckServiceNodes
 		resp.Answer = append(resp.Answer, srvRec)
 
 		// Add the extra record
-		records := d.formatNodeRecord(&node.Node, srvRec.Target, dns.TypeANY)
+		records := d.formatNodeRecord(&node.Node, srvRec.Target, dns.TypeANY, ttl)
 		if records != nil {
 			resp.Extra = append(resp.Extra, records...)
 		}