From 3df6ec09f30d9bcf9e1225d7922d228513ba904d Mon Sep 17 00:00:00 2001
From: bigeagle <justin.w.xd@gmail.com>
Date: Wed, 4 Feb 2015 16:21:08 +0800
Subject: [PATCH] Specify server in server-list-file

---
 examples/china.conf               |  5 ++
 godns.conf => examples/godns.conf |  2 +-
 resolver.go                       | 78 +++++++++++++++++++++----------
 settings.go                       |  6 +--
 4 files changed, 63 insertions(+), 28 deletions(-)
 create mode 100644 examples/china.conf
 rename godns.conf => examples/godns.conf (94%)

diff --git a/examples/china.conf b/examples/china.conf
new file mode 100644
index 0000000..e886276
--- /dev/null
+++ b/examples/china.conf
@@ -0,0 +1,5 @@
+server=8.8.8.8#53
+server=127.0.0.1#5553
+
+server=/baidu.com/114.114.114.114
+# refer https://github.com/felixonmars/dnsmasq-china-list
diff --git a/godns.conf b/examples/godns.conf
similarity index 94%
rename from godns.conf
rename to examples/godns.conf
index ec06255..ea6f535 100644
--- a/godns.conf
+++ b/examples/godns.conf
@@ -12,9 +12,9 @@ host = "127.0.0.1"
 port = 53
 
 [resolv]
+server-list-file = "/etc/godns.d/servers"
 resolv-file = "/etc/resolv.conf"
 timeout = 5  # 5 seconds
-domain-server-file = "/etc/godns.d/servers"
 
 [redis]
 enable = true
diff --git a/resolver.go b/resolver.go
index 438ed4c..41a93da 100644
--- a/resolver.go
+++ b/resolver.go
@@ -4,6 +4,7 @@ import (
 	"bufio"
 	"fmt"
 	"os"
+	"strconv"
 	"strings"
 	"time"
 
@@ -21,30 +22,40 @@ func (e ResolvError) Error() string {
 }
 
 type Resolver struct {
-	config        *dns.ClientConfig
+	servers       []string
 	domain_server *suffixTreeNode
+	config        *ResolvSettings
 }
 
 func NewResolver(c ResolvSettings) *Resolver {
-	var clientConfig *dns.ClientConfig
-	clientConfig, err := dns.ClientConfigFromFile(c.ResolvFile)
-	if err != nil {
-		logger.Printf(":%s is not a valid resolv.conf file\n", c.ResolvFile)
-		logger.Println(err)
-		panic(err)
+	r := &Resolver{
+		servers:       []string{},
+		domain_server: newSuffixTreeRoot(),
+		config:        &c,
 	}
-	clientConfig.Timeout = c.Timeout
 
-	domain_server := newSuffixTreeRoot()
-	r := &Resolver{clientConfig, domain_server}
-
-	if len(c.DomainServerFile) > 0 {
-		r.ReadDomainServerFile(c.DomainServerFile)
+	if len(c.ServerListFile) > 0 {
+		r.ReadServerListFile(c.ServerListFile)
+		// Debug("%v", r.servers)
 	}
+
+	if len(c.ResolvFile) > 0 {
+		clientConfig, err := dns.ClientConfigFromFile(c.ResolvFile)
+		if err != nil {
+			logger.Printf(":%s is not a valid resolv.conf file\n", c.ResolvFile)
+			logger.Println(err)
+			panic(err)
+		}
+		for _, server := range clientConfig.Servers {
+			nameserver := server + ":" + clientConfig.Port
+			r.servers = append(r.servers, nameserver)
+		}
+	}
+
 	return r
 }
 
-func (r *Resolver) ReadDomainServerFile(file string) {
+func (r *Resolver) ReadServerListFile(file string) {
 	buf, err := os.Open(file)
 	if err != nil {
 		panic("Can't open " + file)
@@ -66,15 +77,35 @@ func (r *Resolver) ReadDomainServerFile(file string) {
 		line = strings.TrimSpace(sli[1])
 
 		tokens := strings.Split(line, "/")
-		if len(tokens) != 3 {
-			continue
+		switch len(tokens) {
+		case 3:
+			domain := tokens[1]
+			ip := tokens[2]
+			if !isDomain(domain) || !isIP(ip) {
+				continue
+			}
+			r.domain_server.sinsert(strings.Split(domain, "."), ip)
+		case 1:
+			srv_port := strings.Split(line, "#")
+			if len(srv_port) > 2 {
+				continue
+			}
+
+			ip := ""
+			if ip = srv_port[0]; !isIP(ip) {
+				continue
+			}
+
+			port := "53"
+			if len(srv_port) == 2 {
+				if _, err := strconv.Atoi(srv_port[1]); err != nil {
+					continue
+				}
+				port = srv_port[1]
+			}
+			r.servers = append(r.servers, ip+":"+port)
+
 		}
-		domain := tokens[1]
-		ip := tokens[2]
-		if !isDomain(domain) || !isIP(ip) {
-			continue
-		}
-		r.domain_server.sinsert(strings.Split(domain, "."), ip)
 	}
 
 }
@@ -118,8 +149,7 @@ func (r *Resolver) Nameservers(qname string) []string {
 		ns = append(ns, nameserver)
 	}
 
-	for _, server := range r.config.Servers {
-		nameserver := server + ":" + r.config.Port
+	for _, nameserver := range r.servers {
 		ns = append(ns, nameserver)
 	}
 	return ns
diff --git a/settings.go b/settings.go
index 4c5b153..c4282ad 100644
--- a/settings.go
+++ b/settings.go
@@ -25,9 +25,9 @@ type Settings struct {
 }
 
 type ResolvSettings struct {
-	ResolvFile       string `toml:"resolv-file"`
-	DomainServerFile string `toml:"domain-server-file"`
-	Timeout          int
+	ServerListFile string `toml:"server-list-file"`
+	ResolvFile     string `toml:"resolv-file"`
+	Timeout        int
 }
 
 type DNSServerSettings struct {