From c20903137db2748d3cec9a050462a7e88a786480 Mon Sep 17 00:00:00 2001
From: wweir <wweir@foxmail.com>
Date: Mon, 28 Jan 2019 16:54:27 +0800
Subject: [PATCH] Fix #5, add ipv6 support @release

---
 Makefile            |  4 ++--
 dns/dns.go          |  9 +++++----
 dns/util.go         | 23 +++++++++++++++++------
 proxy/client.go     | 15 ++++++---------
 proxy/http_proxy.go |  5 +++++
 5 files changed, 35 insertions(+), 21 deletions(-)

diff --git a/Makefile b/Makefile
index 07b17c2..26cba6e 100644
--- a/Makefile
+++ b/Makefile
@@ -16,14 +16,14 @@ kill:
 	sudo pkill -9 sower || true
 
 client: build kill
-	sudo $(PWD)/sower -f /usr/local/etc/sower.toml
+	sudo $(PWD)/sower -f conf/sower.toml
 
 server: build kill
 	$(PWD)/sower -n TCP -v 1
 
 run: build kill
 	$(PWD)/sower -n TCP -v 1  &
-	sudo $(PWD)/sower -f /usr/local/etc/sower.toml &
+	sudo $(PWD)/sower -f conf/sower.toml &
 	@sleep 1
 	curl 127.0.0.1
 	@sleep 1
diff --git a/dns/dns.go b/dns/dns.go
index ad03219..fdb7f3a 100644
--- a/dns/dns.go
+++ b/dns/dns.go
@@ -16,8 +16,9 @@ const colon = byte(':')
 
 func StartDNS(dnsServer, listenIP string) {
 	ip := net.ParseIP(listenIP)
-	suggest := &intelliSuggest{listenIP, 2 * time.Second, []string{":80", ":443"}}
+	suggest := &intelliSuggest{listenIP, 2 * time.Second, []string{"80", "443"}}
 	mem.DefaultCache = mem.New(time.Hour)
+	dnsServer = net.JoinHostPort(dnsServer, "53")
 
 	dns.HandleFunc(".", func(w dns.ResponseWriter, r *dns.Msg) {
 		// *Msg r has an TSIG record and it was validated
@@ -39,7 +40,7 @@ func StartDNS(dnsServer, listenIP string) {
 		matchAndServe(w, r, domain, listenIP, dnsServer, ip, suggest)
 	})
 
-	server := &dns.Server{Addr: listenIP + ":53", Net: "udp"}
+	server := &dns.Server{Addr: net.JoinHostPort(listenIP, "53"), Net: "udp"}
 	glog.Fatalln(server.ListenAndServe())
 }
 
@@ -58,7 +59,7 @@ func matchAndServe(w dns.ResponseWriter, r *dns.Msg, domain, listenIP,
 		go mem.Remember(suggest, domain)
 	}
 
-	msg, err := dns.Exchange(r, dnsServer+":53")
+	msg, err := dns.Exchange(r, dnsServer)
 	if msg == nil { // expose any response except nil
 		glog.V(1).Infof("get dns of %s fail: %s", domain, err)
 		return
@@ -83,7 +84,7 @@ func (i *intelliSuggest) GetOne(domain interface{}) (iface interface{}, e error)
 
 	for _, port := range i.ports {
 		// give local dial a hand, make it not so easy to be added into suggestions
-		util.HTTPPing(addr+port, addr, i.timeout/4)
+		util.HTTPPing(addr+port, addr, i.timeout/5)
 		localCh := util.HTTPPing(addr+port, addr, i.timeout)
 		remoteCh := util.HTTPPing(i.listenIP+port, addr, i.timeout)
 
diff --git a/dns/util.go b/dns/util.go
index 52eda50..c3e9788 100644
--- a/dns/util.go
+++ b/dns/util.go
@@ -16,15 +16,19 @@ var (
 )
 
 func init() {
+	host, _, _ := net.SplitHostPort(conf.Conf.ServerAddr)
+
 	//first init
 	blockList = loadRules("block", conf.Conf.BlockList)
-	whiteList = loadRules("white", conf.Conf.WhiteList)
 	suggestList = loadRules("suggest", conf.Conf.BlockList)
+	whiteList = loadRules("white", conf.Conf.WhiteList)
+	whiteList.Add(host)
 
 	conf.OnRefreash = append(conf.OnRefreash, func() error {
 		blockList = loadRules("block", conf.Conf.BlockList)
-		whiteList = loadRules("white", conf.Conf.WhiteList)
 		suggestList = loadRules("suggest", conf.Conf.Suggestions)
+		whiteList = loadRules("white", conf.Conf.WhiteList)
+		whiteList.Add(host)
 		return nil
 	})
 }
@@ -38,9 +42,16 @@ func loadRules(name string, list []string) *util.Node {
 func localA(r *dns.Msg, domain string, localIP net.IP) *dns.Msg {
 	m := new(dns.Msg)
 	m.SetReply(r)
-	m.Answer = []dns.RR{&dns.A{
-		Hdr: dns.RR_Header{Name: domain, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 20},
-		A:   localIP,
-	}}
+	if localIP.To4() != nil {
+		m.Answer = []dns.RR{&dns.A{
+			Hdr: dns.RR_Header{Name: domain, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 20},
+			A:   localIP,
+		}}
+	} else {
+		m.Answer = []dns.RR{&dns.AAAA{
+			Hdr:  dns.RR_Header{Name: domain, Rrtype: dns.TypeAAAA, Class: dns.ClassINET, Ttl: 20},
+			AAAA: localIP,
+		}}
+	}
 	return m
 }
diff --git a/proxy/client.go b/proxy/client.go
index 70e5197..5a03de6 100644
--- a/proxy/client.go
+++ b/proxy/client.go
@@ -2,7 +2,6 @@ package proxy
 
 import (
 	"net"
-	"strings"
 
 	"github.com/golang/glog"
 	"github.com/wweir/sower/proxy/kcp"
@@ -30,14 +29,12 @@ func NewClient(netType string) Client {
 }
 
 func StartClient(netType, server, cipher, password, listenIP string) {
-	connCh := listenLocal(listenIP, []string{":80", ":443"})
+	connCh := listenLocal(listenIP, []string{"80", "443"})
 	client := NewClient(netType)
-	if idx := strings.Index(server, ":"); idx > 0 {
-		ips, err := net.LookupIP(server[:idx])
-		if err != nil || len(ips) == 0 {
-			glog.Fatalln(err, ips)
-		}
-		server = ips[0].String() + server[idx:]
+	if addr, err := net.ResolveTCPAddr("tcp", server); err != nil {
+		glog.Fatalln(err)
+	} else {
+		server = addr.String()
 	}
 
 	glog.Infoln("Client started.")
@@ -61,7 +58,7 @@ func listenLocal(listenIP string, ports []string) <-chan net.Conn {
 	connCh := make(chan net.Conn, 10)
 	for i := range ports {
 		go func(port string) {
-			ln, err := net.Listen("tcp", listenIP+port)
+			ln, err := net.Listen("tcp", net.JoinHostPort(listenIP, port))
 			if err != nil {
 				glog.Fatalln(err)
 			}
diff --git a/proxy/http_proxy.go b/proxy/http_proxy.go
index 376a75c..1d92f62 100644
--- a/proxy/http_proxy.go
+++ b/proxy/http_proxy.go
@@ -14,6 +14,11 @@ import (
 
 func StartHttpProxy(netType, server, cipher, password, addr string) {
 	client := NewClient(netType)
+	if addr, err := net.ResolveTCPAddr("tcp", server); err != nil {
+		glog.Fatalln(err)
+	} else {
+		server = addr.String()
+	}
 
 	srv := &http.Server{
 		Addr: addr,