fix: DNS upstream connection close

.github/sower.hcl

@@ -12,7 +12,7 @@ dns {
 }
 
 socks_5 {
-    addr = ":1080"
+    addr = "127.0.0.1:1080"
 }
 
 router "block" {

.github/sower.service

@@ -6,6 +6,8 @@ After=network.target
 Type=simple
 WorkingDirectory=/etc/sower
 ExecStart=/usr/local/bin/sower -f /etc/sower/sower.hcl
+Restart=on-failure
+RestartSec=1min
 
 [Install]
 WantedBy=multi-user.target

.github/sowerd.service

@@ -7,6 +7,8 @@ Type=simple
 ExecStart=/usr/local/bin/sowerd
 Environment="FAKE_SITE=127.0.0.1:8080"
 Environment="PASSWORD=I_am_Passw0rd"
+Restart=on-failure
+RestartSec=1min
 
 [Install]
 WantedBy=multi-user.target

cmd/sower/main.go

@@ -2,11 +2,13 @@ package main
 
 import (
 	"bufio"
+	"compress/gzip"
 	"io"
 	"net"
 	"net/http"
 	"net/url"
 	"os"
+	"runtime"
 	"strconv"
 	"strings"
 	"time"
@@ -17,7 +19,8 @@ import (
 	"github.com/cristalhq/aconfig/aconfigyaml"
 	"github.com/miekg/dns"
 	"github.com/pkg/errors"
-	"github.com/wweir/deferlog/log"
+	"github.com/sower-proxy/deferlog/log"
+	"github.com/wweir/sower/pkg/suffixtree"
 	"github.com/wweir/sower/router"
 )
 
@@ -38,7 +41,7 @@ var (
 		}
 		Socks5 struct {
 			Disable bool   `default:"false" usage:"disable sock5 proxy"`
-			Addr    string `default:":1080" usage:"socks5 listen address"`
+			Addr    string `default:"127.0.0.1:1080" usage:"socks5 listen address"`
 		} `flag:"socks5"`
 
 		Router struct {
@@ -94,10 +97,12 @@ func init() {
 }
 
 func main() {
-	proxtDial := GenProxyDial(conf.Remote.Type, conf.Remote.Addr, conf.Remote.Password)
-	r := router.NewRouter(conf.DNS.Serve, conf.DNS.Fallback, conf.Router.Country.MMDB, proxtDial)
-	r.SetRules(conf.Router.Block.Rules, conf.Router.Direct.Rules, conf.Router.Proxy.Rules,
-		conf.Router.Country.Rules)
+	proxyDial := GenProxyDial(conf.Remote.Type, conf.Remote.Addr, conf.Remote.Password)
+	r := router.NewRouter(conf.DNS.Serve, conf.DNS.Fallback, conf.Router.Country.MMDB, proxyDial)
+	r.BlockRule = suffixtree.NewNodeFromRules(conf.Router.Block.Rules...)
+	r.DirectRule = suffixtree.NewNodeFromRules(conf.Router.Direct.Rules...)
+	r.ProxyRule = suffixtree.NewNodeFromRules(conf.Router.Proxy.Rules...)
+	r.AddCountryCIDRs(conf.Router.Country.Rules...)
 
 	go func() {
 		if conf.DNS.Disable {
@@ -140,28 +145,30 @@ func main() {
 	}()
 
 	start := time.Now()
-	conf.Router.Block.Rules = append(conf.Router.Block.Rules,
-		loadRules(proxtDial, conf.Router.Block.File, conf.Router.Block.FilePrefix)...)
-	conf.Router.Direct.Rules = append(conf.Router.Direct.Rules,
-		loadRules(proxtDial, conf.Router.Direct.File, conf.Router.Direct.FilePrefix)...)
-	conf.Router.Proxy.Rules = append(conf.Router.Proxy.Rules,
-		loadRules(proxtDial, conf.Router.Proxy.File, conf.Router.Proxy.FilePrefix)...)
-	conf.Router.Country.Rules = append(conf.Router.Country.Rules,
-		loadRules(proxtDial, conf.Router.Country.File, conf.Router.Country.FilePrefix)...)
-	r.SetRules(conf.Router.Block.Rules, conf.Router.Direct.Rules, conf.Router.Proxy.Rules,
-		conf.Router.Country.Rules)
+	loadRule(r.BlockRule, proxyDial, conf.Router.Block.File, conf.Router.Block.FilePrefix)
+	loadRule(r.DirectRule, proxyDial, conf.Router.Direct.File, conf.Router.Direct.FilePrefix)
+	loadRule(r.ProxyRule, proxyDial, conf.Router.Proxy.File, conf.Router.Proxy.FilePrefix)
+	for line := range featchRuleFile(proxyDial, conf.Router.Country.File) {
+		r.AddCountryCIDRs(line)
+	}
 
 	log.Info().
-		Dur("spend", time.Since(start)).
-		Int("blockRule", len(conf.Router.Block.Rules)).
-		Int("directRule", len(conf.Router.Direct.Rules)).
-		Int("proxyRule", len(conf.Router.Proxy.Rules)).
-		Int("countryRule", len(conf.Router.Country.Rules)).
+		Dur("took", time.Since(start)).
+		Uint64("blockRule", r.BlockRule.Count).
+		Uint64("directRule", r.DirectRule.Count).
+		Uint64("proxyRule", r.ProxyRule.Count).
 		Msg("Loaded rules, proxy started")
+	runtime.GC()
 	select {}
 }
 
-func loadRules(proxyDial router.ProxyDialFn, file, linePrefix string) []string {
+func loadRule(rule *suffixtree.Node, proxyDial router.ProxyDialFn, file, linePrefix string) {
+	for line := range featchRuleFile(proxyDial, file) {
+		rule.Add(linePrefix + line)
+	}
+	rule.GC()
+}
+func featchRuleFile(proxyDial router.ProxyDialFn, file string) <-chan string {
 	var loadFn func() (io.ReadCloser, error)
 	if _, err := url.Parse(file); err == nil {
 		// load rule file from remote by HTTP
@@ -176,7 +183,9 @@ func loadRules(proxyDial router.ProxyDialFn, file, linePrefix string) []string {
 		}
 
 		loadFn = func() (io.ReadCloser, error) {
-			resp, err := client.Get(file)
+			req, _ := http.NewRequest(http.MethodGet, file, nil)
+			req.Header.Add("Accept-Encoding", "gzip")
+			resp, err := client.Do(req)
 			if err != nil {
 				return nil, err
 			}
@@ -207,34 +216,37 @@ func loadRules(proxyDial router.ProxyDialFn, file, linePrefix string) []string {
 		time.Sleep(i * i * 100 * time.Millisecond)
 		rc, err = loadFn()
 	}
-	if err != nil {
-		log.Fatal().Err(err).
-			Str("file", file).
-			Msg("load config file")
-	}
-	defer rc.Close()
+	log.InfoFatal(err).
+		Str("file", file).
+		Msg("fetch rule file")
 
-	// parse rule file into rule tree
-	var lines []string
-	br := bufio.NewReader(rc)
-	for {
-		line, _, err := br.ReadLine()
-		if err == io.EOF {
-			break
-		} else if err != nil {
-			log.Error().Err(err).
-				Str("file", file).
-				Msg("read line")
-			return nil
+	ch := make(chan string, 100)
+	go func() {
+		defer rc.Close()
+		defer close(ch)
+		gr, err := gzip.NewReader(rc)
+		log.DebugFatal(err).Msg("gzip reader")
+		defer gr.Close()
+
+		br := bufio.NewReader(gr)
+		for {
+			line, _, err := br.ReadLine()
+			if err == io.EOF {
+				return
+			} else if err != nil {
+				log.Error().Err(err).
+					Str("file", file).
+					Msg("read line")
+				return
+			}
+
+			if strings.TrimSpace(string(line)) == "" {
+				continue
+			}
+
+			ch <- string(line)
 		}
+	}()
 
-		if strings.TrimSpace(string(line)) == "" {
-			continue
-		}
-
-		// use line content as suffix
-		lines = append(lines, linePrefix+string(line))
-	}
-
-	return lines
+	return ch
 }

cmd/sower/proxy.go

@@ -8,14 +8,14 @@ import (
 	"time"
 
 	"github.com/pkg/errors"
-	"github.com/wweir/deferlog/log"
-	"github.com/wweir/sower/pkg/teeconn"
+	"github.com/sower-proxy/conns/relay"
+	"github.com/sower-proxy/conns/teeconn"
+	"github.com/sower-proxy/deferlog/log"
 	"github.com/wweir/sower/router"
 	"github.com/wweir/sower/transport"
 	"github.com/wweir/sower/transport/socks5"
 	"github.com/wweir/sower/transport/sower"
 	"github.com/wweir/sower/transport/trojan"
-	"github.com/wweir/sower/util"
 )
 
 func GenProxyDial(proxyType, proxyHost, proxyPassword string) router.ProxyDialFn {
@@ -97,7 +97,7 @@ func ServeHTTP(ln net.Listener, r *router.Router) {
 	defer rc.Close()
 
 	teeconn.Stop().Reread()
-	util.Relay(teeconn, rc)
+	relay.Relay(teeconn, rc)
 	log.Debug().
 		Str("host", req.Host).
 		Dur("spend", time.Since(start)).
@@ -134,7 +134,7 @@ func ServeHTTPS(ln net.Listener, r *router.Router) {
 	defer rc.Close()
 
 	teeconn.Stop().Reread()
-	util.Relay(teeconn, rc)
+	relay.Relay(teeconn, rc)
 	log.Debug().
 		Str("host", domain).
 		Dur("spend", time.Since(start)).

cmd/sowerd/main.go

@@ -9,12 +9,12 @@ import (
 	"time"
 
 	"github.com/cristalhq/aconfig"
-	"github.com/wweir/deferlog"
-	"github.com/wweir/deferlog/log"
-	"github.com/wweir/sower/pkg/teeconn"
+	"github.com/sower-proxy/conns/relay"
+	"github.com/sower-proxy/conns/teeconn"
+	"github.com/sower-proxy/deferlog"
+	"github.com/sower-proxy/deferlog/log"
 	"github.com/wweir/sower/transport/sower"
 	"github.com/wweir/sower/transport/trojan"
-	"github.com/wweir/sower/util"
 	"golang.org/x/crypto/acme/autocert"
 )
 
@@ -124,7 +124,7 @@ func serve443(ln net.Listener, fakeSite string, sower *sower.Sower, trojan *troj
 	if addr, err = sower.Unwrap(teeconn); err == nil {
 		teeconn.Stop()
 
-		dur, err = util.RelayTo(teeconn, addr.String())
+		dur, err = relay.RelayTo(teeconn, addr.String())
 		return
 	}
 
@@ -133,11 +133,11 @@ func serve443(ln net.Listener, fakeSite string, sower *sower.Sower, trojan *troj
 	if addr, err = trojan.Unwrap(teeconn); err == nil {
 		teeconn.Stop()
 
-		dur, err = util.RelayTo(teeconn, addr.String())
+		dur, err = relay.RelayTo(teeconn, addr.String())
 		return
 	}
 
 	// 3. fallback to fake site
 	teeconn.Stop().Reread()
-	dur, err = util.RelayTo(teeconn, fakeSite)
+	dur, err = relay.RelayTo(teeconn, fakeSite)
 }

go.mod

@@ -1,6 +1,6 @@
 module github.com/wweir/sower
 
-go 1.17
+go 1.18
 
 require (
 	github.com/cristalhq/aconfig v0.16.8
@@ -8,22 +8,26 @@ require (
 	github.com/cristalhq/aconfig/aconfigtoml v0.16.1
 	github.com/cristalhq/aconfig/aconfigyaml v0.16.1
 	github.com/krolaw/dhcp4 v0.0.0-20190909130307-a50d88189771
-	github.com/libp2p/go-reuseport v0.1.0
-	github.com/miekg/dns v1.1.43
-	github.com/oschwald/geoip2-golang v1.5.0
+	github.com/libp2p/go-reuseport v0.2.0
+	github.com/miekg/dns v1.1.48
+	github.com/oschwald/geoip2-golang v1.7.0
 	github.com/pkg/errors v0.9.1
-	github.com/ulule/deepcopier v0.0.0-20200430083143-45decc6639b6
-	github.com/wweir/deferlog v0.1.0
-	golang.org/x/crypto v0.0.0-20211108221036-ceb1ce70b4fa
+	github.com/sower-proxy/conns v0.0.2
+	github.com/sower-proxy/deferlog v1.0.3
+	github.com/sower-proxy/mem v0.0.3-0.20220430030119-3c1d685997ce
+	golang.org/x/crypto v0.0.0-20220427172511-eb4f295cb31f
 )
 
 require (
-	github.com/BurntSushi/toml v0.4.1 // indirect
+	github.com/BurntSushi/toml v0.3.1 // indirect
 	github.com/hashicorp/hcl v1.0.0 // indirect
-	github.com/oschwald/maxminddb-golang v1.8.0 // indirect
-	github.com/rs/zerolog v1.23.0 // indirect
-	golang.org/x/net v0.0.0-20211109214657-ef0fda0de508 // indirect
-	golang.org/x/sys v0.0.0-20211109184856-51b60fd695b3 // indirect
-	golang.org/x/text v0.3.7 // indirect
-	gopkg.in/yaml.v2 v2.4.0 // indirect
+	github.com/oschwald/maxminddb-golang v1.9.0 // indirect
+	github.com/rs/zerolog v1.26.1 // indirect
+	golang.org/x/mod v0.4.2 // indirect
+	golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2 // indirect
+	golang.org/x/sys v0.0.0-20220422013727-9388b58f7150 // indirect
+	golang.org/x/text v0.3.6 // indirect
+	golang.org/x/tools v0.1.7 // indirect
+	golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 // indirect
+	gopkg.in/yaml.v2 v2.3.0 // indirect
 )

go.sum

@@ -1,6 +1,5 @@
+github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
-github.com/BurntSushi/toml v0.4.1 h1:GaI7EiDXDRfa8VshkTj7Fym7ha+y8/XxIgD2okUIjLw=
-github.com/BurntSushi/toml v0.4.1/go.mod h1:CxXYINrC8qIiEnFrOxCa7Jy5BFHlXnUU2pbicEuybxQ=
 github.com/coreos/go-systemd/v22 v22.3.2/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
 github.com/cristalhq/aconfig v0.16.1/go.mod h1:NXaRp+1e6bkO4dJn+wZ71xyaihMDYPtCSvEhMTm/H3E=
 github.com/cristalhq/aconfig v0.16.8 h1:lg8i0XHgfhvsnjNM5q/ou6jIHDRXlbBybjRP9t2fWuw=
@@ -11,7 +10,6 @@ github.com/cristalhq/aconfig/aconfigtoml v0.16.1 h1:q9DJfJqdLvZj7vL3bOIR6QM0ySpA
 github.com/cristalhq/aconfig/aconfigtoml v0.16.1/go.mod h1:VfRDnJBq09TZXiQdnnd4CBsNyj4s7AEDJlCxDSDt5DI=
 github.com/cristalhq/aconfig/aconfigyaml v0.16.1 h1:ghmzWFolCuja6BzBvHkNAn2Qsewu6xqdAXQyWnXjvZw=
 github.com/cristalhq/aconfig/aconfigyaml v0.16.1/go.mod h1:UgM0LkO4TKbC0s/oUP2p5POM6QiQwNvOFXostf6con0=
-github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 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/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
@@ -19,74 +17,73 @@ github.com/hashicorp/hcl v1.0.0 h1:0Anlzjpi4vEasTeNFn2mLJgTSwt0+6sfsiTG8qcWGx4=
 github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
 github.com/krolaw/dhcp4 v0.0.0-20190909130307-a50d88189771 h1:t2c2B9g1ZVhMYduqmANSEGVD3/1WlsrEYNPtVoFlENk=
 github.com/krolaw/dhcp4 v0.0.0-20190909130307-a50d88189771/go.mod h1:0AqAH3ZogsCrvrtUpvc6EtVKbc3w6xwZhkvGLuqyi3o=
-github.com/libp2p/go-reuseport v0.1.0 h1:0ooKOx2iwyIkf339WCZ2HN3ujTDbkK0PjC7JVoP1AiM=
-github.com/libp2p/go-reuseport v0.1.0/go.mod h1:bQVn9hmfcTaoo0c9v5pBhOarsU1eNOBZdaAd2hzXRKU=
-github.com/miekg/dns v1.1.43 h1:JKfpVSCB84vrAmHzyrsxB5NAr5kLoMXZArPSw7Qlgyg=
-github.com/miekg/dns v1.1.43/go.mod h1:+evo5L0630/F6ca/Z9+GAqzhjGyn8/c+TBaOyfEl0V4=
-github.com/oschwald/geoip2-golang v1.5.0 h1:igg2yQIrrcRccB1ytFXqBfOHCjXWIoMv85lVJ1ONZzw=
-github.com/oschwald/geoip2-golang v1.5.0/go.mod h1:xdvYt5xQzB8ORWFqPnqMwZpCpgNagttWdoZLlJQzg7s=
-github.com/oschwald/maxminddb-golang v1.8.0 h1:Uh/DSnGoxsyp/KYbY1AuP0tYEwfs0sCph9p/UMXK/Hk=
-github.com/oschwald/maxminddb-golang v1.8.0/go.mod h1:RXZtst0N6+FY/3qCNmZMBApR19cdQj43/NM9VkrNAis=
+github.com/libp2p/go-reuseport v0.2.0 h1:18PRvIMlpY6ZK85nIAicSBuXXvrYoSw3dsBAR7zc560=
+github.com/libp2p/go-reuseport v0.2.0/go.mod h1:bvVho6eLMm6Bz5hmU0LYN3ixd3nPPvtIlaURZZgOY4k=
+github.com/miekg/dns v1.1.48 h1:Ucfr7IIVyMBz4lRE8qmGUuZ4Wt3/ZGu9hmcMT3Uu4tQ=
+github.com/miekg/dns v1.1.48/go.mod h1:e3IlAVfNqAllflbibAZEWOXOQ+Ynzk/dDozDxY7XnME=
+github.com/oschwald/geoip2-golang v1.7.0 h1:JW1r5AKi+vv2ujSxjKthySK3jo8w8oKWPyXsw+Qs/S8=
+github.com/oschwald/geoip2-golang v1.7.0/go.mod h1:mdI/C7iK7NVMcIDDtf4bCKMJ7r0o7UwGeCo9eiitCMQ=
+github.com/oschwald/maxminddb-golang v1.9.0 h1:tIk4nv6VT9OiPyrnDAfJS1s1xKDQMZOsGojab6EjC1Y=
+github.com/oschwald/maxminddb-golang v1.9.0/go.mod h1:TK+s/Z2oZq0rSl4PSeAEoP0bgm82Cp5HyvYbt8K3zLY=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
-github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
-github.com/rs/xid v1.2.1/go.mod h1:+uKXf+4Djp6Md1KODXJxgGQPKngRmWyn10oCKFzNHOQ=
-github.com/rs/zerolog v1.23.0 h1:UskrK+saS9P9Y789yNNulYKdARjPZuS35B8gJF2x60g=
-github.com/rs/zerolog v1.23.0/go.mod h1:6c7hFfxPOy7TacJc4Fcdi24/J0NKYGzjG8FWRI916Qo=
-github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
-github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
-github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
-github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
-github.com/ulule/deepcopier v0.0.0-20200430083143-45decc6639b6 h1:TtyC78WMafNW8QFfv3TeP3yWNDG+uxNkk9vOrnDu6JA=
-github.com/ulule/deepcopier v0.0.0-20200430083143-45decc6639b6/go.mod h1:h8272+G2omSmi30fBXiZDMkmHuOgonplfKIKjQWzlfs=
-github.com/wweir/deferlog v0.1.0 h1:jzVo2Toct+0Ei29P3BpjVJpNgCLyyMZYhi+zr7L0bLw=
-github.com/wweir/deferlog v0.1.0/go.mod h1:cHSuGypQc2+yDurLohzOdJ2TpaQqGNm/6EhzskmBGZM=
-github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
+github.com/rs/xid v1.3.0/go.mod h1:trrq9SKmegXys3aeAKXMUTdJsYXVwGY3RLcfgqegfbg=
+github.com/rs/zerolog v1.26.1 h1:/ihwxqH+4z8UxyI70wM1z9yCvkWcfz/a3mj48k/Zngc=
+github.com/rs/zerolog v1.26.1/go.mod h1:/wSSJWX7lVrsOwlbyTRSOJvqRlc+WjWlfes+CiJ+tmc=
+github.com/sower-proxy/conns v0.0.2 h1:RRqygPkPaybSI8nrWGC6gc2HVK6Od/yVkyMU/7tz1Tc=
+github.com/sower-proxy/conns v0.0.2/go.mod h1:p3ZM/8MolzW4TN0LaiPJnwnhhif3y9sBi2vQEWHvwhg=
+github.com/sower-proxy/deferlog v1.0.3 h1:PdYrgadbwMesMhLURnuVCKsN7aZRHrVd6tcgYLyOIII=
+github.com/sower-proxy/deferlog v1.0.3/go.mod h1:i0SLfmrMoaLmZFXI8DEWrJbNpCd8alMn1GDHPZjWIsw=
+github.com/sower-proxy/mem v0.0.3-0.20220430030119-3c1d685997ce h1:vm/OnctGxqnpXO7vSoXhnl9qAhuOUFu4JJbI04u2c/Y=
+github.com/sower-proxy/mem v0.0.3-0.20220430030119-3c1d685997ce/go.mod h1:tM3P+gCAZMVBnOixGVAHS34si9Iwj2HU6M4IVEv4rK0=
+github.com/stretchr/testify v1.7.1 h1:5TQK59W5E3v0r2duFAb7P95B6hEeOyEnHRa8MjYSMTY=
+github.com/yuin/goldmark v1.3.5/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
+github.com/yuin/goldmark v1.4.0/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
-golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
-golang.org/x/crypto v0.0.0-20211108221036-ceb1ce70b4fa h1:idItI2DDfCokpg0N51B2VtiLdJ4vAuXC9fnCb2gACo4=
-golang.org/x/crypto v0.0.0-20211108221036-ceb1ce70b4fa/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
-golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
+golang.org/x/crypto v0.0.0-20211215165025-cf75a172585e/go.mod h1:P+XmwS30IXTQdn5tA2iutPOUgjI07+tq3H3K9MVA1s8=
+golang.org/x/crypto v0.0.0-20220427172511-eb4f295cb31f h1:OeJjE6G4dgCY4PIXvIRQbE8+RX+uXZyGhUy/ksMGJoc=
+golang.org/x/crypto v0.0.0-20220427172511-eb4f295cb31f/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4=
+golang.org/x/mod v0.4.2 h1:Gz96sIWK3OalVv/I/qNygP42zyoKp3xptRVCWRFEBvo=
+golang.org/x/mod v0.4.2/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
 golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
-golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
 golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
-golang.org/x/net v0.0.0-20211109214657-ef0fda0de508 h1:v3NKo+t/Kc3EASxaKZ82lwK6mCf4ZeObQBduYFZHo7c=
-golang.org/x/net v0.0.0-20211109214657-ef0fda0de508/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
+golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4/go.mod h1:p54w0d4576C0XHj96bSt6lcn1PtDYWL6XObtHCRCNQM=
+golang.org/x/net v0.0.0-20210726213435-c6fcb2dbf985/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
+golang.org/x/net v0.0.0-20210805182204-aaa1db679c0d/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
+golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2 h1:CIJ76btIcR3eFI5EgSo6k1qKw9KJexJuRLI9G7Hp5wE=
+golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
-golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20210220032951-036812b2e83c h1:5KslGYwFpkhGh+Q16bwMP3cOontH8FOep7tGV86Y7SQ=
 golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
-golang.org/x/sys v0.0.0-20190228124157-a34e9553db1e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20191224085550-c709ea063b76/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20210119212857-b64e53b001e4/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20210303074136-134d130e1a04/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20210330210617-4fbd30eecc44/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.0.0-20211109184856-51b60fd695b3 h1:T6tyxxvHMj2L1R2kZg0uNMpS8ZhB9lRa9XRGTCSA65w=
-golang.org/x/sys v0.0.0-20211109184856-51b60fd695b3/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20210510120138-977fb7262007/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20220422013727-9388b58f7150 h1:xHms4gcpe1YE7A3yIllJXP16CMAGuqwO2lX1mTyyRRc=
+golang.org/x/sys v0.0.0-20220422013727-9388b58f7150/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
+golang.org/x/text v0.3.6 h1:aRYxNxv6iGQlyVaZmk6ZgYEDa+Jg18DxebPSrd6bg1M=
 golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
-golang.org/x/text v0.3.7 h1:olpwvP2KacW1ZWvsR7uQhoyTYvKAupfQrRGBFM352Gk=
-golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
-golang.org/x/tools v0.1.0/go.mod h1:xkSsbof2nBLbhDlRMhhhyNLN/zl3eTqcnHD5viDpcZ0=
+golang.org/x/tools v0.1.6-0.20210726203631-07bc1bf47fb2/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
+golang.org/x/tools v0.1.7 h1:6j8CgantCy3yc8JGBqkDLMKWqZ0RDU2g1HVgacojGWQ=
+golang.org/x/tools v0.1.7/go.mod h1:LGqMHiF4EqQNHR1JncWGqT5BVaXmza+X+BDGol+dOxo=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 h1:go1bK/D/BFZV2I8cIQd1NKEZ+0owSTG1fDTci4IqFcE=
 golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v2 v2.3.0 h1:clyUAQHOM3G0M3f5vQj7LuJrETvjVot3Z5el9nffUtU=
 gopkg.in/yaml.v2 v2.3.0/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
-gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
-gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=
-gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

pkg/mem/mem.go

@@ -1,139 +0,0 @@
-package mem
-
-import (
-	"errors"
-	"reflect"
-	"sync"
-	"time"
-
-	"github.com/ulule/deepcopier"
-)
-
-// Data define the type which can speed up by mem cache
-type Data interface {
-	Fulfill(key string) (err error)
-}
-
-// Cache is the definition of cache, be careful of the memory usage
-type Cache struct {
-	old     *sync.Map
-	now     *sync.Map
-	barrier *sync.Map
-	rotate  <-chan time.Time
-	rwmutex *sync.RWMutex
-}
-
-// DefaultCache is default cache for surge
-var DefaultCache = New(time.Minute)
-
-// Remember is a surge, it provides a quite simple way to use cache
-func Remember(dst Data, key string) error {
-	return DefaultCache.Remember(dst, key)
-}
-
-// Delete is a surge, it delete a specified data in DefaultCache
-func Delete(dst Data, key string) {
-	DefaultCache.Delete(dst, key)
-}
-
-// New create a cache entity with a custom expiration time
-func New(rotateInterval time.Duration) *Cache {
-	return &Cache{
-		old:     &sync.Map{},
-		now:     &sync.Map{},
-		barrier: &sync.Map{},
-		rotate:  time.NewTicker(rotateInterval).C,
-		rwmutex: &sync.RWMutex{},
-	}
-}
-
-// Remember automatically save and retrieve data from a cache entity
-func (c *Cache) Remember(dst Data, key string) error {
-	rv := reflect.ValueOf(dst)
-	if rv.Kind() != reflect.Ptr {
-		panic("invalid not pointor type: " + reflect.TypeOf(dst).Name())
-	} else if rv.IsNil() {
-		return errors.New("invalid nil pointor")
-	}
-
-	c.rwmutex.RLock()
-	defer c.rwmutex.RUnlock()
-
-	// rotate logic, rwlock just protect fields in Cache, but not field content.
-	// So that, write lock just take a very short time, and simple read lock is
-	// just an atomic action, do not care the performance
-	select {
-	case <-c.rotate:
-		c.old = c.now
-		c.now = &sync.Map{}
-		c.barrier = &sync.Map{}
-	default:
-	}
-
-	// First: load from cache
-	if val, ok := c.now.Load(key); ok {
-		return deepcopier.Copy(val).To(dst)
-	}
-
-	// Second: load from old cache, or waitting the sigle groutine getting data
-	ch := make(chan struct{})
-	if chVal, ok := c.barrier.LoadOrStore(key, ch); ok {
-		close(ch) // the ch is not used
-
-		if val, ok := c.old.Load(key); ok {
-			return deepcopier.Copy(val).To(dst)
-		}
-
-		// type chan:  wait the sigle groutine getting data
-		// type error: already failed
-		if ch, ok = chVal.(chan struct{}); ok {
-			<-ch
-			if val, ok := c.now.Load(key); ok {
-				return deepcopier.Copy(val).To(dst)
-			}
-		}
-
-		val, _ := c.barrier.Load(key)
-		if err, ok := val.(error); ok {
-			return err
-		}
-
-		panic("new value lost, please report a bug")
-	}
-
-	// Third: getting data from CacheType, maybe from db
-	err := dst.Fulfill(key)
-	if err != nil {
-		c.barrier.Store(key, err)
-		return err
-	}
-
-	c.now.Store(key, dst)
-	close(ch) // broadcast, wakeup all waiting groutine
-
-	return nil
-}
-
-// Delete immediately specified the cached content to expire
-func (c *Cache) Delete(dst Data, key string) {
-	c.rwmutex.Lock()
-	defer c.rwmutex.Unlock()
-
-	c.old.Delete(key)
-	c.now.Delete(key)
-	c.barrier.Store(key, errors.New(key+" is deleted"))
-}
-
-// Rotate force refresh cached data
-func (c *Cache) Rotate(reset bool) {
-	c.rwmutex.Lock()
-	defer c.rwmutex.Unlock()
-
-	if reset {
-		c.old = &sync.Map{}
-	} else {
-		c.old = c.now
-	}
-	c.now = &sync.Map{}
-	c.barrier = &sync.Map{}
-}

pkg/suffixtree/suffix_tree.go

@@ -0,0 +1,143 @@
+package suffixtree
+
+import (
+	"strings"
+)
+
+type Node struct {
+	*node
+	sep   string
+	Count uint64
+}
+type node struct {
+	secs     []string
+	subNodes []*node
+}
+
+func NewNodeFromRules(rules ...string) *Node {
+	n := &Node{&node{}, ".", 0}
+	for i := range rules {
+		n.Add(rules[i])
+	}
+
+	n.GC()
+	return n
+}
+
+func (n *node) GC() {
+	if n == nil {
+		return
+	}
+
+	n.secs = GCSlice(n.secs)
+	n.subNodes = GCSlice(n.subNodes)
+	for i := range n.subNodes {
+		n.subNodes[i].GC()
+	}
+}
+
+func (n *Node) String() string {
+	return n.string("", "     ")
+}
+func (n *node) string(prefix, indent string) (out string) {
+	if n == nil {
+		return
+	}
+	for key, val := range n.subNodes {
+		out += prefix + n.secs[key] + "\n" + val.string(prefix+indent, indent)
+	}
+	return
+}
+
+func (n *Node) trim(item string) string {
+	return strings.TrimSuffix(item, n.sep)
+}
+
+func (n *Node) Add(item string) {
+	n.Count++
+	n.add(strings.Split(n.trim(item), n.sep))
+}
+func (n *node) add(secs []string) {
+	length := len(secs)
+	switch length {
+	case 0:
+	case 1:
+		sec := secs[length-1]
+		switch sec {
+		case "", "*", "**":
+			n.secs = append([]string{sec}, n.secs...)
+			n.subNodes = append([]*node{nil}, n.subNodes...)
+		default:
+			n.secs = append(n.secs, sec)
+			n.subNodes = append(n.subNodes, nil)
+		}
+	default:
+		sec := secs[length-1]
+		if sec == "**" { // ** is only allowed in the last sec
+			sec = "*"
+		}
+
+		idx := n.index(sec)
+		if idx == -1 {
+			switch sec {
+			case "", "*", "**":
+				idx = 0
+				n.secs = append([]string{sec}, n.secs...)
+				n.subNodes = append([]*node{{}}, n.subNodes...)
+			default:
+				idx = len(n.secs)
+				n.secs = append(n.secs, sec)
+				n.subNodes = append(n.subNodes, &node{})
+			}
+
+		} else if n.subNodes[idx] == nil {
+			n.subNodes[idx] = &node{}
+			n.subNodes[idx].add([]string{""})
+		}
+
+		n.subNodes[idx].add(secs[:length-1])
+	}
+}
+
+func (n *Node) Match(item string) bool {
+	if n == nil {
+		return false
+	}
+
+	return n.matchSecs(strings.Split(n.trim(item), n.sep), false)
+}
+
+func (n *node) matchSecs(secs []string, fuzzNode bool) bool {
+	length := len(secs)
+	if length == 0 {
+		if n == nil {
+			return true
+		}
+		return n.index("") != -1 || n.index("**") != -1
+	}
+
+	if idx := n.index(secs[length-1]); idx >= 0 {
+		if n.subNodes[idx].matchSecs(secs[:length-1], false) {
+			return true
+		}
+	}
+	if idx := n.index("*"); idx >= 0 {
+		if n.subNodes[idx].matchSecs(secs[:length-1], true) {
+			return true
+		}
+	}
+	return n.index("**") >= 0
+}
+
+// index return the sec index in node, or -1 if not found
+func (n *node) index(sec string) int {
+	if n == nil {
+		return -1
+	}
+	for s := range n.secs {
+		if n.secs[s] == sec {
+			return s
+		}
+	}
+	return -1
+}

pkg/suffixtree/suffix_tree_test.go

@@ -1,9 +1,9 @@
-package util_test
+package suffixtree_test
 
 import (
 	"testing"
 
-	"github.com/wweir/sower/util"
+	"github.com/wweir/sower/pkg/suffixtree"
 )
 
 func TestNode_Match(t *testing.T) {
@@ -13,18 +13,18 @@ func TestNode_Match(t *testing.T) {
 	}
 	tests := []struct {
 		name  string
-		node  *util.Node
+		node  *suffixtree.Node
 		tests []test
 	}{{
 		"simple",
-		util.NewNodeFromRules("a.wweir.cc", "b.wweir.cc"),
+		suffixtree.NewNodeFromRules("a.wweir.cc", "b.wweir.cc"),
 		[]test{
 			{"a.wweir.cc", true},
 			{"b.wweir.cc", true},
 		},
 	}, {
 		"parent",
-		util.NewNodeFromRules("wweir.cc", "a.wweir.cc"),
+		suffixtree.NewNodeFromRules("wweir.cc", "a.wweir.cc"),
 		[]test{
 			{"wweir.cc", true},
 			{"a.wweir.cc", true},
@@ -32,7 +32,7 @@ func TestNode_Match(t *testing.T) {
 		},
 	}, {
 		"fuzz1",
-		util.NewNodeFromRules("wweir.cc", "a.wweir.cc", "*.wweir.cc"),
+		suffixtree.NewNodeFromRules("wweir.cc", "a.wweir.cc", "*.wweir.cc"),
 		[]test{
 			{"wweir.cc", true},
 			{"a.wweir.cc", true},
@@ -41,7 +41,7 @@ func TestNode_Match(t *testing.T) {
 		},
 	}, {
 		"fuzz2",
-		util.NewNodeFromRules("a.*.cc", "c.wweir.*"),
+		suffixtree.NewNodeFromRules("a.*.cc", "c.wweir.*"),
 		[]test{
 			{"wweir.cc", false},
 			{"a.wweir.cc", true},
@@ -50,7 +50,7 @@ func TestNode_Match(t *testing.T) {
 		},
 	}, {
 		"fuzz3",
-		util.NewNodeFromRules("*.*.cc", "iamp.*.*"),
+		suffixtree.NewNodeFromRules("*.*.cc", "iamp.*.*"),
 		[]test{
 			{"wweir.cc", false},
 			{"a.wweir.cc", true},
@@ -59,7 +59,7 @@ func TestNode_Match(t *testing.T) {
 		},
 	}, {
 		"fuzz4",
-		util.NewNodeFromRules("**.cc", "a.**.com", "**.wweir.*"),
+		suffixtree.NewNodeFromRules("**.cc", "a.**.com", "**.wweir.*"),
 		[]test{
 			{"wweir.cc", true},
 			{"a.wweir.cc", true},
@@ -68,6 +68,12 @@ func TestNode_Match(t *testing.T) {
 			{"b.fuzz.com", false},
 			{"www.wweir.com", true},
 		},
+	}, {
+		"github",
+		suffixtree.NewNodeFromRules("**.github.com"),
+		[]test{
+			{"github.com", true},
+		},
 	}}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {

pkg/suffixtree/util.go

@@ -0,0 +1,6 @@
+package suffixtree
+
+func GCSlice[T any](arr []T) []T {
+	old := arr
+	return append(make([]T, 0, len(old)), old...)
+}

pkg/teeconn/tee_conn.go

@@ -1,56 +0,0 @@
-package teeconn
-
-import (
-	"io"
-	"net"
-)
-
-type Conn struct {
-	net.Conn
-
-	buf    []byte
-	offset int
-	stop   bool // read
-	err    error
-}
-
-func New(c net.Conn) *Conn {
-	return &Conn{Conn: c}
-}
-
-func (t *Conn) Reread() {
-	t.offset = 0
-}
-func (t *Conn) Reset() {
-	t.buf = []byte{}
-	t.offset = 0
-}
-func (t *Conn) Stop() *Conn {
-	t.stop = true
-	return t
-}
-
-func (t *Conn) Read(b []byte) (n int, err error) {
-	length := len(t.buf) - t.offset
-	if length > 0 {
-		n = copy(b, t.buf[t.offset:])
-		t.offset += n
-		return n, t.err
-	}
-
-	n, t.err = t.Conn.Read(b)
-	if !t.stop {
-		t.buf = append(t.buf, b[:n]...)
-		t.offset += n
-	}
-
-	return n, t.err
-}
-
-func (t *Conn) Write(b []byte) (n int, err error) {
-	if t.stop {
-		return t.Conn.Write(b)
-	}
-
-	return 0, io.ErrShortWrite
-}

router/country.go

@@ -5,7 +5,7 @@ import (
 	"net"
 	"time"
 
-	"github.com/wweir/deferlog/log"
+	"github.com/sower-proxy/deferlog/log"
 )
 
 func (r *Router) localSite(domain string) bool {

router/dns.go

@@ -2,10 +2,9 @@ package router
 
 import (
 	"net"
-	"time"
 
 	"github.com/miekg/dns"
-	"github.com/wweir/deferlog/log"
+	"github.com/sower-proxy/deferlog/log"
 )
 
 func (r *Router) ServeDNS(w dns.ResponseWriter, req *dns.Msg) {
@@ -18,27 +17,40 @@ func (r *Router) ServeDNS(w dns.ResponseWriter, req *dns.Msg) {
 	domain := req.Question[0].Name
 	// 1. rule_based( block > direct > proxy )
 	switch {
-	case r.blockRule.Match(domain):
+	case r.BlockRule.Match(domain):
 		_ = w.WriteMsg(r.dnsFail(req, dns.RcodeNameError))
 		return
 
-	case r.directRule.Match(domain):
+	case r.DirectRule.Match(domain):
 
-	case r.proxyRule.Match(domain):
+	case r.ProxyRule.Match(domain):
 		_ = w.WriteMsg(r.dnsProxyA(domain, r.dns.serveIP, req))
 		return
 	}
 
 	// 2. direct with cache, do not fallback to proxy to avoid side-effect
-	c := &dnsCache{Router: r, Req: req}
-	if err := r.dns.cache.Remember(c, req.Question[0].String()); err != nil {
-		_ = w.WriteMsg(r.dnsFail(req, dns.RcodeServerFailure))
-		return
+	conn := <-r.dns.connCh
+	resp, rtt, err := r.dns.ExchangeWithConn(req, conn)
+	if err != nil {
+		resp, rtt, err = r.dns.ExchangeWithConn(req, conn)
 	}
 
-	c.Resp.SetReply(req)
-	c.Resp.Compress = true
-	_ = w.WriteMsg(c.Resp)
+	log.DebugWarn(err).
+		Dur("rtt", rtt).
+		Str("question", req.Question[0].String()).
+		Msg("exchange dns record")
+	if err != nil {
+		_ = w.WriteMsg(r.dnsFail(req, dns.RcodeServerFailure))
+		conn.Close()
+
+	} else {
+		_ = w.WriteMsg(resp)
+		select {
+		case r.dns.connCh <- conn:
+		default:
+			conn.Close()
+		}
+	}
 }
 
 func (r *Router) dnsFail(req *dns.Msg, rcode int) *dns.Msg {
@@ -65,29 +77,3 @@ func (r *Router) dnsProxyA(domain string, localIP net.IP, req *dns.Msg) *dns.Msg
 	}
 	return m
 }
-
-type dnsCache struct {
-	*Router
-	Req, Resp *dns.Msg
-}
-
-func (r *dnsCache) Fulfill(question string) (err error) {
-	conn := <-r.dns.connCh
-
-	var rtt time.Duration
-	r.Resp, rtt, err = r.dns.ExchangeWithConn(r.Req, conn)
-	if err != nil {
-		r.Resp, rtt, err = r.dns.ExchangeWithConn(r.Req, conn)
-	}
-	log.DebugWarn(err).
-		Dur("rtt", rtt).
-		Str("question", question).
-		Msg("exchange dns record")
-
-	select {
-	case r.dns.connCh <- conn:
-	default:
-		conn.Close()
-	}
-	return err
-}

router/ping.go

@@ -5,12 +5,11 @@ import (
 	"sync"
 	"time"
 
-	"github.com/wweir/deferlog/log"
+	"github.com/sower-proxy/deferlog/log"
+	"github.com/sower-proxy/mem"
 )
 
-var pingClient = http.Client{
-	Timeout: 2 * time.Second,
-}
+var accessCache = mem.NewCache(time.Hour, httpPing)
 
 func (r *Router) isAccess(domain string, port uint16) bool {
 	switch port {
@@ -20,16 +19,15 @@ func (r *Router) isAccess(domain string, port uint16) bool {
 		return false
 	}
 
-	p := &ping{}
-	_ = r.accessCache.Remember(p, domain)
-	return p.isAccess
+	ok, _ := accessCache.Get(domain)
+	return ok
 }
 
-type ping struct {
-	isAccess bool
+var pingClient = http.Client{
+	Timeout: 2 * time.Second,
 }
 
-func (p *ping) Fulfill(key string) error {
+func httpPing(key string) (bool, error) {
 	wg := sync.WaitGroup{}
 	wg.Add(2)
 	var err80, err443 error
@@ -49,6 +47,5 @@ func (p *ping) Fulfill(key string) error {
 			Msg("Failed to ping")
 	}
 
-	p.isAccess = (err80 == nil && err443 == nil)
-	return nil
+	return err80 == nil && err443 == nil, nil
 }

router/router.go

@@ -8,27 +8,25 @@ import (
 	"github.com/miekg/dns"
 	geoip2 "github.com/oschwald/geoip2-golang"
 	"github.com/pkg/errors"
-	"github.com/wweir/deferlog"
-	"github.com/wweir/deferlog/log"
+	"github.com/sower-proxy/conns/relay"
+	"github.com/sower-proxy/deferlog"
+	"github.com/sower-proxy/deferlog/log"
 	"github.com/wweir/sower/pkg/dhcp"
-	"github.com/wweir/sower/pkg/mem"
-	"github.com/wweir/sower/util"
+	"github.com/wweir/sower/pkg/suffixtree"
 )
 
 type ProxyDialFn func(network, host string, port uint16) (net.Conn, error)
 type Router struct {
-	blockRule   *util.Node
-	directRule  *util.Node
-	proxyRule   *util.Node
-	ProxyDial   ProxyDialFn
-	accessCache *mem.Cache
+	BlockRule  *suffixtree.Node
+	DirectRule *suffixtree.Node
+	ProxyRule  *suffixtree.Node
+	ProxyDial  ProxyDialFn
 
 	dns struct {
 		dns.Client
 		fallbackDNS string
 		serveIP     net.IP
 		connCh      chan *dns.Conn
-		cache       *mem.Cache
 	}
 
 	country struct {
@@ -39,14 +37,12 @@ type Router struct {
 
 func NewRouter(serveIP, fallbackDNS, mmdbFile string, proxyDial ProxyDialFn) *Router {
 	r := Router{
-		ProxyDial:   proxyDial,
-		accessCache: mem.New(time.Hour), // TODO: config
+		ProxyDial: proxyDial,
 	}
 
 	r.dns.serveIP = net.ParseIP(serveIP)
 	r.dns.fallbackDNS = fallbackDNS
 	r.dns.connCh = make(chan *dns.Conn, 1)
-	r.dns.cache = mem.New(5 * time.Minute) // Tll: 10 minutes
 	go r.dialDNSConn()
 
 	var err error
@@ -56,20 +52,15 @@ func NewRouter(serveIP, fallbackDNS, mmdbFile string, proxyDial ProxyDialFn) *Ro
 	return &r
 }
 
-func (r *Router) SetRules(blockList, directList, proxyList, directCIDRs []string) {
-
-	r.blockRule = util.NewNodeFromRules(blockList...)
-	r.directRule = util.NewNodeFromRules(directList...)
-	r.proxyRule = util.NewNodeFromRules(proxyList...)
-
-	r.country.cidrs = make([]*net.IPNet, 0, len(directCIDRs))
-	for _, cidr := range directCIDRs {
+func (r *Router) AddCountryCIDRs(cidrs ...string) {
+	for _, cidr := range cidrs {
 		_, ipnet, err := net.ParseCIDR(cidr)
 		if err != nil {
 			log.Error().Err(err).Msg("Failed to parse CIDR")
 		}
 		r.country.cidrs = append(r.country.cidrs, ipnet)
 	}
+	r.country.cidrs = suffixtree.GCSlice(r.country.cidrs)
 }
 
 func (r *Router) dialDNSConn() {
@@ -79,7 +70,7 @@ func (r *Router) dialDNSConn() {
 			Str("DNS", server).
 			Str("fallback", r.dns.fallbackDNS).
 			Msg("get DNS server")
-		if server == "" {
+		if server == "" || server == r.dns.serveIP.String() {
 			server = r.dns.fallbackDNS
 		}
 
@@ -111,13 +102,13 @@ func (r *Router) RouteHandle(conn net.Conn, domain string, port uint16) (err err
 	// 2. detect_based( CN IP || access site )
 	// 3. fallback( proxy )
 	switch {
-	case r.blockRule.Match(domain):
+	case r.BlockRule.Match(domain):
 		return nil
 
-	case r.directRule.Match(domain):
+	case r.DirectRule.Match(domain):
 		return r.DirectHandle(conn, addr)
 
-	case r.proxyRule.Match(domain):
+	case r.ProxyRule.Match(domain):
 		return r.ProxyHandle(conn, domain, port)
 
 	case r.localSite(domain), r.isAccess(domain, port):
@@ -135,11 +126,11 @@ func (r *Router) ProxyHandle(conn net.Conn, domain string, port uint16) error {
 	}
 	defer rc.Close()
 
-	util.Relay(conn, rc)
+	relay.Relay(conn, rc)
 	return nil
 }
 
 func (r *Router) DirectHandle(conn net.Conn, addr string) error {
-	dur, err := util.RelayTo(conn, addr)
+	dur, err := relay.RelayTo(conn, addr)
 	return errors.Wrapf(err, "spend (%s)", dur)
 }

transport/transport_test.go

@@ -5,8 +5,8 @@ import (
 	"strings"
 	"testing"
 
-	"github.com/wweir/deferlog/log"
-	"github.com/wweir/sower/pkg/teeconn"
+	"github.com/sower-proxy/conns/teeconn"
+	"github.com/sower-proxy/deferlog/log"
 	"github.com/wweir/sower/transport/sower"
 	"github.com/wweir/sower/transport/trojan"
 )

util/relay.go

@@ -1,50 +0,0 @@
-package util
-
-import (
-	"io"
-	"net"
-	"sync"
-	"sync/atomic"
-	"time"
-
-	"github.com/pkg/errors"
-)
-
-func RelayTo(conn net.Conn, addr string) (dur time.Duration, err error) {
-	if _, _, err := net.SplitHostPort(addr); err != nil {
-		addr = net.JoinHostPort(addr, "80")
-	}
-
-	start := time.Now()
-	rc, err := net.DialTimeout("tcp", addr, 5*time.Second)
-	if err != nil {
-		return time.Since(start), errors.WithStack(err)
-	}
-	defer rc.Close()
-
-	Relay(conn, rc)
-	return time.Since(start), nil
-}
-
-func Relay(conn1, conn2 net.Conn) {
-	wg := &sync.WaitGroup{}
-	exitFlag := new(int32)
-	wg.Add(2)
-	go redirect(conn2, conn1, wg, exitFlag)
-	redirect(conn1, conn2, wg, exitFlag)
-	wg.Wait()
-}
-func redirect(dst, src net.Conn, wg *sync.WaitGroup, exitFlag *int32) {
-
-	// io.Copy(dst, io.TeeReader(src, os.Stdout))
-	io.Copy(dst, src)
-
-	if atomic.CompareAndSwapInt32(exitFlag, 0, 1) {
-		// wakeup blocked goroutine
-		now := time.Now()
-		src.SetDeadline(now)
-		dst.SetDeadline(now)
-	}
-
-	wg.Done()
-}

util/suffix_tree.go

@@ -1,108 +0,0 @@
-package util
-
-import (
-	"strings"
-	"sync"
-)
-
-type Node struct {
-	node
-	sep string
-	*sync.RWMutex
-}
-type node struct {
-	node map[string]*node
-}
-
-func NewNodeFromRules(rules ...string) *Node {
-	n := &Node{node{node: map[string]*node{}}, ".", &sync.RWMutex{}}
-	for i := range rules {
-		n.Add(rules[i])
-	}
-	return n
-}
-
-func (n *Node) String() string {
-	n.RLock()
-	defer n.RUnlock()
-	return n.string("", "     ")
-}
-func (n *node) string(prefix, indent string) (out string) {
-	for key, val := range n.node {
-		out += prefix + key + "\n" + val.string(prefix+indent, indent)
-	}
-	return
-}
-
-func (n *Node) trim(item string) string {
-	return strings.TrimSuffix(item, n.sep)
-}
-
-func (n *Node) Add(item string) {
-	n.Lock()
-	defer n.Unlock()
-	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]] = &node{node: map[string]*node{"": {}}}
-	default:
-		sec := secs[length-1]
-		if sec == "**" {
-			sec = "*"
-		}
-
-		subNode, ok := n.node[sec]
-		if !ok {
-			subNode = &node{node: map[string]*node{}}
-			n.node[sec] = subNode
-		}
-		subNode.add(secs[:length-1])
-	}
-}
-
-func (n *Node) Match(item string) bool {
-	if n == nil {
-		return false
-	}
-
-	n.RLock()
-	defer n.RUnlock()
-	return n.matchSecs(strings.Split(n.trim(item), n.sep), false)
-}
-
-func (n *node) matchSecs(secs []string, fuzzNode bool) bool {
-	length := len(secs)
-	if length == 0 {
-		if _, ok := n.node[""]; ok {
-			return true
-		}
-		if _, ok := n.node["**"]; ok {
-			return true
-		}
-		if _, ok := n.node["*"]; ok {
-			return !fuzzNode
-		}
-		return false
-	}
-
-	if n, ok := n.node[secs[length-1]]; ok {
-		if n.matchSecs(secs[:length-1], false) {
-			return true
-		}
-	}
-	if n, ok := n.node["*"]; ok {
-		if n.matchSecs(secs[:length-1], true) {
-			return true
-		}
-	}
-	if _, ok := n.node["**"]; ok {
-		return true
-	}
-
-	return false
-}