Update middleware.go

This commit is contained in:
Jason
2019-08-19 14:08:48 +08:00
parent caae0f04be
commit bc453cfbc1

View File

@@ -4,18 +4,13 @@ import (
"net" "net"
"strings" "strings"
trie "github.com/xjasonlyu/tun2socks/common/domain-trie"
"github.com/xjasonlyu/tun2socks/common/fakeip"
"github.com/xjasonlyu/tun2socks/log"
D "github.com/miekg/dns" D "github.com/miekg/dns"
T "github.com/xjasonlyu/tun2socks/common/domain-trie"
F "github.com/xjasonlyu/tun2socks/common/fakeip"
) )
var ( func dnsExchange(backendDNS []string, r *D.Msg) (msg *D.Msg) {
BackendDNS []string
)
func dnsExchange(r *D.Msg) (msg *D.Msg) {
defer func() { defer func() {
if msg == nil { if msg == nil {
// empty DNS response // empty DNS response
@@ -29,16 +24,17 @@ func dnsExchange(r *D.Msg) (msg *D.Msg) {
c := new(D.Client) c := new(D.Client)
c.Net = "tcp" c.Net = "tcp"
for _, dns := range BackendDNS { for _, dns := range backendDNS {
msg, _, _ = c.Exchange(r, dns) msg, _, _ = c.Exchange(r, dns)
if msg != nil { if msg != nil {
// success, exit query.
break break
} }
} }
return msg return msg
} }
func resolve(hosts *trie.Trie, pool *fakeip.Pool, r *D.Msg) (msg *D.Msg) { func resolve(hosts *T.Trie, pool *F.Pool, backendDNS []string, r *D.Msg) (msg *D.Msg) {
defer func() { defer func() {
if msg != nil { if msg != nil {
msg.SetReply(r) msg.SetReply(r)
@@ -51,14 +47,13 @@ func resolve(hosts *trie.Trie, pool *fakeip.Pool, r *D.Msg) (msg *D.Msg) {
q := r.Question[0] q := r.Question[0]
if q.Qtype != D.TypeA || q.Qclass != D.ClassINET { if q.Qtype != D.TypeA || q.Qclass != D.ClassINET {
log.Debugf("DNS Query: %v %v %v", q.Name, q.Qclass, q.Qtype) return dnsExchange(backendDNS, r)
return dnsExchange(r)
} }
return fakeResolve(pool, r) return fakeResolve(pool, r)
} }
func fakeResolve(pool *fakeip.Pool, r *D.Msg) *D.Msg { func fakeResolve(pool *F.Pool, r *D.Msg) *D.Msg {
q := r.Question[0] q := r.Question[0]
host := strings.TrimRight(q.Name, ".") host := strings.TrimRight(q.Name, ".")
@@ -73,7 +68,7 @@ func fakeResolve(pool *fakeip.Pool, r *D.Msg) *D.Msg {
return msg return msg
} }
func hostResolve(hosts *trie.Trie, r *D.Msg) *D.Msg { func hostResolve(hosts *T.Trie, r *D.Msg) *D.Msg {
if hosts == nil { if hosts == nil {
return nil return nil
} }
@@ -114,11 +109,3 @@ func hostResolve(hosts *trie.Trie, r *D.Msg) *D.Msg {
setMsgTTL(msg, dnsDefaultTTL) setMsgTTL(msg, dnsDefaultTTL)
return msg return msg
} }
func newHandler(hosts *trie.Trie, pool *fakeip.Pool) D.HandlerFunc {
return func(w D.ResponseWriter, r *D.Msg) {
msg := resolve(hosts, pool, r)
w.WriteMsg(msg)
return
}
}