zabov/dns_client.go

package main

import (
	"fmt"
	"log"
	"time"

	"math/rand"
	"strings"

	"github.com/miekg/dns"
)

func SingleResolve(query *dns.Msg, d string, lfqdn string, ch chan *dns.Msg) {
	c := new(dns.Client)

	c.ReadTimeout = 500 * time.Millisecond
	c.WriteTimeout = 500 * time.Millisecond

	in, _, err := c.Exchange(query, d)
	if err != nil {
		fmt.Printf("SingleResolve: Problem with DNS %s : %s\n", d, err.Error())
		go incrementStats("DNS Problems "+d, 1)
		ch <- nil
	} else {
		go incrementStats(d, 1)
		Rcode := in.MsgHdr.Rcode
		in.SetReply(query)
		in.MsgHdr.Rcode = Rcode
		in.Authoritative = true
		in.Compress = true
		go DomainCache(lfqdn, in)
		if ZabovDebug {
			log.Println("SingleResolve: OK:", d, lfqdn)
		}
		ch <- in
	}
}

func ParallelResolve(query *dns.Msg, config string, lfqdn string) *dns.Msg {

	var resCur *dns.Msg
	var res *dns.Msg
	ch := make(chan *dns.Msg)

	dnss := oneTimeDNS(config, ZabovConfigs[config].ZabovParallelQueries)

	for _, d := range dnss {
		if ZabovDebug {
			log.Println("ParallelResolve: running SingleResolve on:", d, lfqdn)
		}
		go SingleResolve(query, d, lfqdn, ch)
	}

	if ZabovDebug {
		log.Println("ParallelResolve: wait for results...")
	}
	for range dnss {
		resCur = <-ch
		if resCur != nil {

			if res == nil {
				if ZabovDebug {
					log.Println("ParallelResolve: got first result!")
				}
				res = resCur
			} else if resCur.Rcode == dns.RcodeSuccess {
				if ZabovDebug {
					log.Println("ParallelResolve: got next result, RcodeSuccess, replacing previous...")
				}
				res = resCur
			} else {
				if ZabovDebug {
					log.Println("ParallelResolve: got next result, discarding...")
				}
			}
			if res.Rcode == dns.RcodeSuccess {
				break
			}
		}
	}
	return res
}

//ForwardQuery forwards the query to the upstream server
//first server to answer wins
//accepts config name to select the UP DNS source list
func ForwardQuery(query *dns.Msg, config string, nocache bool) *dns.Msg {
	if ZabovDebug {
		log.Println("ForwardQuery: nocache", nocache)
	}
	go incrementStats("ForwardQueries", 1)

	r := new(dns.Msg)
	r.SetReply(query)
	r.Authoritative = true

	fqdn := strings.TrimRight(query.Question[0].Name, ".")

	lfqdn := fmt.Sprintf("%d", query.Question[0].Qtype) + "." + fqdn
	if !nocache {
		if cached := GetDomainFromCache(lfqdn); cached != nil {
			go incrementStats("CacheHit", 1)
			Rcode := cached.MsgHdr.Rcode
			cached.SetReply(query)
			cached.MsgHdr.Rcode = Rcode
			cached.Authoritative = true
			if ZabovDebug {
				log.Println("ForwardQuery: CacheHit")
			}
			cached.Compress = true
			return cached
		}
	}

	for {
		// round robin with retry

		// local responder should always be available also if no internet connection
		if !NetworkUp && localresponderConfigName != config {
			time.Sleep(10 * time.Second)
			go incrementStats("Network Problems ", 1)
			continue
		}

		in := ParallelResolve(query, config, lfqdn)

		if in != nil {
			return in
		}

	}

}

func init() {

	fmt.Println("DNS client engine starting")
	NetworkUp = checkNetworkUp()

	if NetworkUp {
		fmt.Println("[OK]: Network is UP")
	} else {
		fmt.Println("[KO] Network is DOWN: system will check again in 2 minutes")
	}

}

func oneTimeDNS(config string, count int) (dns []string) {

	if count == 0 {
		count = 1
	}
	rand.Seed(time.Now().UnixNano())

	upl := ZabovConfigs[config].ZabovDNSArray

	if len(upl) < 1 {

		if len(ZabovLocalResponder) > 0 {
			fmt.Println("No DNS defined, fallback to local responder:", ZabovLocalResponder)
			return []string{ZabovLocalResponder}
		}
		fmt.Println("No DNS defined, using default 127.0.0.53:53. Hope it works!")
		return []string{"127.0.0.53:53"}
	}

	n := rand.Intn(128*len(upl)) % len(upl)

	res := []string{}
	for i := 0; i < count; i++ {
		res = append(res, upl[(n+i)%len(upl)])
	}

	check := make(map[string]int)
	for _, val := range res {
		check[val] = 1
	}

	res = []string{}
	for d, _ := range check {
		res = append(res, d)
	}
	return res
}