finally. fix hosts query

README.MD

@@ -59,7 +59,10 @@ If multi `namerserver` set at resolv.conf, the upsteam server will try in order
 
 #### hosts
 
+Force resolv domain to assigned ip, support two types hosts configuration: locale file and remote redis
 
+#hosts file#  
+can be 
 
 
 #### cache

godns.conf

@@ -34,6 +34,9 @@ expire = 600  # 10 minutes
 maxcount = 100000
 
 [hosts]
+#If set false, will not query hosts file and redis hosts record
+enable = true
 host-file = "/etc/hosts"
 redis-key = "godns:hosts"
+ttl = 600
 

handler.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"github.com/miekg/dns"
+	"net"
 	"time"
 )
 
@@ -18,6 +19,7 @@ func (q *Question) String() string {
 type GODNSHandler struct {
 	resolver *Resolver
 	cache    Cache
+	hosts    Hosts
 }
 
 func NewHandler() *GODNSHandler {
@@ -59,18 +61,36 @@ func NewHandler() *GODNSHandler {
 		logger.Printf("Invalid cache backend %s", cacheConfig.Backend)
 		panic("Invalid cache backend")
 	}
-	return &GODNSHandler{resolver, cache}
+
+	hosts := NewHosts(settings.Hosts, settings.Redis)
+
+	return &GODNSHandler{resolver, cache, hosts}
 }
 
-func (h *GODNSHandler) do(net string, w dns.ResponseWriter, req *dns.Msg) {
+func (h *GODNSHandler) do(Net string, w dns.ResponseWriter, req *dns.Msg) {
 
 	q := req.Question[0]
-	Q := Question{q.Name, dns.TypeToString[q.Qtype], dns.ClassToString[q.Qclass]}
+	Q := Question{UnFqdn(q.Name), dns.TypeToString[q.Qtype], dns.ClassToString[q.Qclass]}
 
 	Debug("Question:　%s", Q.String())
 
-	key := KeyGen(Q)
+	// Query hosts
+	if settings.Hosts.Enable && h.isIPQuery(q) {
+		if ip, ok := h.hosts.Get(Q.qname); ok {
+			m := new(dns.Msg)
+			m.SetReply(req)
+			rr_header := dns.RR_Header{Name: q.Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: settings.Hosts.TTL}
+			a := &dns.A{rr_header, net.ParseIP(ip)}
+			m.Answer = append(m.Answer, a)
+			w.WriteMsg(m)
+			Debug("%s found in hosts", Q.qname)
+			return
+		}
+
+	}
+
 	// Only query cache when qtype == 'A' , qclass == 'IN'
+	key := KeyGen(Q)
 	if h.isIPQuery(q) {
 		mesg, err := h.cache.Get(key)
 		if err != nil {
@@ -84,7 +104,7 @@ func (h *GODNSHandler) do(net string, w dns.ResponseWriter, req *dns.Msg) {
 
 	}
 
-	mesg, err := h.resolver.Lookup(net, req)
+	mesg, err := h.resolver.Lookup(Net, req)
 
 	if err != nil {
 		Debug("%s", err)
@@ -117,3 +137,10 @@ func (h *GODNSHandler) DoUDP(w dns.ResponseWriter, req *dns.Msg) {
 func (h *GODNSHandler) isIPQuery(q dns.Question) bool {
 	return q.Qtype == dns.TypeA && q.Qclass == dns.ClassINET
 }
+
+func UnFqdn(s string) string {
+	if dns.IsFqdn(s) {
+		return s[:len(s)-1]
+	}
+	return s
+}

hosts.go

@@ -1,24 +1,89 @@
 package main
 
 import (
+	"bufio"
 	"github.com/hoisie/redis"
-	// "github.com/miekg/dns"
-	"fmt"
-	"io/ioutil"
+	"os"
+	"regexp"
+	"strings"
 )
 
-type HostsQueryFaild struct {
-	domain string
+type Hosts struct {
+	FileHosts  map[string]string
+	RedisHosts *RedisHosts
 }
 
-func (e HostsQueryFaild) Error() string {
-	return e.domain + " hosts match failed"
+func NewHosts(hs HostsSettings, rs RedisSettings) Hosts {
+	fileHosts := &FileHosts{hs.HostsFile}
+	redis := &redis.Client{Addr: rs.Addr(), Db: rs.DB, Password: rs.Password}
+	redisHosts := &RedisHosts{redis, hs.RedisKey}
+
+	hosts := Hosts{fileHosts.GetAll(), redisHosts}
+	return hosts
+
 }
 
-func readLocalHostsFile(file string) map[string]string {
+/*
+1. Resolve hosts file only one times
+2. Request redis on every get called, may lead to performance lose serious
+3. Match local /etc/hosts file first, remote redis records second
+*/
+
+func (h *Hosts) Get(domain string) (ip string, ok bool) {
+	if ip, ok = h.FileHosts[domain]; ok {
+		return
+	}
+	if ip, ok = h.RedisHosts.Get(domain); ok {
+		return
+	}
+	return "", false
+}
+
+func (h *Hosts) GetAll() map[string]string {
+
+	m := make(map[string]string)
+	for domain, ip := range h.RedisHosts.GetAll() {
+		m[domain] = ip
+	}
+	for domain, ip := range h.FileHosts {
+		m[domain] = ip
+	}
+	return m
+}
+
+type RedisHosts struct {
+	redis *redis.Client
+	key   string
+}
+
+func (r *RedisHosts) GetAll() map[string]string {
 	var hosts = make(map[string]string)
-	f, _ := os.Open(file)
-	scanner := bufio.NewScanner(f)
+	r.redis.Hgetall(r.key, hosts)
+	return hosts
+}
+
+func (r *RedisHosts) Get(domain string) (ip string, ok bool) {
+	b, err := r.redis.Hget(r.key, domain)
+	return string(b), err == nil
+}
+
+func (r *RedisHosts) Set(domain, ip string) (bool, error) {
+	return r.redis.Hset(r.key, domain, []byte(ip))
+}
+
+type FileHosts struct {
+	file string
+}
+
+func (f *FileHosts) GetAll() map[string]string {
+	var hosts = make(map[string]string)
+
+	buf, err := os.Open(f.file)
+	if err != nil {
+		panic("Can't open " + f.file)
+	}
+
+	scanner := bufio.NewScanner(buf)
 	for scanner.Scan() {
 
 		line := scanner.Text()
@@ -39,7 +104,7 @@ func readLocalHostsFile(file string) map[string]string {
 
 		domain := sli[len(sli)-1]
 		ip := sli[0]
-		if !isDomain(domain) || !isIP(ip) {
+		if !f.isDomain(domain) || !f.isIP(ip) {
 			continue
 		}
 
@@ -48,12 +113,12 @@ func readLocalHostsFile(file string) map[string]string {
 	return hosts
 }
 
-func isDomain(domain string) bool {
+func (f *FileHosts) isDomain(domain string) bool {
 	match, _ := regexp.MatchString("^[a-z]", domain)
 	return match
 }
 
-func isIP(ip string) bool {
+func (f *FileHosts) isIP(ip string) bool {
 	match, _ := regexp.MatchString("^[1-9]", ip)
 	return match
 }

resolver.go

@@ -41,7 +41,7 @@ func (r *Resolver) Lookup(net string, req *dns.Msg) (message *dns.Msg, err error
 			Debug("%s failed to get an valid answer on %s", qname, nameserver)
 			continue
 		}
-		Debug("%s resolv on %s ttl: %d", qname, nameserver, rtt)
+		Debug("%s resolv on %s ttl: %d", UnFqdn(qname), nameserver, rtt)
 		return r, nil
 	}
 	return nil, ResolvError{qname, r.Nameservers()}

settings.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"github.com/BurntSushi/toml"
 	"os"
+	"strconv"
 )
 
 var (
@@ -39,6 +40,10 @@ type RedisSettings struct {
 	Password string
 }
 
+func (s RedisSettings) Addr() string {
+	return s.Host + ":" + strconv.Itoa(s.Port)
+}
+
 type LogSettings struct {
 	File string
 }
@@ -50,8 +55,10 @@ type CacheSettings struct {
 }
 
 type HostsSettings struct {
+	Enable    bool
 	HostsFile string `toml:"host-file"`
 	RedisKey  string `toml:"redis-key"`
+	TTL       uint32 `toml:"ttl"`
 }
 
 func init() {