From 9792ef0b9709c144c479312351a0845c2eebd7c5 Mon Sep 17 00:00:00 2001
From: Jason <xjasonlyu@gmail.com>
Date: Mon, 19 Aug 2019 17:45:19 +0800
Subject: [PATCH] update fakedns backend DNS opt

---
 cmd/main_fakedns.go             |  6 +++++-
 component/fakedns/middleware.go | 16 +++++++++++-----
 component/fakedns/resolver.go   |  9 +++------
 3 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/cmd/main_fakedns.go b/cmd/main_fakedns.go
index f6ed19b..b0c2fb4 100644
--- a/cmd/main_fakedns.go
+++ b/cmd/main_fakedns.go
@@ -19,11 +19,15 @@ func init() {
 	registerInitFn(func() {
 		if *args.EnableFakeDNS {
 			var err error
-			fakeDNS, err = fakedns.NewResolver(*args.FakeDNSAddr, *args.Hosts, *args.BackendDNS)
+			fakeDNS, err = fakedns.NewResolver(*args.FakeDNSAddr, *args.Hosts)
 			if err != nil {
 				log.Fatalf("Create fake DNS server failed: %v", err)
 			}
 
+			// Register backend DNS
+			fakedns.RegisterDNS(*args.BackendDNS)
+
+			// Start FakeDNS
 			if err := fakeDNS.Start(); err != nil {
 				log.Fatalf("Start fake DNS failed: %v", err)
 			}
diff --git a/component/fakedns/middleware.go b/component/fakedns/middleware.go
index 7fe1088..5a82df6 100644
--- a/component/fakedns/middleware.go
+++ b/component/fakedns/middleware.go
@@ -12,7 +12,13 @@ import (
 	F "github.com/xjasonlyu/tun2socks/common/fakeip"
 )
 
-func dnsExchange(backendDNS []string, r *D.Msg) (msg *D.Msg) {
+var backendDNS []string
+
+func RegisterDNS(dns string) {
+	backendDNS = append(backendDNS, strings.Split(dns, ",")...)
+}
+
+func dnsExchange(r *D.Msg) (msg *D.Msg) {
 	defer func() {
 		if msg == nil {
 			// empty DNS response
@@ -37,7 +43,7 @@ func dnsExchange(backendDNS []string, r *D.Msg) (msg *D.Msg) {
 	return msg
 }
 
-func resolve(hosts *T.Trie, pool *F.Pool, backendDNS []string, r *D.Msg) (msg *D.Msg) {
+func resolve(hosts *T.Trie, pool *F.Pool, r *D.Msg) (msg *D.Msg) {
 	defer func() {
 		if msg != nil {
 			msg.SetReply(r)
@@ -50,7 +56,7 @@ func resolve(hosts *T.Trie, pool *F.Pool, backendDNS []string, r *D.Msg) (msg *D
 
 	q := r.Question[0]
 	if q.Qtype != D.TypeA || q.Qclass != D.ClassINET {
-		return dnsExchange(backendDNS, r)
+		return dnsExchange(r)
 	}
 
 	return fakeResolve(pool, r)
@@ -113,9 +119,9 @@ func hostResolve(hosts *T.Trie, r *D.Msg) *D.Msg {
 	return msg
 }
 
-func newHandler(hosts *T.Trie, pool *F.Pool, backendDNS []string) handler {
+func newHandler(hosts *T.Trie, pool *F.Pool) handler {
 	return func(w D.ResponseWriter, r *D.Msg) {
-		msg := resolve(hosts, pool, backendDNS, r)
+		msg := resolve(hosts, pool, r)
 		w.WriteMsg(msg)
 		return
 	}
diff --git a/component/fakedns/resolver.go b/component/fakedns/resolver.go
index 8177dc4..91861da 100644
--- a/component/fakedns/resolver.go
+++ b/component/fakedns/resolver.go
@@ -25,7 +25,6 @@ const (
 type handler = D.HandlerFunc
 
 type Resolver struct {
-	b []string
 	h handler
 	p *F.Pool
 	t *T.Trie
@@ -88,7 +87,7 @@ func (r *Resolver) Resolve(request []byte) ([]byte, error) {
 		return nil, errors.New("cannot handle dns query: invalid question length")
 	}
 
-	msg := resolve(r.t, r.p, r.b, req)
+	msg := resolve(r.t, r.p, req)
 	if msg == nil {
 		return nil, errors.New("cannot resolve dns query: msg is nil")
 	}
@@ -99,7 +98,7 @@ func (r *Resolver) Resolve(request []byte) ([]byte, error) {
 	return resp, nil
 }
 
-func NewResolver(a, h, b string) (*Resolver, error) {
+func NewResolver(a, h string) (*Resolver, error) {
 	_, ipnet, _ := net.ParseCIDR(dnsFakeIPRange)
 
 	pool, err := F.New(ipnet, dnsCacheSize)
@@ -124,10 +123,8 @@ func NewResolver(a, h, b string) (*Resolver, error) {
 		return tree
 	}(h)
 
-	backendDNS := strings.Split(b, ",")
-	handler := newHandler(tree, pool, backendDNS)
+	handler := newHandler(tree, pool)
 	return &Resolver{
-		b:         backendDNS,
 		h:         handler,
 		p:         pool,
 		t:         tree,