feat: url & keyword filter

refactor & fix

README.md

@@ -112,10 +112,8 @@ See [benchmark](benchmark/readme.md)
 - [x] TTL of URL
 - [x] rate limit
 - ~~[ ] management(auth)~~
-  - [x] remove(domain)
-  - [ ] domain blacklist
-  - [ ] domain whitelist
-  - [ ] keyword filter
+  - [x] domain filter
+  - [x] keyword filter
 - [ ] statistic
   - [ ] URL status
   - [ ] runtime metrics

config.yml

@@ -1,11 +1,22 @@
 debug: false
-shortener:
-  store:
-    dbpath: ./yasuser.db
-    dbtype: bolt
-    redis: redis://localhost:6379
+store:
+  dbpath: ./yasuser.db
+  dbtype: bolt
+  redis: redis://localhost:6379
 server:
   domain: https://u.kfd.me
   port: 8084
+  limit: 10
   pprof: false
-  gaid: 62244864-8
+  gaid: 62244864-8
+filter:
+  domain:
+    whitelist:
+      - kfd.me
+    blacklist:
+      - t66y.com
+  keyword:
+    whitelist:
+      - yasuser
+    blacklist:
+      - gg

config/config.go

@@ -26,15 +26,22 @@ type StoreConfig struct {
 	Redis  string `default:"redis://localhost:6379"`
 }
 
-type ShortenerConfig struct {
-	Store StoreConfig
+type list struct {
+	WhiteList []string
+	BlackList []string
+}
+
+type Filter struct {
+	Domain  list
+	Keyword list
 }
 
 type Config struct {
-	Debug     bool `default:"false"`
-	Shortener ShortenerConfig
-	Server    SrvConfig
-	Auth      string `default:"passwd"`
+	Debug  bool   `default:"false"`
+	Auth   string `default:"password"`
+	Store  StoreConfig
+	Server SrvConfig
+	Filter Filter
 }
 
 func New() *Config {

filter/errors.go

@@ -4,7 +4,6 @@ import "errors"
 
 // errors
 var (
-	ErrInBlackList = errors.New("domain in blacklist")
-	ErrBadKeyword  = errors.New("url containes bad keyword")
-	ErrRemoved     = errors.New("domain removed")
+	ErrBadDomain  = errors.New("domain in blacklist")
+	ErrBadKeyword = errors.New("url contains bad keyword")
 )

filter/filter.go

@@ -1,8 +1,71 @@
 package filter
 
+import (
+	"net/url"
+	"strings"
+
+	"github.com/wrfly/yasuser/config"
+)
+
 type Filter interface {
-	Removed(domain string) bool
-	InBlackList(domain string) bool
-	InWhiteList(domain string) bool
-	BadKeyword(url string) bool
+	OK(*url.URL) error
+}
+
+type list struct {
+	blacklist map[string]bool
+	whitelist map[string]bool
+}
+
+type urlFilter struct {
+	domain  list
+	keyword list
+}
+
+func (f *urlFilter) OK(u *url.URL) error {
+	// bypass domain
+	if f.domain.whitelist[u.Hostname()] {
+		return nil
+	}
+
+	// bypass keyword
+	for x := range f.keyword.whitelist {
+		if strings.Contains(u.Path, x) {
+			return nil
+		}
+	}
+
+	// bad domain
+	if f.domain.blacklist[u.Hostname()] {
+		return ErrBadDomain
+	}
+
+	// bad keyword
+	for x := range f.keyword.blacklist {
+		if strings.Contains(u.Path, x) {
+			return ErrBadKeyword
+		}
+	}
+
+	return nil
+}
+
+func makeList(slice []string) map[string]bool {
+	x := make(map[string]bool, len(slice))
+	for _, s := range slice {
+		x[s] = true
+	}
+	return x
+}
+
+func New(conf config.Filter) Filter {
+	return &urlFilter{
+		domain: list{
+			whitelist: makeList(conf.Domain.WhiteList),
+			blacklist: makeList(conf.Domain.BlackList),
+		},
+		keyword: list{
+			whitelist: makeList(conf.Keyword.WhiteList),
+			blacklist: makeList(conf.Keyword.BlackList),
+		},
+	}
 }

main.go

@@ -10,6 +10,7 @@ import (
 	"gopkg.in/urfave/cli.v2"
 
 	"github.com/wrfly/yasuser/config"
+	"github.com/wrfly/yasuser/filter"
 	"github.com/wrfly/yasuser/routes"
 	"github.com/wrfly/yasuser/shortener"
 )
@@ -67,7 +68,11 @@ func main() {
 				gin.SetMode(gin.ReleaseMode)
 			}
 
-			err := routes.Serve(conf.Server, shortener.New(conf.Shortener))
+			err := routes.Serve(
+				conf.Server,
+				shortener.New(conf.Store),
+				filter.New(conf.Filter),
+			)
 			if err != nil {
 				logrus.Error(err)
 			}

routes/asset/asset.go

@@ -1,6 +1,6 @@
 /*
 CODE GENERATED BY "github.com/wrfly/bindata" 
-@2019-01-13T23:46:56+08:00
+@2019-01-22T00:14:59+08:00
 
 Files:
 	/
@@ -28,7 +28,7 @@ var _file_0 = &file{
 		isDir: true,
 		size:  4096,
 		mode:  os.FileMode(2147484157),
-		mTime: time.Unix(1547393379, 0),
+		mTime: time.Unix(1548085371, 0),
 		cType: "",
 	},
 	path:  "/",
@@ -478,7 +478,7 @@ var _file_3 = &file{
 		isDir: false,
 		size:  1817,
 		mode:  os.FileMode(436),
-		mTime: time.Unix(1532867014, 0),
+		mTime: time.Unix(1548085371, 0),
 		cType: "text/html; charset=utf-8",
 	},
 	path:  "/index.html",
@@ -520,7 +520,7 @@ var _file_4 = &file{
 		isDir: false,
 		size:  857,
 		mode:  os.FileMode(436),
-		mTime: time.Unix(1532868516, 0),
+		mTime: time.Unix(1548085371, 0),
 		cType: "text/css; charset=utf-8",
 	},
 	path:  "/main.css",

routes/handler.go

@@ -15,7 +15,7 @@ import (
 	"github.com/wrfly/yasuser/config"
 	"github.com/wrfly/yasuser/filter"
 	"github.com/wrfly/yasuser/routes/asset"
-	stner "github.com/wrfly/yasuser/shortener"
+	s "github.com/wrfly/yasuser/shortener"
 	"github.com/wrfly/yasuser/types"
 )
 
@@ -24,35 +24,41 @@ const (
 	maxCustomLength = 60
 )
 
-var validCustomURI *regexp.Regexp
+var (
+	validCustomURI *regexp.Regexp
+	indexTemplate  *template.Template
+)
 
 func init() {
 	validURI, err := regexp.Compile("^[a-zA-Z0-9][a-zA-Z0-9_+-]+$")
 	if err != nil {
 		panic(err)
 	}
 	validCustomURI = validURI
+
+	a, err := asset.Data.Asset("/index.html")
+	if err != nil {
+		panic(err)
+	}
+	indexTemplate = a.Template()
+
 }
 
 type server struct {
 	domain string
 	gaID   string
 	limit  int64
 
-	stener        stner.Shortener
-	indexTemplate *template.Template
-	fileMap       map[string]bool
-	filter        filter.Filter
+	stener  s.Shortener
+	fileMap map[string]bool
+	filter  filter.Filter
 
 	host string
 	tb   map[string]tokenbucket.Bucket
 }
 
-func newServer(conf config.SrvConfig, shortener stner.Shortener) server {
-	a, err := asset.Data.Asset("/index.html")
-	if err != nil {
-		panic(err)
-	}
+func newServer(conf config.SrvConfig,
+	shortener s.Shortener, filter filter.Filter) server {
 
 	u, err := url.Parse(conf.Domain)
 	if err != nil {
@@ -67,9 +73,7 @@ func newServer(conf config.SrvConfig, shortener stner.Shortener) server {
 		limit:   conf.Limit,
 		fileMap: make(map[string]bool),
 		tb:      make(map[string]tokenbucket.Bucket, 0),
-		// TODO: create a filter
-		filter:        nil,
-		indexTemplate: a.Template(),
+		filter:  filter,
 	}
 	for _, a := range asset.Data.List() {
 		srv.fileMap[a.Name()] = true
@@ -88,7 +92,7 @@ func (s *server) handleIndex() gin.HandlerFunc {
 				s.domain, "http://longlonglong.com/long/long/long?a=1&b=2"))
 		} else {
 			// visit from a web browser
-			s.indexTemplate.Execute(c.Writer, map[string]string{
+			indexTemplate.Execute(c.Writer, map[string]string{
 				"domain": s.domain,
 				"gaID":   s.gaID,
 			})
@@ -226,31 +230,14 @@ func (s *server) invalidURL(URL string) error {
 		return err
 	}
 
-	domain := u.Hostname()
-	if s.filter.InWhiteList(domain) {
-		return nil // bypass
-	}
-
-	if domain == s.host {
+	if u.Hostname() == s.host {
 		return types.ErrSameHost
 	}
-
-	if s.filter.InBlackList(domain) {
-		return filter.ErrInBlackList
-	}
-
-	if s.filter.BadKeyword(u.Path) {
-		return filter.ErrBadKeyword
-	}
-
-	if s.filter.Removed(domain) {
-		return filter.ErrRemoved
-	}
-
 	switch u.Scheme {
 	case "http", "https", "ftp", "tcp":
-		return nil
 	default:
 		return types.ErrScheme
 	}
+
+	return s.filter.OK(u)
 }

routes/serve.go

@@ -14,7 +14,8 @@ import (
 	"github.com/sirupsen/logrus"
 
 	"github.com/wrfly/yasuser/config"
-	stner "github.com/wrfly/yasuser/shortener"
+	"github.com/wrfly/yasuser/filter"
+	s "github.com/wrfly/yasuser/shortener"
 )
 
 const MAX_URL_LENGTH = 1e3
@@ -26,11 +27,12 @@ var urlBufferPool = sync.Pool{
 }
 
 // Serve routes
-func Serve(conf config.SrvConfig, shortener stner.Shortener) error {
+func Serve(conf config.SrvConfig,
+	shortener s.Shortener, filter filter.Filter) error {
 	sigChan := make(chan os.Signal)
 	signal.Notify(sigChan, os.Interrupt, os.Kill)
 
-	srv := newServer(conf, shortener)
+	srv := newServer(conf, shortener, filter)
 
 	e := gin.New()