Optimize suggest logic

sower-proxy · Feb 27, 2019 · 192b4f7 · 192b4f7
1 parent 2adefe0
commit 192b4f7
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 32 deletions.
diff --git a/conf/conf.go b/conf/conf.go
@@ -15,6 +15,7 @@ import (
 	"github.com/wweir/sower/util"
 )
 
+// Conf define the config items
 var Conf = struct {
 	ConfigFile string
 	NetType    string `toml:"net_type"`
@@ -34,19 +35,17 @@ var Conf = struct {
 	Suggestions []string `toml:"suggestions"`
 	Verbose     int      `toml:"verbose"`
 }{}
-var mu = &sync.Mutex{}
 
+// OnRefreash will be executed while init and write new config
 var OnRefreash = []func() error{
 	func() (err error) {
-		mu.Lock()
-		defer mu.Unlock()
-
 		f, err := os.OpenFile(Conf.ConfigFile, os.O_RDONLY, 0644)
 		if err != nil {
 			return err
 		}
 		defer f.Close()
 
+		//safe refresh config
 		file := Conf.ConfigFile
 		if err = toml.NewDecoder(f).Decode(&Conf); err != nil {
 			return err
@@ -62,13 +61,9 @@ var OnRefreash = []func() error{
 
 			switch runtime.GOOS {
 			case "windows":
-				if err := exec.CommandContext(ctx, "cmd", "/c", Conf.ClearDNSCache).Run(); err != nil {
-					glog.Errorln(err)
-				}
+				return exec.CommandContext(ctx, "cmd", "/c", Conf.ClearDNSCache).Run()
 			default:
-				if err := exec.CommandContext(ctx, "sh", "-c", Conf.ClearDNSCache).Run(); err != nil {
-					glog.Errorln(err)
-				}
+				return exec.CommandContext(ctx, "sh", "-c", Conf.ClearDNSCache).Run()
 			}
 		}
 		return nil
@@ -89,36 +84,41 @@ func init() {
 	}
 }
 
+// mu keep synchronized add rule(write), do not care read while write
+var mu = &sync.Mutex{}
+
+// AddSuggest add new domain into suggest rules
 func AddSuggest(domain string) {
 	mu.Lock()
 	defer mu.Unlock()
 
 	Conf.Suggestions = append(Conf.Suggestions, domain)
 	Conf.Suggestions = util.NewReverseSecSlice(Conf.Suggestions).Sort().Uniq()
 
-	// safe write
-	f, err := os.OpenFile(Conf.ConfigFile+"~", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
-	if err != nil {
-		glog.Errorln(err)
-		return
-	}
+	{ // safe write
+		f, err := os.OpenFile(Conf.ConfigFile+"~", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
+		if err != nil {
+			glog.Errorln(err)
+			return
+		}
 
-	if err := toml.NewEncoder(f).ArraysWithOneElementPerLine(true).Encode(Conf); err != nil {
-		glog.Errorln(err)
+		if err := toml.NewEncoder(f).ArraysWithOneElementPerLine(true).Encode(Conf); err != nil {
+			glog.Errorln(err)
+			f.Close()
+			return
+		}
 		f.Close()
-		return
-	}
-	f.Close()
 
-	if err = os.Rename(Conf.ConfigFile+"~", Conf.ConfigFile); err != nil {
-		glog.Errorln(err)
-		return
+		if err = os.Rename(Conf.ConfigFile+"~", Conf.ConfigFile); err != nil {
+			glog.Errorln(err)
+			return
+		}
 	}
 
 	// reload config
 	for i := range OnRefreash {
 		if err := OnRefreash[i](); err != nil {
-			glog.Fatalln(err)
+			glog.Errorln(err)
 		}
 	}
 }
diff --git a/dns/dhcp.go b/dns/dhcp.go
@@ -29,7 +29,7 @@ func GetDefaultDNSServer() string {
 	for _, iface := range ifaces {
 		conn, err := NewConn(iface.IP.String() + ":68")
 		if err != nil { // maybe in use
-			glog.V(1).Infoln(err)
+			glog.Errorln(err)
 			continue
 		}
 		defer conn.Close()

diff --git a/dns/dns.go b/dns/dns.go
@@ -67,24 +67,25 @@ func matchAndServe(w dns.ResponseWriter, r *dns.Msg, domain, listenIP, dnsServer
 	dhcpCh chan struct{}, ipNet net.IP, suggest *intelliSuggest) {
 
 	inWriteList := whiteList.Match(domain)
-
 	if !inWriteList && (blockList.Match(domain) || suggestList.Match(domain)) {
 		glog.V(2).Infof("match %s suss", domain)
 		w.WriteMsg(localA(r, domain, ipNet))
 		return
 	}
 
-	if !inWriteList {
-		go mem.Remember(suggest, domain)
-	}
-
 	msg, err := dns.Exchange(r, dnsServer)
 	if err != nil && dhcpCh != nil {
 		select {
 		case dhcpCh <- struct{}{}:
 		default:
 		}
 	}
+
+	// trigger suggest logic except dhcp dns server error
+	if !inWriteList {
+		go mem.Remember(suggest, domain)
+	}
+
 	if msg == nil { // expose any response except nil
 		glog.V(1).Infof("get dns of %s fail: %s", domain, err)
 		return

diff --git a/dns/util.go b/dns/util.go
@@ -23,12 +23,14 @@ func init() {
 	suggestList = loadRules("suggest", conf.Conf.Suggestions)
 	whiteList = loadRules("white", conf.Conf.WhiteList)
 	whiteList.Add(host)
+	glog.V(1).Infoln("load config")
 
 	conf.OnRefreash = append(conf.OnRefreash, func() error {
 		blockList = loadRules("block", conf.Conf.BlockList)
 		suggestList = loadRules("suggest", conf.Conf.Suggestions)
 		whiteList = loadRules("white", conf.Conf.WhiteList)
 		whiteList.Add(host)
+		glog.V(1).Infoln("reload config")
 		return nil
 	})
 }

diff --git a/proxy/client.go b/proxy/client.go
@@ -36,7 +36,6 @@ func StartClient(netType, server, cipher, password, listenIP string) {
 	glog.Infoln("Client started.")
 	for {
 		conn := <-connCh
-		glog.V(1).Infof("new conn from (%s) to (%s)", conn.RemoteAddr(), server)
 
 		if !resolved {
 			if addr, err := net.ResolveTCPAddr("tcp", server); err != nil {
@@ -46,6 +45,7 @@ func StartClient(netType, server, cipher, password, listenIP string) {
 				resolved = true
 			}
 		}
+		glog.V(1).Infof("new conn from (%s) to (%s)", conn.RemoteAddr(), server)
 
 		rc, err := client.Dial(server)
 		if err != nil {