From a86b9939314b90b0c39d9c8f59a925b772a23a90 Mon Sep 17 00:00:00 2001
From: fangdingjun <fangdingjun@gmail.com>
Date: Sat, 10 Oct 2015 10:38:21 +0800
Subject: [PATCH] add more functions

1. support ipv6 address in blacklist
2. listen tcp port
3. add remote address on debug log
---
 server.go | 30 +++++++++++++++++++++++++++---
 1 file changed, 27 insertions(+), 3 deletions(-)

diff --git a/server.go b/server.go
index daf58ef..a5dfdc2 100644
--- a/server.go
+++ b/server.go
@@ -52,6 +52,7 @@ func in_blacklist(m *dns.Msg) bool {
 	}
 
 	for _, rr := range m.Answer {
+		/* A */
 		if t, ok := rr.(*dns.A); ok {
 			ip := t.A.String()
 			if _, ok1 := Blacklist_ips[ip]; ok1 {
@@ -59,9 +60,20 @@ func in_blacklist(m *dns.Msg) bool {
 				return true
 			}
 		}
+
+		/* AAAA */
+		if t, ok := rr.(*dns.AAAA); ok {
+			ip := t.AAAA.String()
+			if _, ok1 := Blacklist_ips[ip]; ok1 {
+				logger.Debug("%s is in blacklist\n", ip)
+				return true
+			}
+		}
 	}
+
 	return false
 }
+
 func handleRoot(w dns.ResponseWriter, r *dns.Msg) {
 	var err error
 	var res *dns.Msg
@@ -79,7 +91,8 @@ func handleRoot(w dns.ResponseWriter, r *dns.Msg) {
 			msg.SetReply(r)
 			msg.Answer = append(msg.Answer, rr)
 			w.WriteMsg(msg)
-			logger.Debug("query %s %s %s, reply from hosts\n",
+			logger.Debug("%s query %s %s %s, reply from hosts\n",
+				w.RemoteAddr(),
 				domain,
 				dns.ClassToString[r.Question[0].Qclass],
 				dns.TypeToString[r.Question[0].Qtype],
@@ -92,7 +105,8 @@ func handleRoot(w dns.ResponseWriter, r *dns.Msg) {
 		done = 0
 		for _, sv := range Servers {
 			if sv.match(domain) {
-				logger.Debug("query %s %s %s, forward to %s:%s\n",
+				logger.Debug("%s query %s %s %s, forward to %s:%s\n",
+					w.RemoteAddr(),
 					domain,
 					dns.ClassToString[r.Question[0].Qclass],
 					dns.TypeToString[r.Question[0].Qtype],
@@ -112,7 +126,8 @@ func handleRoot(w dns.ResponseWriter, r *dns.Msg) {
 
 		if done != 1 {
 			res, err = DefaultServer.query(r)
-			logger.Debug("query %s %s %s, use default server %s:%s\n",
+			logger.Debug("%s query %s %s %s, use default server %s:%s\n",
+				w.RemoteAddr(),
 				domain,
 				dns.ClassToString[r.Question[0].Qclass],
 				dns.TypeToString[r.Question[0].Qtype],
@@ -146,6 +161,15 @@ func main() {
 
 	logger.Info("Listen on %s\n", bind_addr)
 
+	go func() {
+		/* listen tcp */
+		err := dns.ListenAndServe(bind_addr, "tcp", nil)
+		if err != nil {
+			log.Fatal(err)
+		}
+	}()
+
+	/* listen udp */
 	err := dns.ListenAndServe(bind_addr, "udp", nil)
 	if err != nil {
 		log.Fatal(err)