app/ratelimit.go

v1.5.3
skunky-art/app/ratelimit.go history · blame · raw

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