feat: support force setting upstream dns server

cmd/sower/main.go

@@ -40,6 +40,7 @@ var (
 			Serve      string `usage:"dns server ip"`
 			Serve6     string `usage:"dns server ipv6, eg: ::1"`
 			ServeIface string `usage:"use the IP in the net interface, if serve ip not setted. eg: eth0"`
+			Upstream   string `usage:"upstream dns server"`
 			Fallback   string `default:"223.5.5.5" required:"true" usage:"fallback dns server"`
 		}
 		Socks5 struct {
@@ -136,8 +137,12 @@ func init() {
 }
 
 func main() {
-	proxyDial := GenProxyDial(conf.Remote.Type, conf.Remote.Addr, conf.Remote.Password, conf.DNS.Fallback)
-	r := router.NewRouter(conf.DNS.Serve, conf.DNS.Fallback, conf.Router.Country.MMDB, proxyDial)
+	upstreamDNS := conf.DNS.Upstream
+	if upstreamDNS == "" {
+		upstreamDNS = conf.DNS.Fallback
+	}
+	proxyDial := GenProxyDial(conf.Remote.Type, conf.Remote.Addr, conf.Remote.Password, upstreamDNS)
+	r := router.NewRouter(conf.DNS.Serve, conf.DNS.Upstream, conf.DNS.Fallback, conf.Router.Country.MMDB, proxyDial)
 	r.BlockRule = suffixtree.NewNodeFromRules(conf.Router.Block.Rules...)
 	r.DirectRule = suffixtree.NewNodeFromRules(conf.Router.Direct.Rules...)
 	r.ProxyRule = suffixtree.NewNodeFromRules(conf.Router.Proxy.Rules...)

cmd/sower/proxy.go

@@ -19,7 +19,7 @@ import (
 	"github.com/wweir/sower/transport/trojan"
 )
 
-func GenProxyDial(proxyType, proxyHost, proxyPassword, dnsFallback string) router.ProxyDialFn {
+func GenProxyDial(proxyType, proxyHost, proxyPassword, dns string) router.ProxyDialFn {
 	var proxy transport.Transport
 	var dialFn func() (net.Conn, error)
 
@@ -28,7 +28,7 @@ func GenProxyDial(proxyType, proxyHost, proxyPassword, dnsFallback string) route
 			PreferGo: true,
 			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
 				d := net.Dialer{}
-				c, err := d.DialContext(ctx, "udp", net.JoinHostPort(dnsFallback, "53"))
+				c, err := d.DialContext(ctx, "udp", net.JoinHostPort(dns, "53"))
 				if err != nil {
 					log.Warn().Err(err).Msg("dial fallback dns failed, use default dns setting")
 					c, err = d.DialContext(ctx, network, address)

router/dns.go

@@ -2,6 +2,7 @@ package router
 
 import (
 	"net"
+	"sync/atomic"
 
 	"github.com/miekg/dns"
 	"github.com/sower-proxy/deferlog/log"
@@ -66,18 +67,21 @@ func (r *Router) dnsProxyA(domain string, localIP net.IP, req *dns.Msg) *dns.Msg
 
 var (
 	upstreamAddrs []string
-	upstreamIndex = -1
-	queryCount    = 0
+	upstreamIndex int32 = -1
+	queryCount    int32 = 0
 )
 
 func (r *Router) Exchange(req *dns.Msg) (_ *dns.Msg, err error) {
-	queryCount++
-	if upstreamIndex < 0 {
-		dnsIPs, err := dhcp.GetDNSServer()
-		log.Err(err).
-			Int("queryCount", queryCount).
-			Strs("dns", dnsIPs).
-			Msg("get dns server")
+	atomic.AddInt32(&queryCount, 1)
+	if atomic.LoadInt32(&upstreamIndex) < 0 {
+		dnsIPs := []string{r.dns.upstreamDNS}
+		if r.dns.upstreamDNS == "" {
+			dnsIPs, err = dhcp.GetDNSServer()
+			log.Err(err).
+				Int32("queryCount", atomic.LoadInt32(&queryCount)).
+				Strs("dns", dnsIPs).
+				Msg("get dns server")
+		}
 
 		addrs := make([]string, 0, len(dnsIPs)+1)
 		addrs = append(addrs, net.JoinHostPort(r.dns.fallbackDNS, "53"))
@@ -87,18 +91,18 @@ func (r *Router) Exchange(req *dns.Msg) (_ *dns.Msg, err error) {
 			}
 		}
 
-		queryCount = 0
+		atomic.StoreInt32(&queryCount, 0)
 		upstreamAddrs = addrs
-		upstreamIndex = len(addrs) - 1
-		log.Info().Str("ip", upstreamAddrs[upstreamIndex]).Msg("use upstream dns")
+		atomic.StoreInt32(&upstreamIndex, int32(len(addrs)-1))
+		log.Info().Str("ip", upstreamAddrs[atomic.LoadInt32(&upstreamIndex)]).Msg("use upstream dns")
 	}
 
-	resp, err := dns.Exchange(req, upstreamAddrs[upstreamIndex])
+	resp, err := dns.Exchange(req, upstreamAddrs[atomic.LoadInt32(&upstreamIndex)])
 	if err != nil {
-		upstreamIndex--
-		if upstreamIndex >= 0 {
-			log.Info().Str("ip", upstreamAddrs[upstreamIndex]).Msg("use upstream dns")
-			resp, err = dns.Exchange(req, upstreamAddrs[upstreamIndex])
+		atomic.AddInt32(&upstreamIndex, -1)
+		if atomic.LoadInt32(&upstreamIndex) >= 0 {
+			log.Info().Str("ip", upstreamAddrs[atomic.LoadInt32(&upstreamIndex)]).Msg("use upstream dns")
+			resp, err = dns.Exchange(req, upstreamAddrs[atomic.LoadInt32(&upstreamIndex)])
 		}
 	}
 

router/router.go

@@ -14,30 +14,34 @@ import (
 	"github.com/wweir/sower/pkg/suffixtree"
 )
 
-type ProxyDialFn func(network, host string, port uint16) (net.Conn, error)
-type Router struct {
-	BlockRule  *suffixtree.Node
-	DirectRule *suffixtree.Node
-	ProxyRule  *suffixtree.Node
-	ProxyDial  ProxyDialFn
+type (
+	ProxyDialFn func(network, host string, port uint16) (net.Conn, error)
+	Router      struct {
+		BlockRule  *suffixtree.Node
+		DirectRule *suffixtree.Node
+		ProxyRule  *suffixtree.Node
+		ProxyDial  ProxyDialFn
 
-	dns struct {
-		dns.Client
-		fallbackDNS string
-		serveIP     net.IP
+		dns struct {
+			dns.Client
+			upstreamDNS string
+			fallbackDNS string
+			serveIP     net.IP
+		}
+
+		country struct {
+			*geoip2.Reader
+			cidrs []*net.IPNet
+		}
 	}
+)
 
-	country struct {
-		*geoip2.Reader
-		cidrs []*net.IPNet
-	}
-}
-
-func NewRouter(serveIP, fallbackDNS, mmdbFile string, proxyDial ProxyDialFn) *Router {
+func NewRouter(serveIP, upstreamDNS, fallbackDNS, mmdbFile string, proxyDial ProxyDialFn) *Router {
 	r := Router{
 		ProxyDial: proxyDial,
 	}
 
+	r.dns.upstreamDNS = upstreamDNS
 	r.dns.serveIP = net.ParseIP(serveIP)
 	r.dns.fallbackDNS = fallbackDNS