godns/resolver/resolver.go

package resolver

import (
	"errors"
	"fmt"
	log "github.com/sirupsen/logrus"
	"github.com/spf13/viper"
	"io"
	"meow.tf/joker/godns/utils"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"crypto/tls"
	"github.com/miekg/dns"
)

type ResolvError struct {
	qname, net  string
	nameservers []*Nameserver
}

func (e ResolvError) Error() string {
	nameservers := make([]string, len(e.nameservers))

	for i, nameserver := range e.nameservers {
		nameservers[i] = nameserver.address
	}

	return fmt.Sprintf("%s resolv failed on %s (%s)", e.qname, strings.Join(nameservers, "; "), e.net)
}

type RResp struct {
	msg        *dns.Msg
	nameserver *Nameserver
	rtt        time.Duration
}

// Resolver contains a list of nameservers, domain-specific nameservers, and dns clients
type Resolver struct {
	servers      []*Nameserver
	domainServer *suffixTreeNode
	config       *Settings

	clients map[string]*dns.Client
	clientLock sync.RWMutex
}

// NewResolver initializes a resolver from the specified settings
func NewResolver(c Settings) *Resolver {
	r := &Resolver{
		servers:      make([]*Nameserver, 0),
		domainServer: newSuffixTreeRoot(),
		config:       &c,
	}

	if len(c.ServerListFile) > 0 {
		err := r.ReadServerListFile(c.ServerListFile)

		if err != nil {
			log.WithError(err).Fatalln("Unable to read server list file")
		}
	}

	if len(c.ResolvFile) > 0 {
		clientConfig, err := dns.ClientConfigFromFile(c.ResolvFile)

		if err != nil {
			log.WithError(err).Fatalln("not a valid resolv.conf file")
		}

		for _, server := range clientConfig.Servers {
			r.servers = append(r.servers, &Nameserver{net: "udp", address: net.JoinHostPort(server, clientConfig.Port)})
		}
	}

	return r
}

// server is a configuration struct for server lists
type server struct {
	// Type is the nameserver type (https, udp, tcp-tls), optional
	Type string
	Server string
	// Optional host for passing to TLS Config
	Host string
	Domains []string
}

// parseServerListFile loads a YAML server list file.
func (r *Resolver) parseServerListFile(buf io.Reader) error {
	v := viper.New()

	var err error

	v.SetConfigType("yaml")

	if err = v.ReadConfig(buf); err != nil {
		return err
	}

	list := make([]server, 0)

	if err = v.UnmarshalKey("servers", &list); err != nil {
		return err
	}

	for _, server := range list {
		nameserver := &Nameserver{
			net: determineNet(server.Type, server.Server),
			address: server.Server,
		}

		if len(server.Domains) > 0 {
			for _, domain := range server.Domains {
				r.domainServer.sinsert(strings.Split(domain, "."), nameserver)
			}

			continue
		}

		r.servers = append(r.servers, nameserver)
	}

	return nil
}

// ReadServerListFile loads a list of server list files.
func (r *Resolver) ReadServerListFile(files []string) error {
	for _, file := range files {
		buf, err := os.Open(file)

		if err != nil {
			return err
		}

		err = r.parseServerListFile(buf)

		buf.Close()

		if err != nil {
			return err
		}
	}

	return nil
}

// Lookup will ask each nameserver in top-to-bottom fashion, starting a new request
// in every second, and return as early as possbile (have an answer).
// It returns an error if no request has succeeded.
func (r *Resolver) Lookup(net string, req *dns.Msg) (message *dns.Msg, err error) {
	if net == "udp" && r.config.SetEDNS0 {
		req = req.SetEdns0(65535, true)
	}

	qname := req.Question[0].Name

	res := make(chan *RResp, 1)
	var wg sync.WaitGroup
	L := func(resolver *Resolver, nameserver *Nameserver) {
		defer wg.Done()

		c, err := resolver.resolverFor(net, nameserver)

		if err != nil {
			log.WithError(err).Warn("resolver failed to resolve")
			return
		}

		r, rtt, err := c.Exchange(req, nameserver.address)

		if err != nil {
			log.WithFields(log.Fields{
				"error": err,
				"question": qname,
				"nameserver": nameserver.address,
			}).Warn("Socket error encountered")
			return
		}
		// If SERVFAIL happen, should return immediately and try another upstream resolver.
		// However, other Error code like NXDOMAIN is an clear response stating
		// that it has been verified no such domain existas and ask other resolvers
		// would make no sense. See more about #20
		if r != nil && r.Rcode != dns.RcodeSuccess {
			log.WithFields(log.Fields{
				"question": qname,
				"nameserver": nameserver.address,
			}).Warn("Nameserver failed to get a valid answer")

			if r.Rcode == dns.RcodeServerFailure {
				return
			}
		}
		re := &RResp{r, nameserver, rtt}
		select {
		case res <- re:
		default:
		}
	}

	ticker := time.NewTicker(time.Duration(r.config.Interval) * time.Millisecond)
	defer ticker.Stop()
	// Start lookup on each nameserver top-down, in every second
	nameservers := r.Nameservers(qname)
	for _, nameserver := range nameservers {
		wg.Add(1)
		go L(r, nameserver)
		// but exit early, if we have an answer
		select {
		case re := <-res:
			log.WithFields(log.Fields{
				"question": utils.UnFqdn(qname),
				"nameserver": re.nameserver.address,
				"rtt": re.rtt,
			}).Debug("Resolve")
			return re.msg, nil
		case <-ticker.C:
			continue
		}
	}
	// wait for all the namservers to finish
	wg.Wait()
	select {
	case re := <-res:
		log.WithFields(log.Fields{
			"question": utils.UnFqdn(qname),
			"nameserver": re.nameserver.address,
			"rtt": re.rtt,
		}).Debug("Resolve")
		return re.msg, nil
	default:
		return nil, ResolvError{qname, net, nameservers}
	}
}

func (r *Resolver) resolverFor(network string, n *Nameserver) (*dns.Client, error) {
	key := network

	// Use HTTPS if network is https, or TLS to force secure connections
	if n.net == "https" {
		key = n.net
	} else if n.net == "tcp-tls" {
		key = n.net + ":" + n.address
	}

	r.clientLock.RLock()
	client, exists := r.clients[key]
	r.clientLock.RUnlock()

	if exists {
		return client, nil
	}

	if n.net != "tcp" && n.net != "tcp-tls" && n.net != "https" && n.net != "udp" {
		return nil, errors.New("unknown network type")
	}

	timeout := r.Timeout()

	client = &dns.Client{
		Net: n.net,
		ReadTimeout: timeout,
		WriteTimeout: timeout,
	}

	if n.net == "tcp-tls" {
		host := n.host

		if host == "" {
			var err error

			host, _, err = net.SplitHostPort(n.address)

			if err != nil {
				host = n.address
			}
		}

		client.TLSConfig = &tls.Config{
			ServerName: host,
		}
	}

	r.clientLock.Lock()
	r.clients[key] = client
	r.clientLock.Lock()

	return client, nil
}

// Nameservers return the array of nameservers, with port number appended.
// '#' in the name is treated as port separator, as with dnsmasq.

func (r *Resolver) Nameservers(qname string) []*Nameserver {
	queryKeys := strings.Split(qname, ".")
	queryKeys = queryKeys[:len(queryKeys)-1] // ignore last '.'

	if v, found := r.domainServer.search(queryKeys); found {
		log.WithFields(log.Fields{
			"question": qname,
			"upstream": v.address,
		}).Debug("Found in domain server list")

		//Ensure query the specific upstream nameserver in async Lookup() function.
		return []*Nameserver{v}
	}

	return r.servers
}

func (r *Resolver) Timeout() time.Duration {
	return time.Duration(r.config.Timeout) * time.Second
}

func determineNet(t, server string) string {
	if t != "" {
		return t
	}

	if strings.HasPrefix(server, "https") {
		return "https"
	}

	return "udp"
}