From c1bd7169f5aa5e4a53af11338aab30563563fd41 Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Mon, 21 Apr 2014 15:33:01 -0700
Subject: [PATCH] agent: Improve DNS parser. Fixes #39.

---
 command/agent/dns.go      |  35 ++++++++-----
 command/agent/dns_test.go | 103 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 125 insertions(+), 13 deletions(-)

diff --git a/command/agent/dns.go b/command/agent/dns.go
index d1107a241c..f0b1ed9914 100644
--- a/command/agent/dns.go
+++ b/command/agent/dns.go
@@ -254,31 +254,40 @@ func (d *DNSServer) dispatch(network string, req, resp *dns.Msg) {
 
 	// The last label is either "node", "service" or a datacenter name
 PARSE:
-	if len(labels) == 0 {
+	n := len(labels)
+	if n == 0 {
 		goto INVALID
 	}
-	switch labels[len(labels)-1] {
+	switch labels[n-1] {
 	case "service":
-		// Handle lookup with and without tag
-		switch len(labels) {
-		case 2:
-			d.serviceLookup(network, datacenter, labels[0], "", req, resp)
-		case 3:
-			d.serviceLookup(network, datacenter, labels[1], labels[0], req, resp)
-		default:
+		if n == 1 {
 			goto INVALID
 		}
 
+		// Extract the service
+		service := labels[n-2]
+
+		// Support "." in the label, re-join all the parts
+		tag := ""
+		if n >= 3 {
+			tag = strings.Join(labels[:n-2], ".")
+		}
+
+		// Handle lookup with and without tag
+		d.serviceLookup(network, datacenter, service, tag, req, resp)
+
 	case "node":
-		if len(labels) != 2 {
+		if len(labels) == 1 {
 			goto INVALID
 		}
-		d.nodeLookup(network, datacenter, labels[0], req, resp)
+		// Allow a "." in the node name, just join all the parts
+		node := strings.Join(labels[:n-1], ".")
+		d.nodeLookup(network, datacenter, node, req, resp)
 
 	default:
 		// Store the DC, and re-parse
-		datacenter = labels[len(labels)-1]
-		labels = labels[:len(labels)-1]
+		datacenter = labels[n-1]
+		labels = labels[:n-1]
 		goto PARSE
 	}
 	return
diff --git a/command/agent/dns_test.go b/command/agent/dns_test.go
index 91d9f7bb74..52f9eae738 100644
--- a/command/agent/dns_test.go
+++ b/command/agent/dns_test.go
@@ -124,6 +124,48 @@ func TestDNS_NodeLookup(t *testing.T) {
 	}
 }
 
+func TestDNS_NodeLookup_PeriodName(t *testing.T) {
+	dir, srv := makeDNSServer(t)
+	defer os.RemoveAll(dir)
+	defer srv.agent.Shutdown()
+
+	// Wait for leader
+	time.Sleep(100 * time.Millisecond)
+
+	// Register node with period in name
+	args := &structs.RegisterRequest{
+		Datacenter: "dc1",
+		Node:       "foo.bar",
+		Address:    "127.0.0.1",
+	}
+	var out struct{}
+	if err := srv.agent.RPC("Catalog.Register", args, &out); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	m := new(dns.Msg)
+	m.SetQuestion("foo.bar.node.consul.", dns.TypeANY)
+
+	c := new(dns.Client)
+	addr, _ := srv.agent.config.ClientListener(srv.agent.config.Ports.DNS)
+	in, _, err := c.Exchange(m, addr.String())
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	if len(in.Answer) != 1 {
+		t.Fatalf("Bad: %#v", in)
+	}
+
+	aRec, ok := in.Answer[0].(*dns.A)
+	if !ok {
+		t.Fatalf("Bad: %#v", in.Answer[0])
+	}
+	if aRec.A.String() != "127.0.0.1" {
+		t.Fatalf("Bad: %#v", in.Answer[0])
+	}
+}
+
 func TestDNS_NodeLookup_AAAA(t *testing.T) {
 	dir, srv := makeDNSServer(t)
 	defer os.RemoveAll(dir)
@@ -270,6 +312,67 @@ func TestDNS_ServiceLookup(t *testing.T) {
 	}
 }
 
+func TestDNS_ServiceLookup_TagPeriod(t *testing.T) {
+	dir, srv := makeDNSServer(t)
+	defer os.RemoveAll(dir)
+	defer srv.agent.Shutdown()
+
+	// Wait for leader
+	time.Sleep(100 * time.Millisecond)
+
+	// Register node
+	args := &structs.RegisterRequest{
+		Datacenter: "dc1",
+		Node:       "foo",
+		Address:    "127.0.0.1",
+		Service: &structs.NodeService{
+			Service: "db",
+			Tags:    []string{"v1.master"},
+			Port:    12345,
+		},
+	}
+	var out struct{}
+	if err := srv.agent.RPC("Catalog.Register", args, &out); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	m := new(dns.Msg)
+	m.SetQuestion("v1.master.db.service.consul.", dns.TypeSRV)
+
+	c := new(dns.Client)
+	addr, _ := srv.agent.config.ClientListener(srv.agent.config.Ports.DNS)
+	in, _, err := c.Exchange(m, addr.String())
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	if len(in.Answer) != 1 {
+		t.Fatalf("Bad: %#v", in)
+	}
+
+	srvRec, ok := in.Answer[0].(*dns.SRV)
+	if !ok {
+		t.Fatalf("Bad: %#v", in.Answer[0])
+	}
+	if srvRec.Port != 12345 {
+		t.Fatalf("Bad: %#v", srvRec)
+	}
+	if srvRec.Target != "foo.node.dc1.consul." {
+		t.Fatalf("Bad: %#v", srvRec)
+	}
+
+	aRec, ok := in.Extra[0].(*dns.A)
+	if !ok {
+		t.Fatalf("Bad: %#v", in.Extra[0])
+	}
+	if aRec.Hdr.Name != "foo.node.dc1.consul." {
+		t.Fatalf("Bad: %#v", in.Extra[0])
+	}
+	if aRec.A.String() != "127.0.0.1" {
+		t.Fatalf("Bad: %#v", in.Extra[0])
+	}
+}
+
 func TestDNS_ServiceLookup_Dedup(t *testing.T) {
 	dir, srv := makeDNSServer(t)
 	defer os.RemoveAll(dir)