/
base_limiter.go
87 lines (77 loc) · 1.94 KB
/
base_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
package rlutils
import (
"fmt"
"net/http"
"path/filepath"
"strings"
"time"
"github.com/2manymws/rl"
"github.com/2manymws/rl/counter"
)
const (
RemoteAddrKey = "remote_addr"
HostKey = "host"
)
type BaseLimiter struct {
reqLimit int `mapstructure:"req_limit"`
windowLen time.Duration
targetExtensions map[string]struct{}
onRequestLimit func(*rl.Context, string) http.HandlerFunc
rl.Counter
}
func NewBaseLimiter(
reqLimit int,
windowLen time.Duration,
targetExtensions []string,
onRequestLimit func(*rl.Context, string) http.HandlerFunc,
) BaseLimiter {
ttl := windowLen * 2 // 最低2回分のウィンドウ分のカウンタを維持する
targetExtensionsMap := make(map[string]struct{}, len(targetExtensions))
if len(targetExtensions) > 0 {
for _, ext := range targetExtensions {
if len(ext) > 0 && ext[0] != '.' {
ext = "." + ext
}
targetExtensionsMap[strings.ToLower(ext)] = struct{}{}
}
}
return BaseLimiter{
reqLimit: reqLimit,
windowLen: windowLen,
Counter: counter.New(ttl),
targetExtensions: targetExtensionsMap,
onRequestLimit: onRequestLimit,
}
}
func (l *BaseLimiter) ShouldSetXRateLimitHeaders(r *rl.Context) bool {
return false
}
func (l *BaseLimiter) Name() string {
return "base_limiter"
}
func (l *BaseLimiter) IsTargetRequest(r *http.Request) bool {
return l.isTargetExtensions(r)
}
func (l *BaseLimiter) isTargetExtensions(r *http.Request) bool {
if len(l.targetExtensions) == 0 {
return true
}
extension := strings.ToLower(filepath.Ext(r.URL.Path))
_, ok := l.targetExtensions[extension]
return ok
}
func validateKey(key string) error {
for _, k := range []string{RemoteAddrKey, HostKey} {
if k == key {
return nil
}
}
return fmt.Errorf("invalid key: %s", key)
}
func fillKey(r *http.Request, key string) string {
if key == RemoteAddrKey {
remoteAddr := strings.Split(r.RemoteAddr, ":")[0]
return remoteAddr
}
return r.Host
}