Skip to content

Commit e8e47aa

Browse files
peterqpeterq
committed
add: 基于redis的api调用频率限制middleware
1 parent 6b5023b commit e8e47aa

File tree

8 files changed

+158
-5
lines changed

8 files changed

+158
-5
lines changed

‎server/artisan/redis.go‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
package artisan
2+
3+
import (
4+
"encoding/json"
5+
"github.com/go-redis/cache"
6+
"github.com/peterq/pan-light/server/conf"
7+
"time"
8+
)
9+
10+
var codec = &cache.Codec{
11+
Redis: conf.Redis,
12+
13+
Marshal: func(v interface{}) ([]byte, error) {
14+
return json.Marshal(v)
15+
},
16+
Unmarshal: func(b []byte, v interface{}) error {
17+
return json.Unmarshal(b, v)
18+
},
19+
}
20+
21+
func RedisGet(key string, data interface{}) error {
22+
return codec.Get(key, data)
23+
}
24+
25+
func RedisSet(key string, value interface{}, expiration time.Duration) error {
26+
return codec.Set(&cache.Item{
27+
Key: key,
28+
Object: value,
29+
Expiration: expiration,
30+
})
31+
}

‎server/artisan/web.go‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,9 @@ import (
55
"fmt"
66
"github.com/kataras/iris"
77
"github.com/kataras/iris/context"
8+
"github.com/peterq/pan-light/server/pc-api/middleware"
9+
"strings"
10+
"time"
811
)
912

1013
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) {
6366
}()
6467
ctx.Next()
6568
}
69+
70+
type ThrottleOption struct {
71+
Duration time.Duration // 时间窗口
72+
Number int // 允许操作次数
73+
GetKey func(ctx context.Context) string
74+
}
75+
76+
type throttleState struct {
77+
Time int64 `json:"time"`
78+
Water float64 `json:"water"`
79+
}
80+
81+
func (o ThrottleOption) hit(ctx context.Context) time.Duration {
82+
key := o.GetKey(ctx)
83+
stateKey := "throttle-" + key
84+
85+
var state throttleState
86+
RedisGet(stateKey, &state)
87+
// 计算现有数量
88+
speed := float64(o.Number) / float64(o.Duration/time.Second) // 每秒流失的水
89+
du := time.Duration(time.Now().UnixNano() - state.Time)
90+
water := state.Water - speed*float64(du/time.Second)
91+
if water < 0 {
92+
water = 0
93+
}
94+
water++
95+
if water > float64(o.Number) {
96+
return time.Duration((water-float64(o.Number))/speed) * time.Second
97+
}
98+
state.Time = time.Now().UnixNano()
99+
state.Water = water
100+
err := RedisSet(stateKey, state, o.Duration)
101+
if err != nil {
102+
panic(err)
103+
}
104+
return 0
105+
}
106+
107+
// 频率限制器, 漏斗算法
108+
func Throttle(options ...ThrottleOption) func(ctx context.Context) {
109+
for i, option := range options {
110+
if option.GetKey == nil {
111+
options[i].GetKey = func(ctx context.Context) string {
112+
return strings.Join([]string{
113+
ctx.GetCurrentRoute().Name(),
114+
middleware.CotextLoginInfo(ctx).Uk(),
115+
fmt.Sprint(option.Duration),
116+
fmt.Sprint(option.Number),
117+
}, ".")
118+
}
119+
}
120+
}
121+
return func(ctx context.Context) {
122+
for _, option := range options {
123+
ban := option.hit(ctx)
124+
if ban > 0 {
125+
panic(NewError(fmt.Sprintf("run out of %d call in %s, try after %s",
126+
option.Number, option.Duration, ban), iris.StatusTooManyRequests, nil))
127+
}
128+
}
129+
ctx.Next()
130+
}
131+
}

‎server/conf/conf.go‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package conf
22

33
import (
4+
"github.com/go-redis/redis"
45
"github.com/kataras/iris"
56
"gopkg.in/mgo.v2"
67
"os"
@@ -11,11 +12,13 @@ type conf struct {
1112
AppSecret string
1213
MongodbUri string
1314
Database string
15+
Redis redis.RingOptions
1416
}
1517

1618
var Conf *conf
1719
var IrisConf iris.Configuration
1820
var MongodbSession *mgo.Session
21+
var Redis *redis.Ring
1922

2023
func init() {
2124
confFile, ok := os.LookupEnv("pan_light_server_conf")
@@ -27,7 +30,19 @@ func init() {
2730
AppSecret: getConf("app-secret").(string),
2831
MongodbUri: getConf("mongodb-uri").(string),
2932
Database: getConf("database").(string),
33+
Redis: redis.RingOptions{
34+
Addrs: map[string]string{
35+
"main": getConf("redis.addr").(string),
36+
},
37+
Password: getConf("redis.pwd").(string),
38+
DB: getConf("redis.db").(int),
39+
},
3040
}
41+
connectMongo()
42+
connectRedis()
43+
}
44+
45+
func connectMongo() {
3146
var err error
3247
MongodbSession, err = mgo.Dial(Conf.MongodbUri)
3348
if err != nil {
@@ -36,6 +51,10 @@ func init() {
3651
MongodbSession.Refresh()
3752
}
3853

54+
func connectRedis() {
55+
Redis = redis.NewRing(&Conf.Redis)
56+
}
57+
3958
func getConf(key string) interface{} {
4059
p := strings.Split(key, ".")
4160
cnf := IrisConf.Other
@@ -47,7 +66,11 @@ func getConf(key string) interface{} {
4766
if idx == len(p)-1 {
4867
return parent[name]
4968
}
50-
parent = parent[name].(map[interface{}]interface{})
69+
if idx == 0 {
70+
parent = cnf[name].(map[interface{}]interface{})
71+
} else {
72+
parent = parent[name].(map[interface{}]interface{})
73+
}
5174
}
5275
return nil
5376
}

‎server/go.mod‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@ module github.com/peterq/pan-light/server
22

33
require (
44
github.com/dgrijalva/jwt-go v3.2.0+incompatible
5+
github.com/go-redis/cache v6.4.0+incompatible
6+
github.com/go-redis/redis v6.15.2+incompatible
57
github.com/iris-contrib/middleware v0.0.0-20181021162728-8bd5d51b3c2e
68
github.com/juju/loggo v0.0.0-20190526231331-6e530bcce5d8 // indirect
79
github.com/juju/testing v0.0.0-20190429233213-dfc56b8c09fc // indirect
@@ -13,10 +15,11 @@ require (
1315
github.com/onsi/gomega v1.5.0 // indirect
1416
github.com/pkg/errors v0.8.0
1517
github.com/smartystreets/goconvey v0.0.0-20190330032615-68dc04aab96a // indirect
18+
github.com/vmihailenco/msgpack v4.0.4+incompatible // indirect
1619
github.com/yudai/pp v2.0.1+incompatible // indirect
1720
golang.org/x/net v0.0.0-20190311183353-d8887717615a
1821
golang.org/x/sys v0.0.0-20190602015325-4c4f7f33c9ed // indirect
19-
golang.org/x/tools v0.0.0-20190328211700-ab21143f2384
22+
google.golang.org/appengine v1.6.0 // indirect
2023
gopkg.in/mgo.v2 v2.0.0-20180705113604-9856a29383ce
2124
)
2225

‎server/go.sum‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,10 @@ github.com/fsnotify/fsnotify v1.4.7 h1:IXs+QLmnXW2CcXuY+8Mzv/fWEsPGWxqefPtCP5CnV
2828
github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
2929
github.com/gavv/monotime v0.0.0-20171021193802-6f8212e8d10d h1:oYXrtNhqNKL1dVtKdv8XUq5zqdGVFNQ0/4tvccXZOLM=
3030
github.com/gavv/monotime v0.0.0-20171021193802-6f8212e8d10d/go.mod h1:vmp8DIyckQMXOPl0AQVHt+7n5h7Gb7hS6CUydiV8QeA=
31+
github.com/go-redis/cache v6.4.0+incompatible h1:ZaeoZofvBZmMr8ZKxzFDmkoRTSp8sxHdJlB3e3T6GDA=
32+
github.com/go-redis/cache v6.4.0+incompatible/go.mod h1:XNnMdvlNjcZvHjsscEozHAeOeSE5riG9Fj54meG4WT4=
33+
github.com/go-redis/redis v6.15.2+incompatible h1:9SpNVG76gr6InJGxoZ6IuuxaCOQwDAhzyXg+Bs+0Sb4=
34+
github.com/go-redis/redis v6.15.2+incompatible/go.mod h1:NAIEuMOZ/fxfXJIrKDQDz8wamY7mA7PouImQ2Jvg6kA=
3135
github.com/golang/protobuf v1.2.0 h1:P3YflyNX/ehuJFLhxviNdFxQPkGK5cDcApsge1SqnvM=
3236
github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
3337
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
113117
github.com/stretchr/testify v1.2.2 h1:bSDNvY7ZPG5RlJ8otE/7V6gMiyenm9RtJ7IUVIAoJ1w=
114118
github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
115119
github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
120+
github.com/vmihailenco/msgpack v4.0.4+incompatible h1:dSLoQfGFAo3F6OoNhwUmLwVgaUXK79GlxNBwueZn0xI=
121+
github.com/vmihailenco/msgpack v4.0.4+incompatible/go.mod h1:fy3FlTQTDXWkZ7Bh6AcGMlsjHatGryHQYUTf1ShIgkk=
116122
github.com/xeipuuv/gojsonpointer v0.0.0-20180127040702-4e3ac2762d5f h1:J9EGpcZtP0E/raorCMxlFGSTBrsSlaDGf3jU/qvAE2c=
117123
github.com/xeipuuv/gojsonpointer v0.0.0-20180127040702-4e3ac2762d5f/go.mod h1:N2zxlSyiKSe5eX1tZViRH5QA0qijqEDrYZiPEAiq3wU=
118124
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
131137
golang.org/x/crypto v0.0.0-20181112202954-3d3f9f413869/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4=
132138
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2 h1:VklqNMn3ovrHsnt90PveolxSbWFaJdECFbxSq0Mqo2M=
133139
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
140+
golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
134141
golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
135142
golang.org/x/net v0.0.0-20181023162649-9b4f9f5ad519/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
136143
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=
148155
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
149156
golang.org/x/tools v0.0.0-20190328211700-ab21143f2384 h1:TFlARGu6Czu1z7q93HTxcP1P+/ZFC/IKythI5RzrnRg=
150157
golang.org/x/tools v0.0.0-20190328211700-ab21143f2384/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
158+
google.golang.org/appengine v1.6.0 h1:Tfd7cKwKbFRsI8RMAD3oqqw7JPFRrvFlOsfbgVkjOOw=
159+
google.golang.org/appengine v1.6.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4=
151160
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
152161
gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
153162
gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

‎server/pc-api/handlers.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,6 @@ func handleLogin(ctx context.Context, param map[string]interface{}) (result inte
7373
}
7474

7575
func handleFeedBack(ctx context.Context, param map[string]interface{}) (result interface{}, err error) {
76-
result = ctx.Values().Get(conf.CtxPcLogin).(*middleware.PcLoginInfo).User()
76+
result = middleware.CotextLoginInfo(ctx).User()
7777
return
7878
}

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,3 +99,7 @@ func ParseToken(token string) (claim jwt.MapClaims, err error) {
9999
}
100100
return
101101
}
102+
103+
func CotextLoginInfo(ctx context.Context) *PcLoginInfo {
104+
return ctx.Values().Get(conf.CtxPcLogin).(*PcLoginInfo)
105+
}

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

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,11 @@ package pc_api
22

33
import (
44
"github.com/kataras/iris"
5+
"github.com/kataras/iris/context"
6+
"github.com/kataras/iris/core/router"
57
"github.com/peterq/pan-light/server/artisan"
68
"github.com/peterq/pan-light/server/pc-api/middleware"
9+
"time"
710
)
811

912
func Init(app *iris.Application) {
@@ -12,6 +15,20 @@ func Init(app *iris.Application) {
1215

1316
pc := app.Party("api/pc/")
1417
pc.Use(middleware.PcJwtAuth)
15-
pc.Done()
16-
pc.Post("feedback", artisan.ApiHandler(handleFeedBack))
18+
pc.Use(artisan.Throttle(artisan.ThrottleOption{ // 防止攻击
19+
Duration: time.Second * 5,
20+
Number: 20,
21+
GetKey: func(ctx context.Context) string {
22+
return "pc.api." + middleware.CotextLoginInfo(ctx).Uk()
23+
},
24+
}))
25+
pcAuthRoutes(pc)
26+
}
27+
28+
// 需要登录的api
29+
func pcAuthRoutes(r router.Party) {
30+
r.Post("feedback", artisan.Throttle(artisan.ThrottleOption{
31+
Duration: time.Hour,
32+
Number: 5,
33+
}), artisan.ApiHandler(handleFeedBack))
1734
}

0 commit comments

Comments
 (0)