/
a_rate.go
80 lines (71 loc) · 2.1 KB
/
a_rate.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
package pipe
import (
"github.com/kataras/iris/v12"
"github.com/ulule/limiter/v3"
"github.com/ulule/limiter/v3/drivers/store/memory"
"net/http"
"strconv"
)
type RateLimitPipe struct {
// <limit>-<period>
// * "S": second
// * "M": minute
// * "H": hour
// * "D": day
//
// Examples:
//
// * 5 reqs/second: "5-S"
// * 10 reqs/minute: "10-M"
// * 1000 reqs/hour: "1000-H"
// * 2000 reqs/day: "2000-D"
RatePeriod string `json:"rate_period,omitempty"`
Store string `json:"store,omitempty"` // 存储 暂时 memory
KeyGen *StrExpand `json:"key_gen,omitempty"` // 判断key的来源
WriteHeader bool `json:"write_header,omitempty"` // 是否把rate状态写入到请求header中
}
func (c *RateLimitPipe) genClient() (*limiter.Limiter, error) {
rate, err := limiter.NewRateFromFormatted(c.RatePeriod)
if err != nil {
return nil, err
}
store := memory.NewStore()
return limiter.New(store, rate), nil
}
var (
// RequestRate 限速器 使用 https://github.com/ulule/limiter 实现
// 必传params RateLimitPipe
RequestRate = &RunnerContext[any, *RateLimitPipe, any, any]{
Name: "请求限速",
Key: "request_rate",
call: func(ctx iris.Context, origin any, params *RateLimitPipe, db any, more ...any) *RunResp[any] {
if params == nil {
return NewPipeErr[any](PipePackParamsError)
}
k, err := params.KeyGen.Build()
if err != nil {
return NewPipeErr[any](err)
}
client, err := params.genClient()
if err != nil {
return NewPipeErr[any](err)
}
// 过程报错了 并不是说达到限速了
rateCtx, err := client.Get(ctx, k)
if err != nil {
return NewPipeErr[any](err)
}
if params.WriteHeader {
ctx.Header("X-RateLimit-Limit", strconv.FormatInt(rateCtx.Limit, 10))
ctx.Header("X-RateLimit-Remaining", strconv.FormatInt(rateCtx.Remaining, 10))
ctx.Header("X-RateLimit-Reset", strconv.FormatInt(rateCtx.Reset, 10))
}
// 超出限速了
if rateCtx.Reached {
// 抛出错误
return NewPipeErr[any](PipeRatedError).SetReqCode(http.StatusTooManyRequests)
}
return NewPipeResult[any](nil)
},
}
)