From 9bec3267046283ce588bc3ab23efcb5c1fa501a3 Mon Sep 17 00:00:00 2001 From: drumato Date: Wed, 25 Dec 2024 17:28:57 +0900 Subject: [PATCH] =?UTF-8?q?host=5Flimiter=E3=81=A7host=E3=81=94=E3=81=A8?= =?UTF-8?q?=E3=81=ABreq=20limit=E3=82=92=E8=A8=AD=E5=AE=9A=E3=81=A7?= =?UTF-8?q?=E3=81=8D=E3=82=8B=E3=82=88=E3=81=86=E3=81=AB=E3=81=99=E3=82=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: drumato --- coverage.out | 106 +++++++++++++++++++++++-------------------- host_limiter.go | 10 +++- host_limiter_test.go | 16 ++++++- 3 files changed, 81 insertions(+), 51 deletions(-) diff --git a/coverage.out b/coverage.out index f36be72..e5909bd 100644 --- a/coverage.out +++ b/coverage.out @@ -1,46 +1,52 @@ mode: count -github.com/2manymws/rlutils/base_limiter.go:46.15,51.33 3 36 -github.com/2manymws/rlutils/base_limiter.go:51.33,52.20 1 52 -github.com/2manymws/rlutils/base_limiter.go:52.20,54.4 1 38 -github.com/2manymws/rlutils/base_limiter.go:57.2,67.3 1 36 -github.com/2manymws/rlutils/base_limiter.go:70.57,71.29 1 8 -github.com/2manymws/rlutils/base_limiter.go:71.29,73.32 2 8 -github.com/2manymws/rlutils/base_limiter.go:73.32,74.41 1 7 -github.com/2manymws/rlutils/base_limiter.go:74.41,75.38 1 11 -github.com/2manymws/rlutils/base_limiter.go:75.38,77.6 1 1 -github.com/2manymws/rlutils/base_limiter.go:78.5,78.61 1 11 -github.com/2manymws/rlutils/base_limiter.go:84.51,85.29 1 6 -github.com/2manymws/rlutils/base_limiter.go:85.29,87.29 2 6 -github.com/2manymws/rlutils/base_limiter.go:87.29,88.41 1 5 -github.com/2manymws/rlutils/base_limiter.go:88.41,90.5 1 8 -github.com/2manymws/rlutils/base_limiter.go:95.61,96.29 1 8 -github.com/2manymws/rlutils/base_limiter.go:96.29,98.3 1 8 -github.com/2manymws/rlutils/base_limiter.go:101.61,102.29 1 8 -github.com/2manymws/rlutils/base_limiter.go:102.29,104.3 1 8 -github.com/2manymws/rlutils/base_limiter.go:107.61,108.29 1 8 -github.com/2manymws/rlutils/base_limiter.go:108.29,110.3 1 8 -github.com/2manymws/rlutils/base_limiter.go:113.70,115.2 1 0 -github.com/2manymws/rlutils/base_limiter.go:117.37,119.2 1 0 -github.com/2manymws/rlutils/base_limiter.go:121.61,123.2 1 23 -github.com/2manymws/rlutils/base_limiter.go:125.64,126.34 1 31 -github.com/2manymws/rlutils/base_limiter.go:126.34,128.3 1 24 -github.com/2manymws/rlutils/base_limiter.go:129.2,131.11 3 7 -github.com/2manymws/rlutils/base_limiter.go:134.60,135.31 1 29 -github.com/2manymws/rlutils/base_limiter.go:135.31,137.3 1 24 -github.com/2manymws/rlutils/base_limiter.go:138.2,139.11 2 5 -github.com/2manymws/rlutils/base_limiter.go:142.58,150.4 1 23 -github.com/2manymws/rlutils/base_limiter.go:150.4,151.27 1 66 -github.com/2manymws/rlutils/base_limiter.go:151.27,152.38 1 3 -github.com/2manymws/rlutils/base_limiter.go:152.38,153.36 1 3 -github.com/2manymws/rlutils/base_limiter.go:153.36,155.6 1 3 -github.com/2manymws/rlutils/base_limiter.go:160.2,160.13 1 20 -github.com/2manymws/rlutils/base_limiter.go:162.36,164.53 1 12 -github.com/2manymws/rlutils/base_limiter.go:164.53,165.15 1 22 -github.com/2manymws/rlutils/base_limiter.go:165.15,167.4 1 12 -github.com/2manymws/rlutils/base_limiter.go:169.2,169.43 1 0 -github.com/2manymws/rlutils/base_limiter.go:172.50,173.26 1 7 -github.com/2manymws/rlutils/base_limiter.go:173.26,176.3 2 2 -github.com/2manymws/rlutils/base_limiter.go:177.2,177.15 1 5 +github.com/2manymws/rlutils/base_limiter.go:48.15,53.33 3 37 +github.com/2manymws/rlutils/base_limiter.go:53.33,54.20 1 53 +github.com/2manymws/rlutils/base_limiter.go:54.20,56.4 1 38 +github.com/2manymws/rlutils/base_limiter.go:59.2,70.3 1 37 +github.com/2manymws/rlutils/base_limiter.go:73.57,74.29 1 8 +github.com/2manymws/rlutils/base_limiter.go:74.29,76.32 2 8 +github.com/2manymws/rlutils/base_limiter.go:76.32,77.41 1 7 +github.com/2manymws/rlutils/base_limiter.go:77.41,78.38 1 11 +github.com/2manymws/rlutils/base_limiter.go:78.38,80.6 1 1 +github.com/2manymws/rlutils/base_limiter.go:81.5,81.61 1 11 +github.com/2manymws/rlutils/base_limiter.go:87.51,88.29 1 6 +github.com/2manymws/rlutils/base_limiter.go:88.29,90.29 2 6 +github.com/2manymws/rlutils/base_limiter.go:90.29,91.41 1 5 +github.com/2manymws/rlutils/base_limiter.go:91.41,93.5 1 8 +github.com/2manymws/rlutils/base_limiter.go:98.63,99.29 1 0 +github.com/2manymws/rlutils/base_limiter.go:99.29,101.3 1 0 +github.com/2manymws/rlutils/base_limiter.go:104.61,105.29 1 8 +github.com/2manymws/rlutils/base_limiter.go:105.29,107.3 1 8 +github.com/2manymws/rlutils/base_limiter.go:110.61,111.29 1 8 +github.com/2manymws/rlutils/base_limiter.go:111.29,113.3 1 8 +github.com/2manymws/rlutils/base_limiter.go:116.61,117.29 1 8 +github.com/2manymws/rlutils/base_limiter.go:117.29,119.3 1 8 +github.com/2manymws/rlutils/base_limiter.go:122.70,124.2 1 0 +github.com/2manymws/rlutils/base_limiter.go:126.37,128.2 1 0 +github.com/2manymws/rlutils/base_limiter.go:130.61,132.2 1 24 +github.com/2manymws/rlutils/base_limiter.go:134.63,135.43 1 21 +github.com/2manymws/rlutils/base_limiter.go:135.43,136.12 1 0 +github.com/2manymws/rlutils/base_limiter.go:136.12,138.4 1 0 +github.com/2manymws/rlutils/base_limiter.go:140.2,140.13 1 21 +github.com/2manymws/rlutils/base_limiter.go:142.64,143.34 1 32 +github.com/2manymws/rlutils/base_limiter.go:143.34,145.3 1 25 +github.com/2manymws/rlutils/base_limiter.go:146.2,148.11 3 7 +github.com/2manymws/rlutils/base_limiter.go:151.60,152.31 1 30 +github.com/2manymws/rlutils/base_limiter.go:152.31,154.3 1 25 +github.com/2manymws/rlutils/base_limiter.go:155.2,156.11 2 5 +github.com/2manymws/rlutils/base_limiter.go:159.58,167.4 1 24 +github.com/2manymws/rlutils/base_limiter.go:167.4,168.27 1 69 +github.com/2manymws/rlutils/base_limiter.go:168.27,169.38 1 3 +github.com/2manymws/rlutils/base_limiter.go:169.38,170.36 1 3 +github.com/2manymws/rlutils/base_limiter.go:170.36,172.6 1 3 +github.com/2manymws/rlutils/base_limiter.go:177.2,177.13 1 21 +github.com/2manymws/rlutils/base_limiter.go:179.36,181.53 1 12 +github.com/2manymws/rlutils/base_limiter.go:181.53,182.15 1 22 +github.com/2manymws/rlutils/base_limiter.go:182.15,184.4 1 12 +github.com/2manymws/rlutils/base_limiter.go:186.2,186.43 1 0 +github.com/2manymws/rlutils/base_limiter.go:189.50,190.26 1 7 +github.com/2manymws/rlutils/base_limiter.go:190.26,193.3 2 2 +github.com/2manymws/rlutils/base_limiter.go:194.2,194.15 1 5 github.com/2manymws/rlutils/country_limiter.go:36.28,38.16 2 6 github.com/2manymws/rlutils/country_limiter.go:38.16,40.3 1 0 github.com/2manymws/rlutils/country_limiter.go:41.2,44.30 3 6 @@ -78,16 +84,18 @@ github.com/2manymws/rlutils/get_parameter_limiter.go:42.45,44.2 1 0 github.com/2manymws/rlutils/get_parameter_limiter.go:46.71,47.27 1 4 github.com/2manymws/rlutils/get_parameter_limiter.go:47.27,49.3 1 0 github.com/2manymws/rlutils/get_parameter_limiter.go:50.2,50.36 1 4 -github.com/2manymws/rlutils/get_parameter_limiter.go:50.36,51.32 1 4 +github.com/2manymws/rlutils/get_parameter_limiter.go:50.36,51.32 1 5 github.com/2manymws/rlutils/get_parameter_limiter.go:51.32,57.4 1 3 github.com/2manymws/rlutils/get_parameter_limiter.go:60.2,60.36 1 1 github.com/2manymws/rlutils/get_parameter_limiter.go:63.78,65.2 1 0 -github.com/2manymws/rlutils/host_limiter.go:21.16,30.2 1 1 -github.com/2manymws/rlutils/host_limiter.go:32.37,34.2 1 0 -github.com/2manymws/rlutils/host_limiter.go:36.63,37.27 1 1 -github.com/2manymws/rlutils/host_limiter.go:37.27,39.3 1 0 -github.com/2manymws/rlutils/host_limiter.go:40.2,44.8 1 1 -github.com/2manymws/rlutils/host_limiter.go:47.70,49.2 1 0 +github.com/2manymws/rlutils/host_limiter.go:23.16,33.2 1 2 +github.com/2manymws/rlutils/host_limiter.go:35.37,37.2 1 0 +github.com/2manymws/rlutils/host_limiter.go:39.63,40.27 1 2 +github.com/2manymws/rlutils/host_limiter.go:40.27,42.3 1 0 +github.com/2manymws/rlutils/host_limiter.go:44.2,45.44 2 2 +github.com/2manymws/rlutils/host_limiter.go:45.44,47.3 1 1 +github.com/2manymws/rlutils/host_limiter.go:48.2,52.8 1 2 +github.com/2manymws/rlutils/host_limiter.go:55.70,57.2 1 0 github.com/2manymws/rlutils/ip_limiter.go:22.14,31.2 1 2 github.com/2manymws/rlutils/ip_limiter.go:33.35,35.2 1 0 github.com/2manymws/rlutils/ip_limiter.go:37.61,38.27 1 2 diff --git a/host_limiter.go b/host_limiter.go index 2d1d60f..7a37c78 100644 --- a/host_limiter.go +++ b/host_limiter.go @@ -9,6 +9,8 @@ import ( type HostLimiter struct { BaseLimiter + // ReqLimitPerHost はホストごとに特別なしきい値を設定できるようにします + ReqLimitPerHost map[string]int } // ホストごとにリクエスト数を制限する @@ -26,6 +28,7 @@ func NewHostLimiter( onRequestLimit, setter..., ), + ReqLimitPerHost: make(map[string]int), } } @@ -37,9 +40,14 @@ func (l *HostLimiter) Rule(r *http.Request) (*rl.Rule, error) { if !l.IsTargetRequest(r) { return &rl.Rule{ReqLimit: -1}, nil } + + reqLimit := l.reqLimit + if v, ok := l.ReqLimitPerHost[r.Host]; ok { + reqLimit = v + } return &rl.Rule{ Key: r.Host, - ReqLimit: l.reqLimit, + ReqLimit: reqLimit, WindowLen: l.windowLen, }, nil } diff --git a/host_limiter_test.go b/host_limiter_test.go index 94f4336..c0febac 100644 --- a/host_limiter_test.go +++ b/host_limiter_test.go @@ -14,11 +14,24 @@ func TestHostLimiter(t *testing.T) { cases := []struct { name string host string + reqLimitPerHost map[string]int + expectedReqLimit int expectedToBeLimited bool }{ { name: "Host is limited", host: "api.example.com", + reqLimitPerHost: map[string]int{}, + expectedReqLimit: 5, + expectedToBeLimited: true, + }, + { + name: "when host is more limited", + host: "api.example.com", + reqLimitPerHost: map[string]int{ + "api.example.com": 1, + }, + expectedReqLimit: 1, expectedToBeLimited: true, }, } @@ -38,6 +51,7 @@ func TestHostLimiter(t *testing.T) { nil, nil, ) + limiter.ReqLimitPerHost = tc.reqLimitPerHost // Using the mock counter instead of the real one. limiter.Counter = mockCounter @@ -49,7 +63,7 @@ func TestHostLimiter(t *testing.T) { if tc.expectedToBeLimited { assert.NotNil(t, rule) assert.Equal(t, tc.host, rule.Key) - assert.Equal(t, reqLimit, rule.ReqLimit) + assert.Equal(t, tc.expectedReqLimit, rule.ReqLimit) assert.Equal(t, windowLen, rule.WindowLen) } else { assert.Nil(t, rule)