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
|
package main
import (
"net"
"net/http"
"strings"
"sync"
"time"
)
type limiter struct {
mu sync.Mutex
hits map[string][]time.Time
max int
window time.Duration
}
func newLimiter(max int, window time.Duration) *limiter {
return &limiter{hits: map[string][]time.Time{}, max: max, window: window}
}
func (l *limiter) prune(key string, now time.Time) []time.Time {
cutoff := now.Add(-l.window)
kept := l.hits[key][:0]
for _, t := range l.hits[key] {
if t.After(cutoff) {
kept = append(kept, t)
}
}
if len(kept) == 0 {
delete(l.hits, key)
} else {
l.hits[key] = kept
}
if len(l.hits) > 10000 {
for k, v := range l.hits {
if !v[len(v)-1].After(cutoff) {
delete(l.hits, k)
}
}
}
return kept
}
func (l *limiter) blocked(key string) bool {
l.mu.Lock()
defer l.mu.Unlock()
return len(l.prune(key, time.Now())) >= l.max
}
func (l *limiter) fail(key string) {
l.mu.Lock()
defer l.mu.Unlock()
now := time.Now()
l.hits[key] = append(l.prune(key, now), now)
}
func clientIP(r *http.Request) string {
host, _, err := net.SplitHostPort(r.RemoteAddr)
if err != nil {
host = r.RemoteAddr
}
if ip := net.ParseIP(host); ip != nil && ip.IsLoopback() {
if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
hops := strings.Split(xff, ",")
return strings.TrimSpace(hops[len(hops)-1])
}
}
return host
}
type bucket struct {
tokens float64
last time.Time
}
type buckets struct {
mu sync.Mutex
m map[string]*bucket
rate float64
burst float64
}
func newBuckets(rate, burst float64) *buckets {
return &buckets{m: map[string]*bucket{}, rate: rate, burst: burst}
}
func (b *buckets) allow(key string) bool {
b.mu.Lock()
defer b.mu.Unlock()
now := time.Now()
k, ok := b.m[key]
if !ok {
if len(b.m) >= 20000 {
for name, v := range b.m {
if v.tokens+now.Sub(v.last).Seconds()*b.rate >= b.burst {
delete(b.m, name)
}
}
if len(b.m) >= 20000 {
return false
}
}
k = &bucket{tokens: b.burst, last: now}
b.m[key] = k
}
k.tokens = min(b.burst, k.tokens+now.Sub(k.last).Seconds()*b.rate)
k.last = now
if k.tokens < 1 {
return false
}
k.tokens--
return true
}
func heavyRequest(r *http.Request) bool {
p := r.URL.Path
q := r.URL.Query()
return p == "/search" || strings.Contains(p, "/archive/") || strings.Contains(p, "/blame/") ||
strings.Contains(p, "/interdiff") || (strings.HasSuffix(p, "/search") && q.Get("q") != "") ||
(strings.Contains(p, "/log") && q.Get("q") != "") || p == "/login" || p == "/register"
}
func (s *Server) throttle(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if strings.HasPrefix(r.URL.Path, "/static/") {
next.ServeHTTP(w, r)
return
}
ip := clientIP(r)
if !s.requests.allow(ip) || (heavyRequest(r) && !s.heavy.allow(ip)) {
w.Header().Set("Retry-After", "10")
http.Error(w, "too many requests, slow down", http.StatusTooManyRequests)
return
}
next.ServeHTTP(w, r)
})
}
|