add: 基于redis的api调用频率限制middleware

server/artisan/redis.go

+package artisan
+
+import (
+	"encoding/json"
+	"github.com/go-redis/cache"
+	"github.com/peterq/pan-light/server/conf"
+	"time"
+)
+
+var codec = &cache.Codec{
+	Redis: conf.Redis,
+
+	Marshal: func(v interface{}) ([]byte, error) {
+		return json.Marshal(v)
+	},
+	Unmarshal: func(b []byte, v interface{}) error {
+		return json.Unmarshal(b, v)
+	},
+}
+
+func RedisGet(key string, data interface{}) error {
+	return codec.Get(key, data)
+}
+
+func RedisSet(key string, value interface{}, expiration time.Duration) error {
+	return codec.Set(&cache.Item{
+		Key:        key,
+		Object:     value,
+		Expiration: expiration,
+	})
+}

server/artisan/web.go

@@ -5,6 +5,9 @@ import (
 	"fmt"
 	"github.com/kataras/iris"
 	"github.com/kataras/iris/context"
+	"github.com/peterq/pan-light/server/pc-api/middleware"
+	"strings"
+	"time"
 )
 
 func ApiHandler(handler func(ctx context.Context, param map[string]interface{}) (result interface{}, err error)) func(ctx context.Context) {
@@ -63,3 +66,66 @@ func ApiRecover(ctx context.Context) {
 	}()
 	ctx.Next()
 }
+
+type ThrottleOption struct {
+	Duration time.Duration // 时间窗口
+	Number   int           // 允许操作次数
+	GetKey   func(ctx context.Context) string
+}
+
+type throttleState struct {
+	Time  int64   `json:"time"`
+	Water float64 `json:"water"`
+}
+
+func (o ThrottleOption) hit(ctx context.Context) time.Duration {
+	key := o.GetKey(ctx)
+	stateKey := "throttle-" + key
+
+	var state throttleState
+	RedisGet(stateKey, &state)
+	// 计算现有数量
+	speed := float64(o.Number) / float64(o.Duration/time.Second) // 每秒流失的水
+	du := time.Duration(time.Now().UnixNano() - state.Time)
+	water := state.Water - speed*float64(du/time.Second)
+	if water < 0 {
+		water = 0
+	}
+	water++
+	if water > float64(o.Number) {
+		return time.Duration((water-float64(o.Number))/speed) * time.Second
+	}
+	state.Time = time.Now().UnixNano()
+	state.Water = water
+	err := RedisSet(stateKey, state, o.Duration)
+	if err != nil {
+		panic(err)
+	}
+	return 0
+}
+
+// 频率限制器, 漏斗算法
+func Throttle(options ...ThrottleOption) func(ctx context.Context) {
+	for i, option := range options {
+		if option.GetKey == nil {
+			options[i].GetKey = func(ctx context.Context) string {
+				return strings.Join([]string{
+					ctx.GetCurrentRoute().Name(),
+					middleware.CotextLoginInfo(ctx).Uk(),
+					fmt.Sprint(option.Duration),
+					fmt.Sprint(option.Number),
+				}, ".")
+			}
+		}
+	}
+	return func(ctx context.Context) {
+		for _, option := range options {
+			ban := option.hit(ctx)
+			if ban > 0 {
+				panic(NewError(fmt.Sprintf("run out of %d call in %s, try after %s",
+					option.Number, option.Duration, ban), iris.StatusTooManyRequests, nil))
+			}
+		}
+		ctx.Next()
+	}
+}

server/conf/conf.go

 package conf
 
 import (
+	"github.com/go-redis/redis"
 	"github.com/kataras/iris"
 	"gopkg.in/mgo.v2"
 	"os"
@@ -11,11 +12,13 @@ type conf struct {
 	AppSecret  string
 	MongodbUri string
 	Database   string
+	Redis      redis.RingOptions
 }
 
 var Conf *conf
 var IrisConf iris.Configuration
 var MongodbSession *mgo.Session
+var Redis *redis.Ring
 
 func init() {
 	confFile, ok := os.LookupEnv("pan_light_server_conf")
@@ -27,7 +30,19 @@ func init() {
 		AppSecret:  getConf("app-secret").(string),
 		MongodbUri: getConf("mongodb-uri").(string),
 		Database:   getConf("database").(string),
+		Redis: redis.RingOptions{
+			Addrs: map[string]string{
+				"main": getConf("redis.addr").(string),
+			},
+			Password: getConf("redis.pwd").(string),
+			DB:       getConf("redis.db").(int),
+		},
 	}
+	connectMongo()
+	connectRedis()
+}
+
+func connectMongo() {
 	var err error
 	MongodbSession, err = mgo.Dial(Conf.MongodbUri)
 	if err != nil {
@@ -36,6 +51,10 @@ func init() {
 	MongodbSession.Refresh()
 }
 
+func connectRedis() {
+	Redis = redis.NewRing(&Conf.Redis)
+}
+
 func getConf(key string) interface{} {
 	p := strings.Split(key, ".")
 	cnf := IrisConf.Other
@@ -47,7 +66,11 @@ func getConf(key string) interface{} {
 		if idx == len(p)-1 {
 			return parent[name]
 		}
-		parent = parent[name].(map[interface{}]interface{})
+		if idx == 0 {
+			parent = cnf[name].(map[interface{}]interface{})
+		} else {
+			parent = parent[name].(map[interface{}]interface{})
+		}
 	}
 	return nil
 }

server/go.mod

@@ -2,6 +2,8 @@ module github.com/peterq/pan-light/server
 
 require (
 	github.com/dgrijalva/jwt-go v3.2.0+incompatible
+	github.com/go-redis/cache v6.4.0+incompatible
+	github.com/go-redis/redis v6.15.2+incompatible
 	github.com/iris-contrib/middleware v0.0.0-20181021162728-8bd5d51b3c2e
 	github.com/juju/loggo v0.0.0-20190526231331-6e530bcce5d8 // indirect
 	github.com/juju/testing v0.0.0-20190429233213-dfc56b8c09fc // indirect
@@ -13,10 +15,11 @@ require (
 	github.com/onsi/gomega v1.5.0 // indirect
 	github.com/pkg/errors v0.8.0
 	github.com/smartystreets/goconvey v0.0.0-20190330032615-68dc04aab96a // indirect
+	github.com/vmihailenco/msgpack v4.0.4+incompatible // indirect
 	github.com/yudai/pp v2.0.1+incompatible // indirect
 	golang.org/x/net v0.0.0-20190311183353-d8887717615a
 	golang.org/x/sys v0.0.0-20190602015325-4c4f7f33c9ed // indirect
-	golang.org/x/tools v0.0.0-20190328211700-ab21143f2384
+	google.golang.org/appengine v1.6.0 // indirect
 	gopkg.in/mgo.v2 v2.0.0-20180705113604-9856a29383ce
 )
 

server/go.sum

@@ -28,6 +28,10 @@ github.com/fsnotify/fsnotify v1.4.7 h1:IXs+QLmnXW2CcXuY+8Mzv/fWEsPGWxqefPtCP5CnV
 github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
 github.com/gavv/monotime v0.0.0-20171021193802-6f8212e8d10d h1:oYXrtNhqNKL1dVtKdv8XUq5zqdGVFNQ0/4tvccXZOLM=
 github.com/gavv/monotime v0.0.0-20171021193802-6f8212e8d10d/go.mod h1:vmp8DIyckQMXOPl0AQVHt+7n5h7Gb7hS6CUydiV8QeA=
+github.com/go-redis/cache v6.4.0+incompatible h1:ZaeoZofvBZmMr8ZKxzFDmkoRTSp8sxHdJlB3e3T6GDA=
+github.com/go-redis/cache v6.4.0+incompatible/go.mod h1:XNnMdvlNjcZvHjsscEozHAeOeSE5riG9Fj54meG4WT4=
+github.com/go-redis/redis v6.15.2+incompatible h1:9SpNVG76gr6InJGxoZ6IuuxaCOQwDAhzyXg+Bs+0Sb4=
+github.com/go-redis/redis v6.15.2+incompatible/go.mod h1:NAIEuMOZ/fxfXJIrKDQDz8wamY7mA7PouImQ2Jvg6kA=
 github.com/golang/protobuf v1.2.0 h1:P3YflyNX/ehuJFLhxviNdFxQPkGK5cDcApsge1SqnvM=
 github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/gomodule/redigo v2.0.0+incompatible/go.mod h1:B4C85qUVwatsJoIUNIfCRsp7qO0iAmpGFZ4EELWSbC4=
@@ -113,6 +117,8 @@ github.com/smartystreets/goconvey v0.0.0-20190330032615-68dc04aab96a/go.mod h1:s
 github.com/stretchr/testify v1.2.2 h1:bSDNvY7ZPG5RlJ8otE/7V6gMiyenm9RtJ7IUVIAoJ1w=
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
 github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
+github.com/vmihailenco/msgpack v4.0.4+incompatible h1:dSLoQfGFAo3F6OoNhwUmLwVgaUXK79GlxNBwueZn0xI=
+github.com/vmihailenco/msgpack v4.0.4+incompatible/go.mod h1:fy3FlTQTDXWkZ7Bh6AcGMlsjHatGryHQYUTf1ShIgkk=
 github.com/xeipuuv/gojsonpointer v0.0.0-20180127040702-4e3ac2762d5f h1:J9EGpcZtP0E/raorCMxlFGSTBrsSlaDGf3jU/qvAE2c=
 github.com/xeipuuv/gojsonpointer v0.0.0-20180127040702-4e3ac2762d5f/go.mod h1:N2zxlSyiKSe5eX1tZViRH5QA0qijqEDrYZiPEAiq3wU=
 github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415 h1:EzJWgHovont7NscjpAxXsDA8S8BMYve8Y5+7cuRE7R0=
@@ -131,6 +137,7 @@ golang.org/x/crypto v0.0.0-20181112202954-3d3f9f413869 h1:kkXA53yGe04D0adEYJwEVQ
 golang.org/x/crypto v0.0.0-20181112202954-3d3f9f413869/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2 h1:VklqNMn3ovrHsnt90PveolxSbWFaJdECFbxSq0Mqo2M=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
+golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20181023162649-9b4f9f5ad519/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20181114220301-adae6a3d119a h1:gOpx8G595UYyvj8UK4+OFyY4rx037g3fmfhe5SasG3U=
@@ -148,6 +155,8 @@ golang.org/x/text v0.3.0 h1:g61tztE5qeGQ89tm6NTjjM9VPIm088od1l6aSorWRWg=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/tools v0.0.0-20190328211700-ab21143f2384 h1:TFlARGu6Czu1z7q93HTxcP1P+/ZFC/IKythI5RzrnRg=
 golang.org/x/tools v0.0.0-20190328211700-ab21143f2384/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
+google.golang.org/appengine v1.6.0 h1:Tfd7cKwKbFRsI8RMAD3oqqw7JPFRrvFlOsfbgVkjOOw=
+google.golang.org/appengine v1.6.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

server/pc-api/handlers.go

@@ -73,6 +73,6 @@ func handleLogin(ctx context.Context, param map[string]interface{}) (result inte
 }
 
 func handleFeedBack(ctx context.Context, param map[string]interface{}) (result interface{}, err error) {
-	result = ctx.Values().Get(conf.CtxPcLogin).(*middleware.PcLoginInfo).User()
+	result = middleware.CotextLoginInfo(ctx).User()
 	return
 }

server/pc-api/middleware/pc-jwt.go

@@ -99,3 +99,7 @@ func ParseToken(token string) (claim jwt.MapClaims, err error) {
 	}
 	return
 }
+
+func CotextLoginInfo(ctx context.Context) *PcLoginInfo {
+	return ctx.Values().Get(conf.CtxPcLogin).(*PcLoginInfo)
+}

server/pc-api/pc-api-router.go

@@ -2,8 +2,11 @@ package pc_api
 
 import (
 	"github.com/kataras/iris"
+	"github.com/kataras/iris/context"
+	"github.com/kataras/iris/core/router"
 	"github.com/peterq/pan-light/server/artisan"
 	"github.com/peterq/pan-light/server/pc-api/middleware"
+	"time"
 )
 
 func Init(app *iris.Application) {
@@ -12,6 +15,20 @@ func Init(app *iris.Application) {
 
 	pc := app.Party("api/pc/")
 	pc.Use(middleware.PcJwtAuth)
-	pc.Done()
-	pc.Post("feedback", artisan.ApiHandler(handleFeedBack))
+	pc.Use(artisan.Throttle(artisan.ThrottleOption{ // 防止攻击
+		Duration: time.Second * 5,
+		Number:   20,
+		GetKey: func(ctx context.Context) string {
+			return "pc.api." + middleware.CotextLoginInfo(ctx).Uk()
+		},
+	}))
+	pcAuthRoutes(pc)
+}
+
+// 需要登录的api
+func pcAuthRoutes(r router.Party) {
+	r.Post("feedback", artisan.Throttle(artisan.ThrottleOption{
+		Duration: time.Hour,
+		Number:   5,
+	}), artisan.ApiHandler(handleFeedBack))
 }