host_limiterでhostごとにreq limitを設定できるようにする

Signed-off-by: drumato <[email protected]>
2manymws · Dec 25, 2024 · 9bec326 · 9bec326
1 parent 8bdea3c
commit 9bec326
Show file tree

Hide file tree

Showing 3 changed files with 81 additions and 51 deletions.
diff --git 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
@@ -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
@@ -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)