godns/handler.go

package main

import (
	"crypto/md5"
	"fmt"
	"github.com/miekg/dns"
	"meow.tf/joker/godns/cache"
	"meow.tf/joker/godns/hosts"
	"meow.tf/joker/godns/log"
	"meow.tf/joker/godns/resolver"
	"meow.tf/joker/godns/utils"
	"net"
	"time"
)

const (
	notIPQuery = 0
	_IP4Query  = 4
	_IP6Query  = 6
)

type Handler struct {
	resolver        *resolver.Resolver
	cache, negCache cache.Cache
	hosts           hosts.Hosts
}

func NewHandler(r *resolver.Resolver, resolverCache, negCache cache.Cache, h hosts.Hosts) *Handler {
	return &Handler{r, resolverCache, negCache, h}
}

func (h *Handler) do(Net string, w dns.ResponseWriter, req *dns.Msg) {
	q := req.Question[0]
	question := resolver.Question{Name: utils.UnFqdn(q.Name), Type: dns.TypeToString[q.Qtype], Class: dns.ClassToString[q.Qclass]}

	var remote net.IP
	if Net == "tcp" {
		remote = w.RemoteAddr().(*net.TCPAddr).IP
	} else {
		remote = w.RemoteAddr().(*net.UDPAddr).IP
	}
	log.Info("%s lookup　%s", remote, question.String())

	IPQuery := h.isIPQuery(q)

	// Query hosts
	if h.hosts != nil && IPQuery > 0 {
		if ips, ttl, ok := h.hosts.Get(question.Name, IPQuery); ok {
			m := new(dns.Msg)
			m.SetReply(req)

			switch IPQuery {
			case _IP4Query:
				hdr := dns.RR_Header{
					Name:   q.Name,
					Rrtype: dns.TypeA,
					Class:  dns.ClassINET,
					Ttl:    uint32(ttl / time.Second),
				}
				for _, ip := range ips {
					m.Answer = append(m.Answer, &dns.A{Hdr: hdr, A: ip})
				}
			case _IP6Query:
				hdr := dns.RR_Header{
					Name:   q.Name,
					Rrtype: dns.TypeAAAA,
					Class:  dns.ClassINET,
					Ttl:    uint32(ttl / time.Second),
				}
				for _, ip := range ips {
					m.Answer = append(m.Answer, &dns.AAAA{Hdr: hdr, AAAA: ip})
				}
			}

			w.WriteMsg(m)
			log.Debug("%s found in hosts file", question.Name)
			return
		} else {
			log.Debug("%s didn't found in hosts file", question.Name)
		}
	}

	// Only query cache when qtype == 'A'|'AAAA' , qclass == 'IN'
	key := KeyGen(question)
	if IPQuery > 0 {
		mesg, err := h.cache.Get(key)
		if err != nil {
			if mesg, err = h.negCache.Get(key); err != nil {
				log.Debug("%s didn't hit cache", question.String())
			} else {
				log.Debug("%s hit negative cache", question.String())
				dns.HandleFailed(w, req)
				return
			}
		} else {
			log.Debug("%s hit cache", question.String())
			// we need this copy against concurrent modification of Id
			msg := *mesg
			msg.Id = req.Id
			w.WriteMsg(&msg)
			return
		}
	}

	mesg, err := h.resolver.Lookup(Net, req)

	if err != nil {
		log.Warn("Resolve query error %s", err)
		dns.HandleFailed(w, req)

		// cache the failure, too!
		if err = h.negCache.Set(key, nil); err != nil {
			log.Warn("Set %s negative cache failed: %v", question.String(), err)
		}
		return
	}

	w.WriteMsg(mesg)

	if IPQuery > 0 && len(mesg.Answer) > 0 {
		err = h.cache.Set(key, mesg)
		if err != nil {
			log.Warn("Set %s cache failed: %s", question.String(), err.Error())
		}
		log.Debug("Insert %s into cache", question.String())
	}
}

func (h *Handler) Bind(net string) func(w dns.ResponseWriter, req *dns.Msg) {
	return func(w dns.ResponseWriter, req *dns.Msg) {
		h.do(net, w, req)
	}
}

func (h *Handler) isIPQuery(q dns.Question) int {
	if q.Qclass != dns.ClassINET {
		return notIPQuery
	}

	switch q.Qtype {
	case dns.TypeA:
		return _IP4Query
	case dns.TypeAAAA:
		return _IP6Query
	default:
		return notIPQuery
	}
}

func KeyGen(q resolver.Question) string {
	h := md5.New()
	h.Write([]byte(q.String()))
	x := h.Sum(nil)
	key := fmt.Sprintf("%x", x)
	return key
}