diff --git a/conf/conf.go b/conf/conf.go index e48a3c5..df8b1d2 100644 --- a/conf/conf.go +++ b/conf/conf.go @@ -10,8 +10,8 @@ import ( "strconv" "time" - "github.com/BurntSushi/toml" "github.com/golang/glog" + toml "github.com/pelletier/go-toml" "github.com/wweir/fsnotify" ) @@ -33,7 +33,30 @@ var Conf = struct { BlockList []string `toml:"blocklist"` Suggestions []string `toml:"suggestions"` Verbose int `toml:"verbose"` + + tree *toml.Tree // for suggestions }{} +var OnRefreash = []func() error{ + func() (err error) { + if Conf.tree, err = toml.LoadFile(Conf.ConfigFile); err != nil { + return err + } else if err = Conf.tree.Unmarshal(&Conf); err != nil { + return err + } + + Conf.ClientIPNet = net.ParseIP(Conf.ClientIP) + return flag.Set("v", strconv.Itoa(Conf.Verbose)) + }, + func() error { + if Conf.ClearDNSCache != "" { + ctx, _ := context.WithTimeout(context.TODO(), 5*time.Second) + if err := exec.CommandContext(ctx, "sh", "-c", Conf.ClearDNSCache).Run(); err != nil { + glog.Errorln(err) + } + } + return nil + }, +} func init() { flag.StringVar(&Conf.ConfigFile, "f", filepath.Dir(os.Args[0])+"/sower.toml", "config file location") @@ -60,27 +83,6 @@ func init() { watchConfigFile() } -var OnRefreash = []func() error{func() error { - if _, err := toml.DecodeFile(Conf.ConfigFile, &Conf); err != nil { - return err - } - Conf.ClientIPNet = net.ParseIP(Conf.ClientIP) - - // clear dns cache - if Conf.ClearDNSCache != "" { - ctx, _ := context.WithTimeout(context.TODO(), 5*time.Second) - if err := exec.CommandContext(ctx, "sh", "-c", Conf.ClearDNSCache).Run(); err != nil { - glog.Errorln(err) - } - } - - // for glog - if err := flag.Set("v", strconv.Itoa(Conf.Verbose)); err != nil { - return err - } - return nil -}} - func watchConfigFile() { watcher, err := fsnotify.NewWatcher() if err != nil { diff --git a/dns/dns.go b/dns/dns.go index 9b585de..de6d61f 100644 --- a/dns/dns.go +++ b/dns/dns.go @@ -70,7 +70,8 @@ func bestTry(w dns.ResponseWriter, r *dns.Msg, domain, dnsServer string, ipNet n } func manual(w dns.ResponseWriter, r *dns.Msg, domain, dnsServer string, ipNet net.IP) { - if rule.Match(strings.TrimSuffix(domain, ".")) { + if !writeList.Match(domain) && + (blockList.Match(domain) || suggestList.Match(domain)) { glog.V(2).Infof("match %s suss", domain) w.WriteMsg(localA(r, domain, ipNet)) return diff --git a/dns/suffix_tree.go b/dns/suffix_tree.go index b99b914..989fee2 100644 --- a/dns/suffix_tree.go +++ b/dns/suffix_tree.go @@ -5,16 +5,17 @@ import ( ) type Node struct { + sep string Node map[string]*Node } -func NewNode() *Node { - return &Node{Node: map[string]*Node{}} +func NewNode(sep string) *Node { + return &Node{sep: sep, Node: map[string]*Node{}} } -func NewNodeFromRule(rules ...string) *Node { - node := NewNode() +func NewNodeFromRules(sep string, rules ...string) *Node { + node := NewNode(sep) for i := range rules { - node.Add(strings.Split(rules[i], ".")) + node.Add(rules[i]) } return node } @@ -28,26 +29,32 @@ func (n *Node) string(prefix string) (out string) { } return } +func (n *Node) trim(item string) string { + return strings.TrimSuffix(item, n.sep) +} -func (n *Node) Add(secs []string) { +func (n *Node) Add(item string) { + n.add(strings.Split(n.trim(item), n.sep)) +} +func (n *Node) add(secs []string) { length := len(secs) switch length { case 0: return case 1: - n.Node[secs[length-1]] = NewNode() + n.Node[secs[length-1]] = NewNode(n.sep) default: subNode, ok := n.Node[secs[length-1]] if !ok { - subNode = NewNode() + subNode = NewNode(n.sep) n.Node[secs[length-1]] = subNode } - subNode.Add(secs[:length-1]) + subNode.add(secs[:length-1]) } } -func (n *Node) Match(addr string) bool { - return n.matchSecs(strings.Split(addr, ".")) +func (n *Node) Match(item string) bool { + return n.matchSecs(strings.Split(n.trim(item), n.sep)) } func (n *Node) matchSecs(secs []string) bool { diff --git a/dns/util.go b/dns/util.go index 9e2c9fd..dd4f6f8 100644 --- a/dns/util.go +++ b/dns/util.go @@ -5,15 +5,26 @@ import ( "github.com/wweir/sower/conf" ) -var rule *Node +var ( + blockList *Node + suggestList *Node + writeList = NewNode(".") +) func init() { - rule = NewNodeFromRule(conf.Conf.BlockList...) - glog.V(2).Infof("block rule:\n%s", rule) + //first init + blockList = loadRules("block", conf.Conf.BlockList) + suggestList = loadRules("suggest", conf.Conf.BlockList) conf.OnRefreash = append(conf.OnRefreash, func() error { - rule = NewNodeFromRule(conf.Conf.BlockList...) - glog.V(2).Infof("block rule:\n%s", rule) + blockList = loadRules("block", conf.Conf.BlockList) + suggestList = loadRules("suggest", conf.Conf.Suggestions) return nil }) } + +func loadRules(name string, list []string) *Node { + rule := NewNodeFromRules(".", list...) + glog.V(2).Infof("load %s rule:\n%s", name, rule) + return rule +} diff --git a/go.mod b/go.mod index 58ff807..6a1a418 100644 --- a/go.mod +++ b/go.mod @@ -1,9 +1,10 @@ module github.com/wweir/sower require ( - github.com/BurntSushi/toml v0.3.1 + github.com/BurntSushi/toml v0.3.1 // indirect github.com/bifurcation/mint v0.0.0-20181105071958-a14404e9a861 // indirect github.com/cheekybits/genny v1.0.0 // indirect + github.com/davecgh/go-spew v1.1.1 // indirect github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b github.com/golang/mock v1.2.0 // indirect github.com/hashicorp/golang-lru v0.5.0 // indirect @@ -15,6 +16,7 @@ require ( github.com/miekg/dns v1.1.3 github.com/onsi/ginkgo v1.7.0 // indirect github.com/onsi/gomega v1.4.3 // indirect + github.com/pelletier/go-toml v1.2.0 github.com/pkg/errors v0.8.1 github.com/templexxx/cpufeat v0.0.0-20180724012125-cef66df7f161 // indirect github.com/templexxx/xor v0.0.0-20181023030647-4e92f724b73b // indirect diff --git a/go.sum b/go.sum index f4e5dfb..c45bde5 100644 --- a/go.sum +++ b/go.sum @@ -4,6 +4,8 @@ github.com/bifurcation/mint v0.0.0-20181105071958-a14404e9a861 h1:x17NvoJaphEzay github.com/bifurcation/mint v0.0.0-20181105071958-a14404e9a861/go.mod h1:zVt7zX3K/aDCk9Tj+VM7YymsX66ERvzCJzw8rFCX2JU= github.com/cheekybits/genny v1.0.0 h1:uGGa4nei+j20rOSeDeP5Of12XVm7TGUd4dJA9RDitfE= github.com/cheekybits/genny v1.0.0/go.mod h1:+tQajlRqAUrPI7DOSpB0XAqZYtQakVtB7wXkRAgjxjQ= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/fsnotify/fsnotify v1.4.7 h1:IXs+QLmnXW2CcXuY+8Mzv/fWEsPGWxqefPtCP5CnV9I= github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo= github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b h1:VKtxabqXZkF25pY9ekfRL6a582T4P37/31XEstQ5p58= @@ -33,6 +35,8 @@ github.com/onsi/ginkgo v1.7.0 h1:WSHQ+IS43OoUrWtD1/bbclrwK8TTH5hzp+umCiuxHgs= github.com/onsi/ginkgo v1.7.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= github.com/onsi/gomega v1.4.3 h1:RE1xgDvH7imwFD45h+u2SgIfERHlS2yNG4DObb5BSKU= github.com/onsi/gomega v1.4.3/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY= +github.com/pelletier/go-toml v1.2.0 h1:T5zMGML61Wp+FlcbWjRDT7yAxhJNAiPPLOFECq181zc= +github.com/pelletier/go-toml v1.2.0/go.mod h1:5z9KED0ma1S8pY6P1sdut58dfprrGBbd/94hg7ilaic= github.com/pkg/errors v0.8.1 h1:iURUrRGxPUNPdy5/HRSm+Yj6okJ6UtLINN0Q9M4+h3I= github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/templexxx/cpufeat v0.0.0-20180724012125-cef66df7f161 h1:89CEmDvlq/F7SJEOqkIdNDGJXrQIhuIx9D2DBXjavSU= diff --git a/proxy/client.go b/proxy/client.go index edc929f..e610aed 100644 --- a/proxy/client.go +++ b/proxy/client.go @@ -31,13 +31,18 @@ func NewClient(netType string) Client { func StartClient(netType, server, cipher, password, listenIP string) { connCh := listenLocal(listenIP, []string{":80", ":443"}) client := NewClient(netType) + ips, err := net.LookupIP(server) + if err != nil || len(ips) == 0 { + glog.Fatalln(ips, err) + } + serverAddr := ips[0].String() glog.Infoln("Client started.") for { conn := <-connCh glog.V(1).Infof("new conn from (%s) to (%s)", conn.RemoteAddr(), server) - rc, err := client.Dial(server) + rc, err := client.Dial(serverAddr) if err != nil { conn.Close() glog.Errorln(err)