forked from didip/tollbooth
/
config.go
74 lines (56 loc) · 1.84 KB
/
config.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
// Package config provides data structure to configure rate-limiter.
package config
import (
"sync"
"time"
"golang.org/x/time/rate"
)
// NewLimiter is a constructor for Limiter.
func NewLimiter(max int64, ttl time.Duration) *Limiter {
limiter := &Limiter{Max: max, TTL: ttl}
limiter.MessageContentType = "text/plain; charset=utf-8"
limiter.Message = "You have reached maximum request limit."
limiter.StatusCode = 429
limiter.tokenBuckets = make(map[string]*rate.Limiter)
limiter.IPLookups = []string{"RemoteAddr", "X-Forwarded-For", "X-Real-IP"}
return limiter
}
// Limiter is a config struct to limit a particular request handler.
type Limiter struct {
// This rate limiter identifier
Name string
// HTTP message when limit is reached.
Message string
// Content-Type for Message
MessageContentType string
// HTTP status code when limit is reached.
StatusCode int
// Maximum number of requests to limit per duration.
Max int64
// Duration of rate-limiter.
TTL time.Duration
// List of places to look up IP address.
// Default is "RemoteAddr", "X-Forwarded-For", "X-Real-IP".
// You can rearrange the order as you like.
IPLookups []string
// List of HTTP Methods to limit (GET, POST, PUT, etc.).
// Empty means limit all methods.
Methods []string
// List of HTTP headers to limit.
// Empty means skip headers checking.
Headers map[string][]string
// List of basic auth usernames to limit.
BasicAuthUsers []string
// Throttler struct
tokenBuckets map[string]*rate.Limiter
sync.RWMutex
}
// LimitReached returns a bool indicating if the Bucket identified by key ran out of tokens.
func (l *Limiter) LimitReached(key string) bool {
l.Lock()
defer l.Unlock()
if _, found := l.tokenBuckets[key]; !found {
l.tokenBuckets[key] = rate.NewLimiter(rate.Every(l.TTL), int(l.Max))
}
return !l.tokenBuckets[key].AllowN(time.Now(), 1)
}