app/ratelimit.go
132 lines · 3492 bytes
1package app
2
3import (
4 "net"
5 "net/http"
6 "strings"
7 "sync"
8 "time"
9)
10
11// rateLimiter is a token bucket per client address. It bounds how many page,
12// feed and API requests one client can make, so a single crawler cannot spend
13// the whole upstream budget and turn the DeviantArt throttle into latency for
14// everyone else.
15type rateLimiter struct {
16 perSecond float64
17 burst float64
18 now func() time.Time
19
20 mu sync.Mutex
21 buckets map[string]*bucket
22 inserts int
23}
24
25type bucket struct {
26 tokens float64
27 last time.Time
28}
29
30// bucketIdle is how long a client can go unseen before its bucket is dropped.
31const bucketIdle = 10 * time.Minute
32
33func newRateLimiter(perMinute, burst int) *rateLimiter {
34 return &rateLimiter{
35 perSecond: float64(perMinute) / 60,
36 burst: float64(burst),
37 now: time.Now,
38 buckets: map[string]*bucket{},
39 }
40}
41
42// allow reports whether client may make a request now, spending one token if
43// so. A client's first request finds a full bucket.
44func (l *rateLimiter) allow(client string) bool {
45 l.mu.Lock()
46 defer l.mu.Unlock()
47
48 now := l.now()
49 b := l.buckets[client]
50 if b == nil {
51 b = &bucket{tokens: l.burst, last: now}
52 l.buckets[client] = b
53 l.inserts++
54 if l.inserts%1000 == 0 {
55 l.prune(now)
56 }
57 } else {
58 b.tokens = min(l.burst, b.tokens+now.Sub(b.last).Seconds()*l.perSecond)
59 b.last = now
60 }
61
62 if b.tokens < 1 {
63 return false
64 }
65 b.tokens--
66 return true
67}
68
69// prune drops buckets idle past bucketIdle. The caller holds mu.
70func (l *rateLimiter) prune(now time.Time) {
71 for client, b := range l.buckets {
72 if now.Sub(b.last) > bucketIdle {
73 delete(l.buckets, client)
74 }
75 }
76}
77
78// clientAddr is the address a request is limited by. Behind a reverse proxy
79// every connection arrives from the proxy, so when the connection is from a
80// loopback or private address the client is the rightmost X-Forwarded-For
81// entry, which is the one that proxy appended. A direct client's connection is
82// not from a private address, so its own header is never trusted.
83func clientAddr(r *http.Request) string {
84 host, _, err := net.SplitHostPort(r.RemoteAddr)
85 if err != nil {
86 host = r.RemoteAddr
87 }
88 ip := net.ParseIP(host)
89 if ip == nil || (!ip.IsLoopback() && !ip.IsPrivate()) {
90 return host
91 }
92 xff := r.Header.Get("X-Forwarded-For")
93 if xff == "" {
94 return host
95 }
96 parts := strings.Split(xff, ",")
97 if forwarded := strings.TrimSpace(parts[len(parts)-1]); forwarded != "" {
98 return forwarded
99 }
100 return host
101}
102
103// limited reports whether an endpoint counts against the client's budget.
104// Media, avatars and static assets do not: one listing page loads twenty
105// thumbnails, which would exhaust any sensible page budget.
106func limited(endpoint string) bool {
107 switch endpoint {
108 case "media", "stylesheet", "favicon.ico", "robots.txt":
109 return false
110 }
111 return true
112}
113
114// robotsTXT keeps crawlers off the pages that fan out into unbounded upstream
115// requests. The index, daily deviations and posts stay allowed, so the instance
116// is still discoverable.
117func robotsTXT(base string) string {
118 var b strings.Builder
119 b.WriteString("User-agent: *\n")
120 for _, p := range []string{"search", "api", "group_user", "media", "*?p="} {
121 b.WriteString("Disallow: ")
122 b.WriteString(base)
123 b.WriteString(p)
124 b.WriteString("\n")
125 }
126 b.WriteString("Crawl-delay: 10\n")
127 return b.String()
128}
129
130// daLimiter is the running instance's limiter, or nil when rate-limit.per-minute
131// is 0. Set by ExecuteConfig.
132var daLimiter *rateLimiter