Initial commit

library/net/http/blademaster/middleware/antispam/BUILD

@@ -0,0 +1,64 @@
+package(default_visibility = ["//visibility:public"])
+
+load(
+    "@io_bazel_rules_go//go:def.bzl",
+    "go_test",
+    "go_library",
+)
+
+go_test(
+    name = "go_default_test",
+    srcs = ["antispam_test.go"],
+    embed = [":go_default_library"],
+    rundir = ".",
+    tags = ["automanaged"],
+    deps = [
+        "//library/cache/redis:go_default_library",
+        "//library/container/pool:go_default_library",
+        "//library/net/http/blademaster:go_default_library",
+        "//library/time:go_default_library",
+        "//vendor/github.com/stretchr/testify/assert:go_default_library",
+    ],
+)
+
+go_library(
+    name = "go_default_library",
+    srcs = ["antispam.go"],
+    importpath = "go-common/library/net/http/blademaster/middleware/antispam",
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//library/cache/redis:go_default_library",
+        "//library/ecode:go_default_library",
+        "//library/log:go_default_library",
+        "//library/net/http/blademaster:go_default_library",
+        "//vendor/github.com/pkg/errors:go_default_library",
+    ],
+)
+
+go_test(
+    name = "go_default_xtest",
+    srcs = ["example_test.go"],
+    tags = ["automanaged"],
+    deps = [
+        "//library/cache/redis:go_default_library",
+        "//library/container/pool:go_default_library",
+        "//library/net/http/blademaster:go_default_library",
+        "//library/net/http/blademaster/middleware/antispam:go_default_library",
+        "//library/time:go_default_library",
+    ],
+)
+
+filegroup(
+    name = "package-srcs",
+    srcs = glob(["**"]),
+    tags = ["automanaged"],
+    visibility = ["//visibility:private"],
+)
+
+filegroup(
+    name = "all-srcs",
+    srcs = [":package-srcs"],
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+)

library/net/http/blademaster/middleware/antispam/CHANGELOG.md

@@ -0,0 +1,5 @@
+### business/blademaster/antispam
+
+##### Version 1.0.0
+
+1. 完成基本功能与测试

library/net/http/blademaster/middleware/antispam/CONTRIBUTORS.md

@@ -0,0 +1,6 @@
+# Author
+lintnaghui
+caoguoliang
+ 
+# Reviewer
+maojian

library/net/http/blademaster/middleware/antispam/OWNERS

@@ -0,0 +1,9 @@
+# See the OWNERS docs at https://go.k8s.io/owners
+
+approvers:
+- caoguoliang
+- lintnaghui
+reviewers:
+- caoguoliang
+- lintnaghui
+- maojian

library/net/http/blademaster/middleware/antispam/README.md

@@ -0,0 +1,13 @@
+#### business/blademaster/antispam
+
+##### 项目简介
+
+blademaster 的 antispam middleware，主要用于限制用户的请求频率
+
+##### 编译环境
+
+- **请只用 Golang v1.8.x 以上版本编译执行**
+
+##### 依赖包
+
+- No other dependency

library/net/http/blademaster/middleware/antispam/antispam.go

@@ -0,0 +1,139 @@
+package antispam
+
+import (
+	"fmt"
+	"time"
+
+	"go-common/library/cache/redis"
+	"go-common/library/ecode"
+	"go-common/library/log"
+	bm "go-common/library/net/http/blademaster"
+
+	"github.com/pkg/errors"
+)
+
+const (
+	_prefixRate  = "r_%d_%s_%d"
+	_prefixTotal = "t_%d_%s_%d"
+	// antispam
+	_defSecond = 1
+	_defHour   = 1
+)
+
+// Antispam is a antispam instance.
+type Antispam struct {
+	redis *redis.Pool
+	conf  *Config
+}
+
+// Config antispam config.
+type Config struct {
+	On     bool // switch on/off
+	Second int  // every N second allow N requests.
+	N      int  // one unit allow N requests.
+	Hour   int  // every N hour allow M requests.
+	M      int  // one winodw allow M requests.
+
+	Redis *redis.Config
+}
+
+func (c *Config) validate() error {
+	if c == nil {
+		return errors.New("antispam: empty config")
+	}
+	if c.Second < _defSecond {
+		return errors.New("antispam: invalid Second")
+	}
+	if c.Hour < _defHour {
+		return errors.New("antispam: invalid Hour")
+	}
+	return nil
+}
+
+// New new a antispam service.
+func New(c *Config) (s *Antispam) {
+	if err := c.validate(); err != nil {
+		panic(err)
+	}
+	s = &Antispam{
+		redis: redis.NewPool(c.Redis),
+	}
+	s.Reload(c)
+	return s
+}
+
+// Reload reload antispam config.
+func (s *Antispam) Reload(c *Config) {
+	if err := c.validate(); err != nil {
+		log.Error("Failed to reload antispam: %+v", err)
+		return
+	}
+	s.conf = c
+}
+
+// Rate antispam by user + path.
+func (s *Antispam) Rate(c *bm.Context, second, count int) (err error) {
+	mid, ok := c.Get("mid")
+	if !ok {
+		return
+	}
+	curSecond := int(time.Now().Unix())
+	burst := curSecond - curSecond%second
+	key := rateKey(mid.(int64), c.Request.URL.Path, burst)
+	return s.antispam(c, key, second, count)
+}
+
+// Total antispam by user + path.
+func (s *Antispam) Total(c *bm.Context, hour, count int) (err error) {
+	second := hour * 3600
+	mid, ok := c.Get("mid")
+	if !ok {
+		return
+	}
+	curHour := int(time.Now().Unix() / 3600)
+	burst := curHour - curHour%hour
+	key := totalKey(mid.(int64), c.Request.URL.Path, burst)
+	return s.antispam(c, key, second, count)
+}
+
+func (s *Antispam) antispam(c *bm.Context, key string, interval, count int) error {
+	conn := s.redis.Get(c)
+	defer conn.Close()
+	incred, err := redis.Int64(conn.Do("INCR", key))
+	if err != nil {
+		return nil
+	}
+	if incred == 1 {
+		conn.Do("EXPIRE", key, interval)
+	}
+	if incred > int64(count) {
+		return ecode.LimitExceed
+	}
+	return nil
+}
+
+func rateKey(mid int64, path string, burst int) string {
+	return fmt.Sprintf(_prefixRate, mid, path, burst)
+}
+
+func totalKey(mid int64, path string, burst int) string {
+	return fmt.Sprintf(_prefixTotal, mid, path, burst)
+}
+
+func (s *Antispam) ServeHTTP(ctx *bm.Context) {
+	if err := s.Rate(ctx, s.conf.Second, s.conf.N); err != nil {
+		ctx.JSON(nil, ecode.LimitExceed)
+		ctx.Abort()
+		return
+	}
+	if err := s.Total(ctx, s.conf.Hour, s.conf.M); err != nil {
+		ctx.JSON(nil, ecode.LimitExceed)
+		ctx.Abort()
+		return
+	}
+}
+
+// Handler is antispam handle.
+func (s *Antispam) Handler() bm.HandlerFunc {
+	return s.ServeHTTP
+}

library/net/http/blademaster/middleware/antispam/antispam_test.go

@@ -0,0 +1,108 @@
+package antispam
+
+import (
+	"context"
+	"io/ioutil"
+	"net/http"
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+
+	"go-common/library/cache/redis"
+	"go-common/library/container/pool"
+	bm "go-common/library/net/http/blademaster"
+	xtime "go-common/library/time"
+)
+
+func TestAntiSpamHandler(t *testing.T) {
+	anti := New(
+		&Config{
+			On:     true,
+			Second: 1,
+			N:      1,
+			Hour:   1,
+			M:      1,
+			Redis: &redis.Config{
+				Config: &pool.Config{
+					Active:      10,
+					Idle:        10,
+					IdleTimeout: xtime.Duration(time.Second * 60),
+				},
+				Name:         "test",
+				Proto:        "tcp",
+				Addr:         "172.18.33.60:6889",
+				DialTimeout:  xtime.Duration(time.Second),
+				ReadTimeout:  xtime.Duration(time.Second),
+				WriteTimeout: xtime.Duration(time.Second),
+			},
+		},
+	)
+
+	engine := bm.New()
+	engine.UseFunc(func(c *bm.Context) {
+		mid, _ := strconv.ParseInt(c.Request.Form.Get("mid"), 10, 64)
+		c.Set("mid", mid)
+		c.Next()
+	})
+	engine.Use(anti.Handler())
+	engine.GET("/antispam", func(c *bm.Context) {
+		c.String(200, "pass")
+	})
+	go engine.Run(":18080")
+
+	time.Sleep(time.Millisecond * 50)
+	code, content, err := httpGet("http://127.0.0.1:18080/antispam?mid=11")
+	if err != nil {
+		t.Logf("http get failed, err:=%v", err)
+		t.FailNow()
+	}
+	if code != 200 || string(content) != "pass" {
+		t.Logf("request should pass by limiter, but blocked: %d, %v", code, content)
+		t.FailNow()
+	}
+
+	_, content, err = httpGet("http://127.0.0.1:18080/antispam?mid=11")
+	if err != nil {
+		t.Logf("http get failed, err:=%v", err)
+		t.FailNow()
+	}
+	if string(content) == "pass" {
+		t.Logf("request should block by limiter, but passed")
+		t.FailNow()
+	}
+
+	engine.Server().Shutdown(context.TODO())
+}
+
+func httpGet(url string) (code int, content []byte, err error) {
+	resp, err := http.Get(url)
+	if err != nil {
+		return
+	}
+	defer resp.Body.Close()
+
+	content, err = ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return
+	}
+	code = resp.StatusCode
+	return
+}
+
+func TestConfigValidate(t *testing.T) {
+	var conf *Config
+	assert.Contains(t, conf.validate().Error(), "empty config")
+
+	conf = &Config{
+		Second: 0,
+	}
+	assert.Contains(t, conf.validate().Error(), "invalid Second")
+
+	conf = &Config{
+		Second: 1,
+		Hour:   0,
+	}
+	assert.Contains(t, conf.validate().Error(), "invalid Hour")
+}

library/net/http/blademaster/middleware/antispam/example_test.go

@@ -0,0 +1,45 @@
+package antispam_test
+
+import (
+	"time"
+
+	"go-common/library/cache/redis"
+	"go-common/library/container/pool"
+	"go-common/library/net/http/blademaster"
+	"go-common/library/net/http/blademaster/middleware/antispam"
+	xtime "go-common/library/time"
+)
+
+// This example create a antispam middleware instance and attach to a blademaster engine,
+// it will protect '/ping' API with specified policy.
+// If anyone who requests this API more frequently than 1 req/second or 1 req/hour,
+// a StatusServiceUnavailable error will be raised.
+func Example() {
+	anti := antispam.New(&antispam.Config{
+		On:     true,
+		Second: 1,
+		N:      1,
+		Hour:   1,
+		M:      1,
+		Redis: &redis.Config{
+			Config: &pool.Config{
+				Active:      10,
+				Idle:        10,
+				IdleTimeout: xtime.Duration(time.Second * 60),
+			},
+			Name:         "test",
+			Proto:        "tcp",
+			Addr:         "172.18.33.60:6889",
+			DialTimeout:  xtime.Duration(time.Second),
+			ReadTimeout:  xtime.Duration(time.Second),
+			WriteTimeout: xtime.Duration(time.Second),
+		},
+	})
+
+	engine := blademaster.Default()
+	engine.Use(anti)
+	engine.GET("/ping", func(c *blademaster.Context) {
+		c.String(200, "%s", "pong")
+	})
+	engine.Run(":18080")
+}