-
Notifications
You must be signed in to change notification settings - Fork 0
/
limiter.go
152 lines (131 loc) · 3.5 KB
/
limiter.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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
// Copyright 2020 radiant
//
package ratelimit
import (
"sync"
"time"
"github.com/W3-Engineers-Ltd/Radiant/server/web"
"github.com/W3-Engineers-Ltd/Radiant/server/web/context"
)
// limiterOption is constructor option
type limiterOption func(l *limiter)
type limiter struct {
sync.RWMutex
capacity uint
rate time.Duration
buckets map[string]bucket
bucketFactory func(opts ...bucketOption) bucket
sessionKey func(ctx *context.Context) string
resp RejectionResponse
}
// RejectionResponse stores response information
// for the request rejected by limiter
type RejectionResponse struct {
code int
body string
}
const perRequestConsumedAmount = 1
var defaultRejectionResponse = RejectionResponse{
code: 429,
body: "too many requests",
}
// NewLimiter return FilterFunc, the limiter enables rate limit
// according to the configuration.
func NewLimiter(opts ...limiterOption) web.FilterFunc {
l := &limiter{
buckets: make(map[string]bucket),
sessionKey: defaultSessionKey,
rate: time.Millisecond * 10,
capacity: 100,
bucketFactory: newTokenBucket,
resp: defaultRejectionResponse,
}
for _, o := range opts {
o(l)
}
return func(ctx *context.Context) {
if !l.take(perRequestConsumedAmount, ctx) {
ctx.ResponseWriter.WriteHeader(l.resp.code)
ctx.WriteString(l.resp.body)
}
}
}
// WithSessionKey return limiterOption. WithSessionKey config func
// which defines the request characteristic against the limit is applied
func WithSessionKey(f func(ctx *context.Context) string) limiterOption {
return func(l *limiter) {
l.sessionKey = f
}
}
// WithRate return limiterOption. WithRate config how long it takes to
// generate a token.
func WithRate(r time.Duration) limiterOption {
return func(l *limiter) {
l.rate = r
}
}
// WithCapacity return limiterOption. WithCapacity config the capacity size.
// The bucket with a capacity of n has n tokens after initialization. The capacity
// defines how many requests a client can make in excess of the rate.
func WithCapacity(c uint) limiterOption {
return func(l *limiter) {
l.capacity = c
}
}
// WithBucketFactory return limiterOption. WithBucketFactory customize the
// implementation of Bucket.
func WithBucketFactory(f func(opts ...bucketOption) bucket) limiterOption {
return func(l *limiter) {
l.bucketFactory = f
}
}
// WithRejectionResponse return limiterOption. WithRejectionResponse
// customize the response for the request rejected by the limiter.
func WithRejectionResponse(resp RejectionResponse) limiterOption {
return func(l *limiter) {
l.resp = resp
}
}
func (l *limiter) take(amount uint, ctx *context.Context) bool {
bucket := l.getBucket(ctx)
if bucket == nil {
return true
}
return bucket.take(amount)
}
func (l *limiter) getBucket(ctx *context.Context) bucket {
key := l.sessionKey(ctx)
l.RLock()
b, ok := l.buckets[key]
l.RUnlock()
if !ok {
b = l.createBucket(key)
}
return b
}
func (l *limiter) createBucket(key string) bucket {
l.Lock()
defer l.Unlock()
// double check avoid overwriting
b, ok := l.buckets[key]
if ok {
return b
}
b = l.bucketFactory(withCapacity(l.capacity), withRate(l.rate))
l.buckets[key] = b
return b
}
func defaultSessionKey(ctx *context.Context) string {
return "BEEGO_ALL"
}
func RemoteIPSessionKey(ctx *context.Context) string {
r := ctx.Request
IPAddress := r.Header.Get("X-Real-Ip")
if IPAddress == "" {
IPAddress = r.Header.Get("X-Forwarded-For")
}
if IPAddress == "" {
IPAddress = r.RemoteAddr
}
return IPAddress
}