internal/httpd/apilimit_test.go

 1package httpd
 2
 3import (
 4	"testing"
 5	"time"
 6)
 7
 8func TestAPILimiterBurstThenRefill(t *testing.T) {
 9	l := newAPILimiter(60) // 1/s sustained, burst 60
10
11	// The burst is spendable immediately.
12	for i := 0; i < 60; i++ {
13		if ok, _ := l.allow("u1", false); !ok {
14			t.Fatalf("read %d rejected inside the burst", i)
15		}
16	}
17	ok, wait := l.allow("u1", false)
18	if ok {
19		t.Fatal("burst did not run out")
20	}
21	if wait < time.Second {
22		t.Errorf("Retry-After %v, want at least a second", wait)
23	}
24
25	// Refill is by elapsed time, so a caller recovers without a restart.
26	l.buckets["u1"].last = time.Now().Add(-5 * time.Second)
27	if ok, _ := l.allow("u1", false); !ok {
28		t.Error("bucket did not refill")
29	}
30}
31
32func TestAPILimiterWritesHaveTheirOwnBudget(t *testing.T) {
33	l := newAPILimiter(60) // write burst is 6
34
35	for i := 0; i < 6; i++ {
36		if ok, _ := l.allow("u1", true); !ok {
37			t.Fatalf("write %d rejected inside the write burst", i)
38		}
39	}
40	if ok, _ := l.allow("u1", true); ok {
41		t.Fatal("writes are not separately limited")
42	}
43	// Reads still have budget: a client that hit the write ceiling can
44	// still render a page.
45	if ok, _ := l.allow("u1", false); !ok {
46		t.Error("exhausting writes also blocked reads")
47	}
48}
49
50func TestAPILimiterIsPerCaller(t *testing.T) {
51	l := newAPILimiter(60)
52	for i := 0; i < 60; i++ {
53		l.allow("u1", false)
54	}
55	if ok, _ := l.allow("u1", false); ok {
56		t.Fatal("u1 not limited")
57	}
58	if ok, _ := l.allow("u2", false); !ok {
59		t.Error("one caller's flood limited another")
60	}
61}
62
63// A caller with no budget must not be able to buy more by making the
64// limiter forget them; the sweep only drops buckets that are idle.
65func TestAPILimiterSweepKeepsActiveCallers(t *testing.T) {
66	l := newAPILimiter(60)
67	for i := 0; i < 60; i++ {
68		l.allow("victim", false)
69	}
70	for i := 0; i < 4100; i++ {
71		l.allow(string(rune(i))+"filler", false)
72	}
73	if ok, _ := l.allow("victim", false); ok {
74		t.Error("an active caller's bucket was swept, resetting their budget")
75	}
76}