internal/httpd/apilimit.go

  1package httpd
  2
  3import (
  4	"net"
  5	"net/http"
  6	"strconv"
  7	"strings"
  8	"sync"
  9	"time"
 10)
 11
 12// apiLimiter is a token bucket per caller, with a separate, smaller budget
 13// for writes. Unlike the SSH limiter — which counts only auth failures,
 14// because a busy CLI opens many connections legitimately — this counts
 15// every request: one HTTP call is one command, and a client looping over a
 16// list can issue them far faster than a person can type.
 17//
 18// Keyed by token hash when authenticated, so one caller's budget follows
 19// them across networks, and by IP otherwise, so an unauthenticated flood
 20// cannot mint budget by rotating tokens.
 21type apiLimiter struct {
 22	mu      sync.Mutex
 23	buckets map[string]*bucket
 24	rate    float64 // requests per second, sustained
 25	burst   float64
 26	writes  float64 // sustained write rate, a fraction of rate
 27}
 28
 29type bucket struct {
 30	read, write float64
 31	last        time.Time
 32}
 33
 34func newAPILimiter(perMinute int) *apiLimiter {
 35	if perMinute <= 0 {
 36		perMinute = 120
 37	}
 38	rate := float64(perMinute) / 60
 39	return &apiLimiter{
 40		buckets: map[string]*bucket{},
 41		rate:    rate,
 42		burst:   float64(perMinute),
 43		// Writes are rarer and more expensive; a tenth of the read budget
 44		// is generous for a client and useless for a scraper.
 45		writes: rate / 10,
 46	}
 47}
 48
 49// allow reports whether the caller may make this request, and how long to
 50// wait if not. write requests draw on both buckets: a write is also a
 51// request.
 52func (l *apiLimiter) allow(key string, write bool) (bool, time.Duration) {
 53	l.mu.Lock()
 54	defer l.mu.Unlock()
 55	now := time.Now()
 56
 57	if len(l.buckets) > 4096 {
 58		for k, b := range l.buckets {
 59			if now.Sub(b.last) > 10*time.Minute {
 60				delete(l.buckets, k)
 61			}
 62		}
 63	}
 64
 65	b := l.buckets[key]
 66	if b == nil {
 67		b = &bucket{read: l.burst, write: l.burst / 10, last: now}
 68		l.buckets[key] = b
 69	}
 70	elapsed := now.Sub(b.last).Seconds()
 71	b.last = now
 72	b.read = minf(l.burst, b.read+elapsed*l.rate)
 73	b.write = minf(l.burst/10, b.write+elapsed*l.writes)
 74
 75	if b.read < 1 {
 76		return false, retryAfter(1-b.read, l.rate)
 77	}
 78	if write && b.write < 1 {
 79		return false, retryAfter(1-b.write, l.writes)
 80	}
 81	b.read--
 82	if write {
 83		b.write--
 84	}
 85	return true, 0
 86}
 87
 88func retryAfter(deficit, rate float64) time.Duration {
 89	if rate <= 0 {
 90		return time.Minute
 91	}
 92	d := time.Duration(deficit / rate * float64(time.Second))
 93	if d < time.Second {
 94		return time.Second
 95	}
 96	return d
 97}
 98
 99func minf(a, b float64) float64 {
100	if a < b {
101		return a
102	}
103	return b
104}
105
106// clientIP is the address a request is attributed to. With no trusted
107// proxies configured it is the peer, and forwarded headers are ignored:
108// honouring a client-supplied header would let a caller pick their own
109// bucket. When the peer is a trusted proxy, it is the last
110// X-Forwarded-For hop that is not itself a trusted proxy, so a proxied
111// deployment does not collapse every anonymous caller into one bucket
112// (#136).
113func (s *Server) clientIP(r *http.Request) string {
114	peer := r.RemoteAddr
115	if host, _, err := net.SplitHostPort(peer); err == nil {
116		peer = host
117	}
118	if !s.trustedProxy(peer) {
119		return peer
120	}
121	hops := strings.Split(r.Header.Get("X-Forwarded-For"), ",")
122	for i := len(hops) - 1; i >= 0; i-- {
123		hop := strings.TrimSpace(hops[i])
124		if hop != "" && !s.trustedProxy(hop) {
125			return hop
126		}
127	}
128	return peer
129}
130
131func (s *Server) trustedProxy(addr string) bool {
132	ip := net.ParseIP(addr)
133	if ip == nil {
134		return false
135	}
136	for _, n := range s.proxies {
137		if n.Contains(ip) {
138			return true
139		}
140	}
141	return false
142}
143
144func tooManyRequests(w http.ResponseWriter, wait time.Duration) {
145	w.Header().Set("Retry-After", strconv.Itoa(int(wait.Seconds()+0.5)))
146	apiError(w, http.StatusTooManyRequests,
147		"rate limited; retry in "+strconv.Itoa(int(wait.Seconds()+0.5))+"s")
148}