From eb8080df1860ace9630bb1ac64b72cb7e12dfd76 Mon Sep 17 00:00:00 2001
From: wweir <wweir@foxmail.com>
Date: Wed, 20 Mar 2019 08:45:37 +0800
Subject: [PATCH] Avoid unnecessary dhcp message

---
 dns/dhcp.go |  4 ++--
 dns/dns.go  | 45 ++++++++++++++++++++++++++++++---------------
 2 files changed, 32 insertions(+), 17 deletions(-)

diff --git a/dns/dhcp.go b/dns/dhcp.go
index 528e73f..7f2e801 100644
--- a/dns/dhcp.go
+++ b/dns/dhcp.go
@@ -29,11 +29,11 @@ func GetDefaultDNSServer() (string, error) {
 
 	var conn net.PacketConn
 	if runtime.GOOS == "windows" {
-		if conn, err = reuseport.ListenPacket("udp", iface.IP.String()+":68"); err != nil {
+		if conn, err = reuseport.ListenPacket("udp4", iface.IP.String()+":68"); err != nil {
 			return "", errors.Wrap(err, "listen dhcp")
 		}
 	} else {
-		if conn, err = reuseport.ListenPacket("udp", "0.0.0.0:68"); err != nil {
+		if conn, err = reuseport.ListenPacket("udp4", "0.0.0.0:68"); err != nil {
 			return "", errors.Wrap(err, "listen dhcp")
 		}
 	}
diff --git a/dns/dns.go b/dns/dns.go
index 0de3fb9..a17f24e 100644
--- a/dns/dns.go
+++ b/dns/dns.go
@@ -24,19 +24,7 @@ func StartDNS(dnsServer, listenIP string, suggestCh chan<- string, suggestLevel
 	if dnsServer != "" {
 		dnsServer = net.JoinHostPort(dnsServer, "53")
 	} else {
-		go func() {
-			for {
-				<-dhcpCh
-				host, err := GetDefaultDNSServer()
-				if err != nil {
-					glog.Errorln(err)
-					continue
-				}
-				// atomic action
-				dnsServer = net.JoinHostPort(host, "53")
-				glog.Infoln("set dns server to", host)
-			}
-		}()
+		go dynamicSetUpstreamDNS(listenIP, &dnsServer, dhcpCh)
 		dhcpCh <- struct{}{}
 	}
 
@@ -64,6 +52,33 @@ func StartDNS(dnsServer, listenIP string, suggestCh chan<- string, suggestLevel
 	glog.Fatalln(server.ListenAndServe())
 }
 
+func dynamicSetUpstreamDNS(listenIP string, dnsServer *string, dhcpCh <-chan struct{}) {
+	addr, _ := dns.ReverseAddr(listenIP)
+	msg := &dns.Msg{
+		MsgHdr: dns.MsgHdr{
+			Id:               dns.Id(),
+			RecursionDesired: false,
+		},
+		Question: []dns.Question{{addr, dns.TypeA, dns.ClassINET}},
+	}
+
+	for {
+		<-dhcpCh
+		if _, err := dns.Exchange(msg, *dnsServer); err == nil {
+			continue
+		}
+
+		host, err := GetDefaultDNSServer()
+		if err != nil {
+			glog.Errorln(err)
+			continue
+		}
+
+		// atomic action
+		*dnsServer = net.JoinHostPort(host, "53")
+		glog.Infoln("set dns server to", host)
+	}
+}
 func matchAndServe(w dns.ResponseWriter, r *dns.Msg, domain, listenIP, dnsServer string,
 	dhcpCh chan struct{}, ipNet net.IP, suggest *intelliSuggest) {
 
@@ -85,10 +100,10 @@ func matchAndServe(w dns.ResponseWriter, r *dns.Msg, domain, listenIP, dnsServer
 			default:
 			}
 		}
-		glog.V(1).Infof("get dns of %s fail: %s", domain, err)
+		glog.V(1).Infof("get dns of %s from %s fail: %s", domain, dnsServer, err)
 		return
 	} else if msg == nil { // expose any response except nil
-		glog.V(1).Infof("get dns of %s return empty", domain)
+		glog.V(1).Infof("get dns of %s from %s return empty", domain, dnsServer)
 		return
 	}