From 403520f4b6a7987d31989aa1b38437246b3e1401 Mon Sep 17 00:00:00 2001
From: Alvaro Saurin <alvaro.saurin@gmail.com>
Date: Fri, 13 Feb 2015 13:18:41 +0100
Subject: [PATCH] zettio/weave#366 - move more initialization code to the
 constructor

---
 nameserver/server.go      | 39 +++++++++++++++++++++++++++++----------
 nameserver/server_test.go |  4 ++--
 2 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/nameserver/server.go b/nameserver/server.go
index 22b225bc..18c813bb 100644
--- a/nameserver/server.go
+++ b/nameserver/server.go
@@ -12,7 +12,7 @@ const (
 	DEFAULT_LOCAL_DOMAIN = "weave.local."     // The default name used for the local domain
 	DEFAULT_SERVER_PORT  = 53                 // The default server port
 	DEFAULT_CLI_CFG_FILE = "/etc/resolv.conf" // default "resolv.conf" file to try to load
-	UDPBufSize           = 4096               // bigger than the default 512
+	DEFAULT_UDP_BUFLEN   = 4096               // bigger than the default 512
 )
 
 func makeDNSFailResponse(r *dns.Msg) *dns.Msg {
@@ -81,13 +81,28 @@ func NewDNSServer(config DNSServerConfig) (s *DNSServer, err error) {
 		ifaceName = config.Iface.Name
 	}
 
-	return &DNSServer{
+	s = &DNSServer{
 		config:     &config,
 		mdnsCli:    mdnsClient,
 		mdnsSrv:    mdnsServer,
 		IfaceName:  ifaceName,
 		ListenAddr: fmt.Sprintf(":%d", config.Port),
-	}, nil
+	}
+
+	// create two DNS request multiplexerers, depending on the protocol used by clients
+	// (we use the same protocol for asking upstream servers)
+	mux := func(client *dns.Client) *dns.ServeMux {
+		m := dns.NewServeMux()
+		m.HandleFunc(s.config.LocalDomain, s.queryHandler([]Lookup{s.config.Zone, s.mdnsCli}))
+		m.HandleFunc(RDNS_DOMAIN, s.rdnsHandler([]Lookup{s.config.Zone, s.mdnsCli}, client))
+		m.HandleFunc(".", s.notUsHandler(client))
+		return m
+	}
+
+	s.udpSrv = &dns.Server{Addr: s.ListenAddr, Net: "udp", Handler: mux(&dns.Client{Net: "udp", UDPSize: DEFAULT_UDP_BUFLEN})}
+	s.tcpSrv = &dns.Server{Addr: s.ListenAddr, Net: "tcp", Handler: mux(&dns.Client{Net: "tcp"})}
+
+	return
 }
 
 // Start the DNS server
@@ -108,15 +123,19 @@ func (s *DNSServer) Start() error {
 		return m
 	}
 
-	mdnsServer, err := NewMDNSServer(s.zone)
-	CheckFatal(err)
+	s.udpSrv = &dns.Server{Addr: s.ListenAddr, Net: "udp", Handler: mux(&dns.Client{Net: "udp", UDPSize: DEFAULT_UDP_BUFLEN})}
+	s.tcpSrv = &dns.Server{Addr: s.ListenAddr, Net: "tcp", Handler: mux(&dns.Client{Net: "tcp"})}
 
-	err = mdnsServer.Start(s.iface)
-	CheckFatal(err)
+	return
+}
 
-	address := fmt.Sprintf(":%d", s.port)
-	s.udpSrv = &dns.Server{Addr: address, Net: "udp", Handler: mux(&dns.Client{Net: "udp", UDPSize: UDPBufSize})}
-	s.tcpSrv = &dns.Server{Addr: address, Net: "tcp", Handler: mux(&dns.Client{Net: "tcp"})}
+// Start the DNS server
+func (s *DNSServer) Start() error {
+	Info.Printf("Using mDNS on %s", s.IfaceName)
+	err := s.mdnsCli.Start(s.config.Iface)
+	CheckFatal(err)
+	err = s.mdnsSrv.Start(s.config.Iface, s.config.LocalDomain)
+	CheckFatal(err)
 
 	wg := new(sync.WaitGroup)
 	wg.Add(2)
diff --git a/nameserver/server_test.go b/nameserver/server_test.go
index d132ae13..542bc1e8 100644
--- a/nameserver/server_test.go
+++ b/nameserver/server_test.go
@@ -73,7 +73,7 @@ func TestUDPDNSServer(t *testing.T) {
 	time.Sleep(100 * time.Millisecond) // Allow sever goroutine to start
 
 	c := new(dns.Client)
-	c.UDPSize = UDPBufSize
+	c.UDPSize = DEFAULT_UDP_BUFLEN
 	m := new(dns.Msg)
 	m.SetQuestion(successTestName, dns.TypeA)
 	m.RecursionDesired = true
@@ -198,7 +198,7 @@ func TestTCPDNSServer(t *testing.T) {
 
 	t.Logf("Creating a UDP and a TCP client")
 	uc := new(dns.Client)
-	uc.UDPSize = UDPBufSize
+	uc.UDPSize = DEFAULT_UDP_BUFLEN
 	tc := new(dns.Client)
 	tc.Net = "tcp"