forked from Onyx-Protocol/Onyx
-
Notifications
You must be signed in to change notification settings - Fork 0
/
limit.go
60 lines (51 loc) · 1.15 KB
/
limit.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
// Package limit provides a rate limiting HTTP handler.
package limit
import (
"net/http"
"sync"
"golang.org/x/time/rate"
)
type handler struct {
next http.Handler
limited http.Handler
f func(*http.Request) string
freq rate.Limit
burst int
bucketMu sync.Mutex // protects the following
buckets map[string]*rate.Limiter
}
func Handler(next, limited http.Handler, freq, burst int, f func(*http.Request) string) http.Handler {
return &handler{
next: next,
limited: limited,
f: f,
freq: rate.Limit(freq),
burst: burst,
buckets: make(map[string]*rate.Limiter),
}
}
func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
id := h.f(r)
if !h.bucket(id).Allow() {
h.limited.ServeHTTP(w, r)
return
}
h.next.ServeHTTP(w, r)
}
func (h *handler) bucket(id string) *rate.Limiter {
h.bucketMu.Lock()
bucket, ok := h.buckets[id]
if !ok {
bucket = rate.NewLimiter(h.freq, h.burst)
h.buckets[id] = bucket
}
h.bucketMu.Unlock()
return bucket
}
func RemoteAddrID(r *http.Request) string {
return r.RemoteAddr
}
func AuthUserID(r *http.Request) string {
user, _, _ := r.BasicAuth()
return user
}