1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
|
package handlers
import (
"net/http"
"net/http/httptest"
"testing"
"github.com/stretchr/testify/require"
"gitlab.com/gitlab-org/gitlab-pages/internal/config"
"gitlab.com/gitlab-org/gitlab-pages/internal/testhelpers"
)
var next = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusNoContent)
})
func TestRatelimiter(t *testing.T) {
tt := map[string]struct {
firstRemoteAddr string
firstTarget string
secondRemoteAddr string
secondTarget string
expectedSecondCode int
}{
"rejected_by_ip": {
firstRemoteAddr: "10.0.0.1",
firstTarget: "https://domain.gitlab.io",
secondRemoteAddr: "10.0.0.1",
secondTarget: "https://different.gitlab.io",
expectedSecondCode: http.StatusTooManyRequests,
},
"rejected_by_domain": {
firstRemoteAddr: "10.0.0.1",
firstTarget: "https://domain.gitlab.io",
secondRemoteAddr: "10.0.0.2",
secondTarget: "https://domain.gitlab.io",
expectedSecondCode: http.StatusTooManyRequests,
},
"different_ip_and_domain_passes": {
firstRemoteAddr: "10.0.0.1",
firstTarget: "https://domain.gitlab.io",
secondRemoteAddr: "10.0.0.2",
secondTarget: "https://different.gitlab.io",
expectedSecondCode: http.StatusNoContent,
},
}
for name, tc := range tt {
t.Run(name, func(t *testing.T) {
conf := config.RateLimit{
SourceIPLimitPerSecond: 0.1,
SourceIPBurst: 1,
DomainLimitPerSecond: 0.1,
DomainBurst: 1,
}
handler := Ratelimiter(next, &conf)
r1 := httptest.NewRequest(http.MethodGet, tc.firstTarget, nil)
r1.RemoteAddr = tc.firstRemoteAddr
firstCode, _ := testhelpers.PerformRequest(t, handler, r1)
require.Equal(t, http.StatusNoContent, firstCode)
r2 := httptest.NewRequest(http.MethodGet, tc.secondTarget, nil)
r2.RemoteAddr = tc.secondRemoteAddr
secondCode, _ := testhelpers.PerformRequest(t, handler, r2)
require.Equal(t, tc.expectedSecondCode, secondCode)
})
}
}
|
