From a2be79ddcb5181d0548c3e62912ceefa80b0b5fc Mon Sep 17 00:00:00 2001
From: George <zhoreeq@users.noreply.github.com>
Date: Sun, 8 Mar 2020 00:59:44 -0500
Subject: [PATCH] Update config, signal handling and refactoring

---
 cmd/meshname/main.go   |  4 +--
 cmd/meshnamed/main.go  | 19 ++++++++++-
 src/meshname/server.go | 72 +++++++++++++++++++++++++-----------------
 3 files changed, 63 insertions(+), 32 deletions(-)

diff --git a/cmd/meshname/main.go b/cmd/meshname/main.go
index 8a20a32..bbad1c7 100644
--- a/cmd/meshname/main.go
+++ b/cmd/meshname/main.go
@@ -3,8 +3,8 @@ package main
 import (
 	"fmt"
 	"net"
-	"strings"
 	"os"
+	"strings"
 
 	"github.com/zhoreeq/meshname/src/meshname"
 )
@@ -28,7 +28,7 @@ func main() {
 			fmt.Println("Invalid domain")
 			return
 		}
-		subDomain := labels[len(labels) - 2]
+		subDomain := labels[len(labels)-2]
 		if len(subDomain) != 26 {
 			fmt.Println("Invalid subdomain length")
 			return
diff --git a/cmd/meshnamed/main.go b/cmd/meshnamed/main.go
index 5d5063e..c99c3a5 100644
--- a/cmd/meshnamed/main.go
+++ b/cmd/meshnamed/main.go
@@ -3,8 +3,10 @@ package main
 import (
 	"net"
 	"os"
+	"os/signal"
 	"fmt"
 	"flag"
+	"syscall"
 
 	"github.com/gologme/log"
 
@@ -47,9 +49,24 @@ func main() {
 			os.Exit(1)
 		}
 
-		s.Init(logger, meshname.MeshnameOptions{ListenAddr: *listenAddr, ConfigPath: *useconffile, ValidSubnet: validSubnet})
+		s.Init(logger, *listenAddr, *useconffile, validSubnet)
 		s.Start()
+
+		c := make(chan os.Signal, 1)
+		r := make(chan os.Signal, 1)
+		signal.Notify(c, os.Interrupt, syscall.SIGTERM)
+		signal.Notify(r, os.Interrupt, syscall.SIGHUP)
+		defer s.Stop()
+		for {
+			select {
+			case _ = <-c:
+				goto exit
+			case _ = <-r:
+				s.UpdateConfig()
+			}
+		}
 	default:
 		flag.PrintDefaults()
 	}
+exit:
 }
diff --git a/src/meshname/server.go b/src/meshname/server.go
index 17a2b42..988e00b 100644
--- a/src/meshname/server.go
+++ b/src/meshname/server.go
@@ -49,24 +49,19 @@ func GenConf(target string) (string, error) {
 }
 
 type MeshnameServer struct {
-	validSubnet *net.IPNet
-	log *log.Logger
+	validSubnet                *net.IPNet
+	log                        *log.Logger
 	listenAddr, zoneConfigPath string
-	zoneConfig map[string][]dns.RR
-	dnsClient *dns.Client
+	zoneConfig                 map[string][]dns.RR
+	dnsClient                  *dns.Client
+	dnsServer                  *dns.Server
 }
 
-type MeshnameOptions struct {
-	ListenAddr, ConfigPath string
-	ValidSubnet *net.IPNet
-}
-
-func (s *MeshnameServer) Init(log *log.Logger, options interface{}) {
-	mnoptions := options.(MeshnameOptions)
+func (s *MeshnameServer) Init(log *log.Logger, listenAddr string, zoneConfigPath string, validSubnet *net.IPNet) {
 	s.log = log
-	s.listenAddr = mnoptions.ListenAddr
-	s.validSubnet = mnoptions.ValidSubnet
-	s.zoneConfigPath = mnoptions.ConfigPath
+	s.listenAddr = listenAddr
+	s.validSubnet = validSubnet
+	s.zoneConfigPath = zoneConfigPath
 	s.zoneConfig = make(map[string][]dns.RR)
 	if s.dnsClient == nil {
 		s.dnsClient = new(dns.Client)
@@ -115,11 +110,19 @@ func (s *MeshnameServer) LoadConfig() {
 	s.log.Infoln("Meshname config loaded:", s.zoneConfigPath)
 }
 
-func (s *MeshnameServer) Start() {
-	dnsServer := &dns.Server{Addr: s.listenAddr, Net: "udp"}
-	s.log.Infoln("Started meshnamed on:", s.listenAddr)
+func (s *MeshnameServer) Stop() error {
+	if s.dnsServer != nil {
+		s.dnsServer.Shutdown()
+	}
+	return nil
+}
+
+func (s *MeshnameServer) Start() error {
+	s.dnsServer = &dns.Server{Addr: s.listenAddr, Net: "udp"}
 	dns.HandleFunc(DomainZone, s.handleRequest)
-	dnsServer.ListenAndServe()
+	go s.dnsServer.ListenAndServe()
+	s.log.Infoln("Started meshnamed on:", s.listenAddr)
+	return nil
 }
 
 func (s *MeshnameServer) handleRequest(w dns.ResponseWriter, r *dns.Msg) {
@@ -135,24 +138,23 @@ func (s *MeshnameServer) handleRequest(w dns.ResponseWriter, r *dns.Msg) {
 		}
 		subDomain := labels[len(labels)-2]
 
-		resolvedAddr, err := IPFromDomain(subDomain)
-		if err != nil {
-			s.log.Debugln(err)
-			continue
-		}
-		if !s.validSubnet.Contains(resolvedAddr) {
-			s.log.Debugln("Error: subnet doesn't match")
-			continue
-		}
 		if records, ok := s.zoneConfig[subDomain]; ok {
 			for _, rec := range records {
 				if h := rec.Header(); h.Name == q.Name && h.Rrtype == q.Qtype && h.Class == q.Qclass {
 					m.Answer = append(m.Answer, rec)
 				}
 			}
-		} else if ra := w.RemoteAddr().String(); strings.HasPrefix(ra, "[::1]:") || strings.HasPrefix(ra, "127.0.0.1:") {
-			// TODO prefix whitelists ?
+		} else if s.isRemoteLookupAllowed(w.RemoteAddr()) {
 			// do remote lookups only for local clients
+			resolvedAddr, err := IPFromDomain(subDomain)
+			if err != nil {
+				s.log.Debugln(err)
+				continue
+			}
+			if !s.validSubnet.Contains(resolvedAddr) {
+				s.log.Debugln("Error: subnet doesn't match")
+				continue
+			}
 			remoteLookups[resolvedAddr.String()] = append(remoteLookups[resolvedAddr.String()], q)
 		}
 	}
@@ -170,3 +172,15 @@ func (s *MeshnameServer) handleRequest(w dns.ResponseWriter, r *dns.Msg) {
 	w.WriteMsg(m)
 }
 
+func (s *MeshnameServer) isRemoteLookupAllowed(addr net.Addr) bool {
+	// TODO prefix whitelists ?
+	ra := addr.String()
+	return strings.HasPrefix(ra, "[::1]:") || strings.HasPrefix(ra, "127.0.0.1:")
+}
+
+func (s *MeshnameServer) UpdateConfig() error {
+	s.Stop()
+	s.LoadConfig()
+	s.Start()
+	return nil
+}