app/ratelimit_test.go

  1package app
  2
  3import (
  4	"net/http"
  5	"net/http/httptest"
  6	"strings"
  7	"testing"
  8	"time"
  9)
 10
 11func TestRateLimiterAllowsBurstThenRefuses(t *testing.T) {
 12	l := newRateLimiter(60, 3)
 13	now := time.Now()
 14	l.now = func() time.Time { return now }
 15
 16	for i := range 3 {
 17		if !l.allow("c") {
 18			t.Fatalf("request %d refused inside the burst", i+1)
 19		}
 20	}
 21	if l.allow("c") {
 22		t.Error("request past the burst allowed")
 23	}
 24	if !l.allow("other") {
 25		t.Error("a different client was refused by the first one's budget")
 26	}
 27}
 28
 29func TestRateLimiterRefillsAtTheConfiguredRate(t *testing.T) {
 30	l := newRateLimiter(60, 1) // one token per second
 31	now := time.Now()
 32	l.now = func() time.Time { return now }
 33
 34	l.allow("c")
 35	if l.allow("c") {
 36		t.Fatal("second immediate request allowed with burst 1")
 37	}
 38	now = now.Add(500 * time.Millisecond)
 39	if l.allow("c") {
 40		t.Error("allowed after half a second, want a full second per token")
 41	}
 42	now = now.Add(600 * time.Millisecond)
 43	if !l.allow("c") {
 44		t.Error("refused after more than a second, want one token refilled")
 45	}
 46}
 47
 48func TestRateLimiterPrunesIdleBuckets(t *testing.T) {
 49	l := newRateLimiter(60, 1)
 50	now := time.Now()
 51	l.now = func() time.Time { return now }
 52
 53	l.allow("old")
 54	now = now.Add(bucketIdle + time.Minute)
 55	l.prune(now)
 56
 57	if _, ok := l.buckets["old"]; ok {
 58		t.Error("idle bucket survived a prune")
 59	}
 60}
 61
 62func TestClientAddrTrustsForwardedForOnlyBehindAProxy(t *testing.T) {
 63	cases := []struct{ remote, xff, want string }{
 64		{"127.0.0.1:1234", "203.0.113.5", "203.0.113.5"},
 65		{"10.0.0.2:1234", "198.51.100.7, 203.0.113.5", "203.0.113.5"},
 66		{"127.0.0.1:1234", "", "127.0.0.1"},
 67		{"203.0.113.9:1234", "198.51.100.7", "203.0.113.9"},
 68		{"[::1]:1234", "203.0.113.5", "203.0.113.5"},
 69	}
 70	for _, c := range cases {
 71		r := httptest.NewRequest(http.MethodGet, "/", nil)
 72		r.RemoteAddr = c.remote
 73		if c.xff != "" {
 74			r.Header.Set("X-Forwarded-For", c.xff)
 75		}
 76		if got := clientAddr(r); got != c.want {
 77			t.Errorf("remote %s xff %q: got %s, want %s", c.remote, c.xff, got, c.want)
 78		}
 79	}
 80}
 81
 82func TestRobotsTXTCarriesTheBaseURI(t *testing.T) {
 83	out := robotsTXT("/art/")
 84	for _, want := range []string{"Disallow: /art/search\n", "Disallow: /art/api\n", "Disallow: /art/*?p=\n", "Crawl-delay: 10\n"} {
 85		if !strings.Contains(out, want) {
 86			t.Errorf("robots.txt lacks %q:\n%s", want, out)
 87		}
 88	}
 89}
 90
 91// TestHandlerLimitsPagesButNotMedia drives the real handler: a client past its
 92// budget gets 429 on a page with Retry-After and no Cache-Control, while its
 93// stylesheet, favicon and avatar requests are never counted.
 94func TestHandlerLimitsPagesButNotMedia(t *testing.T) {
 95	limiter, uri := daLimiter, CFG.URI
 96	daLimiter = newRateLimiter(60, 2)
 97	CFG.URI = "/"
 98	defer func() { daLimiter, CFG.URI = limiter, uri }()
 99
100	orig := fetchAvatar
101	fetchAvatar = func(string, rune) (string, error) { return "png", nil }
102	defer func() { fetchAvatar = orig }()
103
104	from := func(target string) *httptest.ResponseRecorder {
105		loadTemplates()
106		rec := httptest.NewRecorder()
107		req := httptest.NewRequest(http.MethodGet, target, nil)
108		req.RemoteAddr = "203.0.113.5:4000"
109		Handler()(rec, req)
110		return rec
111	}
112
113	from("/about")
114	from("/about")
115	third := from("/about")
116	if third.Code != 429 {
117		t.Fatalf("third page request: status %d, want 429", third.Code)
118	}
119	if third.Header().Get("Retry-After") == "" || third.Header().Get("Cache-Control") != "" {
120		t.Errorf("429 headers: Retry-After %q, Cache-Control %q", third.Header().Get("Retry-After"), third.Header().Get("Cache-Control"))
121	}
122
123	for _, target := range []string{"/stylesheet", "/favicon.ico", "/media/emojitar/alice?type=a", "/robots.txt"} {
124		if rec := from(target); rec.Code != 200 {
125			t.Errorf("%s: status %d while limited, want 200 (exempt)", target, rec.Code)
126		}
127	}
128}
129
130func TestHandlerServesRobotsTXT(t *testing.T) {
131	uri := CFG.URI
132	CFG.URI = "/"
133	defer func() { CFG.URI = uri }()
134
135	rec := serve(t, "/robots.txt")
136	if rec.Code != 200 || !strings.HasPrefix(rec.Body.String(), "User-agent: *\n") {
137		t.Errorf("robots.txt: status %d body %q", rec.Code, rec.Body.String())
138	}
139	if ct := rec.Header().Get("Content-Type"); !strings.HasPrefix(ct, "text/plain") {
140		t.Errorf("Content-Type %q, want text/plain", ct)
141	}
142}