From 529786acf6a0a5cf0a93f15d586e2d451df5df5a Mon Sep 17 00:00:00 2001
From: nadoo <287492+nadoo@users.noreply.github.com>
Date: Sun, 26 Aug 2018 22:36:14 +0800
Subject: [PATCH] dns: add `dnsalwaystcp` config. #49

---
 conf.go              |  1 +
 dns/client.go        | 13 +++++++------
 proxy/forwarder.go   |  2 +-
 strategy/strategy.go |  8 +++-----
 4 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/conf.go b/conf.go
index b2c32f1..5929f9d 100644
--- a/conf.go
+++ b/conf.go
@@ -50,6 +50,7 @@ func confInit() {
 	flag.StringVar(&conf.RulesDir, "rules-dir", "", "rule file folder")
 
 	flag.StringVar(&conf.DNS, "dns", "", "dns forwarder server listen address")
+	flag.BoolVar(&conf.DNSConfig.AlwaysTCP, "dnsalwaystcp", false, "always use tcp to query upstream dns servers no matter there's a forwarder or not")
 	flag.StringSliceUniqVar(&conf.DNSConfig.Servers, "dnsserver", []string{"8.8.8.8:53"}, "remote dns server")
 	flag.IntVar(&conf.DNSConfig.Timeout, "dnstimeout", 3, "timeout value used in multiple dnsservers switch(seconds)")
 	flag.IntVar(&conf.DNSConfig.MaxTTL, "dnsmaxttl", 1800, "maximum TTL value for entries in the CACHE(seconds)")
diff --git a/dns/client.go b/dns/client.go
index 81049c6..10f7fe3 100644
--- a/dns/client.go
+++ b/dns/client.go
@@ -18,11 +18,12 @@ type HandleFunc func(Domain, ip string) error
 
 // Config for dns
 type Config struct {
-	Servers []string
-	Timeout int
-	MaxTTL  int
-	MinTTL  int
-	Records []string
+	Servers   []string
+	Timeout   int
+	MaxTTL    int
+	MinTTL    int
+	Records   []string
+	AlwaysTCP bool
 }
 
 // Client is a dns client struct
@@ -133,7 +134,7 @@ func (c *Client) exchange(qname string, reqBytes []byte, preferTCP bool) (server
 	}
 
 	// If client uses udp and no forwarders specified, use udp
-	if !preferTCP && dialer.Addr() == "DIRECT" {
+	if !preferTCP && !c.config.AlwaysTCP && dialer.Addr() == "DIRECT" {
 		network = "udp"
 	}
 
diff --git a/proxy/forwarder.go b/proxy/forwarder.go
index 641e947..ecc8dd1 100644
--- a/proxy/forwarder.go
+++ b/proxy/forwarder.go
@@ -89,7 +89,7 @@ func (f *Forwarder) Dial(network, addr string) (c net.Conn, err error) {
 		f.IncFailures()
 		if f.Failures() >= f.MaxFailures() {
 			f.Disable()
-			log.F("[forwarder] %s reaches maxfailures, set to disabled", f.addr)
+			log.F("[forwarder] %s reaches maxfailures, set to DISABLED", f.addr)
 		}
 	}
 
diff --git a/strategy/strategy.go b/strategy/strategy.go
index 177e378..b97b16a 100644
--- a/strategy/strategy.go
+++ b/strategy/strategy.go
@@ -165,16 +165,14 @@ func (d *Dialer) onStatusChanged(fwdr *proxy.Forwarder) {
 	defer d.mu.Unlock()
 
 	if fwdr.Enabled() {
-		log.F("[strategy] %s changed from Disabled to Enabled ", fwdr.Addr())
+		log.F("[strategy] %s changed status from Disabled to Enabled ", fwdr.Addr())
 		if fwdr.Priority() == d.Priority() {
 			d.available = append(d.available, fwdr)
 		} else if fwdr.Priority() > d.Priority() {
 			d.initAvailable()
 		}
-	}
-
-	if !fwdr.Enabled() {
-		log.F("[strategy] %s changed from Enabled to Disabled", fwdr.Addr())
+	} else {
+		log.F("[strategy] %s changed status from Enabled to Disabled", fwdr.Addr())
 		for i, f := range d.available {
 			if f == fwdr {
 				d.available[i], d.available = d.available[len(d.available)-1], d.available[:len(d.available)-1]