From dd6ca5ab925e5b234e502f0997812841c4c25f05 Mon Sep 17 00:00:00 2001
From: Michael Bridgen <mikeb@squaremobius.net>
Date: Thu, 8 Jan 2015 12:42:41 +0000
Subject: [PATCH] Factor out the particulars of looking up A or PTR queries

---
 nameserver/mdns_server.go      | 54 ++++++++++++++++++++++++++--------
 nameserver/mdns_server_test.go |  8 ++---
 2 files changed, 45 insertions(+), 17 deletions(-)

diff --git a/nameserver/mdns_server.go b/nameserver/mdns_server.go
index 71b9f2a1..0516036d 100644
--- a/nameserver/mdns_server.go
+++ b/nameserver/mdns_server.go
@@ -50,23 +50,51 @@ func (s *MDNSServer) Start(ifi *net.Interface) error {
 		return err
 	}
 
-	handleMDNS := func(w dns.ResponseWriter, r *dns.Msg) {
-		// Ignore answers to other questions
-		if len(r.Answer) == 0 && len(r.Question) > 0 {
-			q := r.Question[0]
-			if ip, err := s.zone.LookupLocal(q.Name); err == nil {
-				m := makeAddressReply(r, &q, []net.IP{ip})
-				if err = s.sendResponse(m); err != nil {
-					Warning.Printf("Error writing to %s", w)
+	handler := func(qtype uint16, lookup func(*dns.Msg, *dns.Question) *dns.Msg) dns.HandlerFunc {
+		return func(w dns.ResponseWriter, r *dns.Msg) {
+			// Handle only questions, ignore answers. We might also
+			// ignore questions that arise locally (i.e., that come
+			// from an IP we think is local), but in the interest of
+			// avoiding complication, and easier testing, this is
+			// elided on the assumption that the client wouldn't ask
+			// if it already knew the answer, and if it does ask,
+			// it'll be happy to get an answer.
+			if len(r.Answer) == 0 && len(r.Question) > 0 {
+				q := &r.Question[0]
+				if q.Qtype == qtype {
+					if m := lookup(r, q); m != nil {
+						if err = s.sendResponse(m); err != nil {
+							Warning.Printf("Error writing to %s", w)
+						}
+					} else {
+						Debug.Printf("No local answer for mDNS query %s", q.Name)
+					}
 				}
-			} else if s.addrIsLocal(w.RemoteAddr()) {
-				// ignore this - it's our own query received via multicast
-			} else {
-				Debug.Printf("Failed MDNS lookup for %s", q.Name)
 			}
 		}
 	}
-	go dns.ActivateAndServe(nil, conn, dns.HandlerFunc(handleMDNS))
+
+	handleLocal := handler(dns.TypeA, func(r *dns.Msg, q *dns.Question) *dns.Msg {
+		if ip, err := s.zone.LookupLocal(q.Name); err == nil {
+			return makeAddressReply(r, q, []net.IP{ip})
+		} else {
+			return nil
+		}
+	})
+
+	handleReverse := handler(dns.TypePTR, func(r *dns.Msg, q *dns.Question) *dns.Msg {
+		if name, err := s.zone.ReverseLookupLocal(q.Name); err == nil {
+			return makePTRReply(r, q, []string{name})
+		} else {
+			return nil
+		}
+	})
+
+	mux := dns.NewServeMux()
+	mux.HandleFunc(LOCAL_DOMAIN, handleLocal)
+	mux.HandleFunc(RDNS_DOMAIN, handleReverse)
+
+	go dns.ActivateAndServe(nil, conn, mux)
 	return err
 }
 
diff --git a/nameserver/mdns_server_test.go b/nameserver/mdns_server_test.go
index 894ed120..9b1bde54 100644
--- a/nameserver/mdns_server_test.go
+++ b/nameserver/mdns_server_test.go
@@ -34,7 +34,7 @@ func TestServerSimpleQuery(t *testing.T) {
 	log.Println("TestServerSimpleQuery starting")
 	var zone = new(ZoneDb)
 	ip, _, _ := net.ParseCIDR(testAddr1)
-	zone.AddRecord(containerID, "test.weave.", ip)
+	zone.AddRecord(containerID, "test.weave.local.", ip)
 
 	mdnsServer, err := NewMDNSServer(zone)
 	wt.AssertNoErr(t, err)
@@ -67,15 +67,15 @@ func TestServerSimpleQuery(t *testing.T) {
 
 	time.Sleep(100 * time.Millisecond) // Allow for server to get going
 
-	sendQuery("test.weave.", dns.TypeA)
+	sendQuery("test.weave.local.", dns.TypeA)
 
 	time.Sleep(time.Second)
 
 	if receivedCount != 1 {
-		t.Fatal("Unexpected result count for test.weave", receivedCount)
+		t.Fatal("Unexpected result count for test.weave.local", receivedCount)
 	}
 	if !receivedAddr.Equal(ip) {
-		t.Fatal("Unexpected result for test.weave", receivedAddr)
+		t.Fatal("Unexpected result for test.weave.local", receivedAddr)
 	}
 
 	receivedCount = 0