From 8e504da15a7a58e6441c4a6209c7052d73d7b739 Mon Sep 17 00:00:00 2001
From: kenshinx <kenshin.acs@gmail.com>
Date: Thu, 12 Feb 2015 14:54:02 +0800
Subject: [PATCH] Adaptive with IPV6 in hosts records.

---
 handler.go |  6 +-----
 hosts.go   | 25 ++++++++++++++++++-------
 2 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/handler.go b/handler.go
index 28ffa20..2e812b9 100644
--- a/handler.go
+++ b/handler.go
@@ -1,7 +1,6 @@
 package main
 
 import (
-	"net"
 	"sync"
 	"time"
 
@@ -87,8 +86,7 @@ func (h *GODNSHandler) do(Net string, w dns.ResponseWriter, req *dns.Msg) {
 
 	// Query hosts
 	if settings.Hosts.Enable && IPQuery > 0 {
-		if sip, ok := h.hosts.Get(Q.qname); ok {
-			var ip net.IP
+		if ip, ok := h.hosts.Get(Q.qname, IPQuery); ok {
 			m := new(dns.Msg)
 			m.SetReply(req)
 
@@ -100,7 +98,6 @@ func (h *GODNSHandler) do(Net string, w dns.ResponseWriter, req *dns.Msg) {
 					Class:  dns.ClassINET,
 					Ttl:    settings.Hosts.TTL,
 				}
-				ip = net.ParseIP(sip).To4()
 				a := &dns.A{rr_header, ip}
 				m.Answer = append(m.Answer, a)
 			case _IP6Query:
@@ -110,7 +107,6 @@ func (h *GODNSHandler) do(Net string, w dns.ResponseWriter, req *dns.Msg) {
 					Class:  dns.ClassINET,
 					Ttl:    settings.Hosts.TTL,
 				}
-				ip = net.ParseIP(sip).To16()
 				aaaa := &dns.AAAA{rr_header, ip}
 				m.Answer = append(m.Answer, aaaa)
 			}
diff --git a/hosts.go b/hosts.go
index 24a16ee..1a1d83c 100644
--- a/hosts.go
+++ b/hosts.go
@@ -36,18 +36,29 @@ func NewHosts(hs HostsSettings, rs RedisSettings) Hosts {
 2. Fetch hosts records from /etc/hosts file and redis per minute
 */
 
-func (h *Hosts) Get(domain string) (ip string, ok bool) {
+func (h *Hosts) Get(domain string, family int) (ip net.IP, ok bool) {
 
-	if ip, ok = h.fileHosts.Get(domain); ok {
-		return
+	var sip string
+
+	if sip, ok = h.fileHosts.Get(domain); !ok {
+		if h.redisHosts != nil {
+			sip, ok = h.redisHosts.Get(domain)
+		}
 	}
 
-	if h.redisHosts != nil {
-		ip, ok = h.redisHosts.Get(domain)
-		return
+	if sip == "" {
+		return nil, false
 	}
 
-	return ip, false
+	switch family {
+	case _IP4Query:
+		ip = net.ParseIP(sip).To4()
+	case _IP6Query:
+		ip = net.ParseIP(sip).To16()
+	default:
+		return nil, false
+	}
+	return ip, (ip != nil)
 }
 
 func (h *Hosts) refresh() {